MACE

目录

• MACE
  • 目录
  • 关于 MACE
  • 文档
  • 安装
    • 通过 pip 安装
    • 从源码安装
  • 使用
    • 训练
    • 评估
  • 教程
  • 使用 cuEquivariance 进行 CUDA 加速
  • Weights and Biases 用于实验跟踪
  • 预训练的基础模型
    • MACE-MP：材料项目力场
      • 在 ASE 中的示例用法
    • MACE-OFF：可迁移的有机力场
      • 在 ASE 中的示例用法
    • MACE-Polar：静电基础模型
    • 微调基础模型
    • 最新推荐的基础模型
  • 缓存
  • 开发
  • 参考文献
  • 联系方式
  • 许可证

关于 MACE

MACE 提供快速且精确的机器学习原子间势能函数，采用高阶等变消息传递机制。

本仓库包含由 Ilyes Batatia、Gregor Simm、David Kovacs 以及 Gabor Csanyi 团队和朋友们（见贡献者）共同开发的 MACE 参考实现。

此外还提供：

• MACE in JAX，目前在评估阶段速度大约是 PyTorch 版本的两倍，但为了获得最佳性能，建议使用 PyTorch 进行训练。
• MACE 层，用于构建适用于任意三维点云的高阶等变图神经网络。

文档

部分文档可在以下网址查阅：https://mace-docs.readthedocs.io

安装

1. 系统要求

• Python >= 3.8 （若使用 openMM，请使用 Python = 3.9）
• PyTorch >= 1.12 （PyTorch 2.1 不支持以 float64 进行训练，但 2.2 及更高版本支持；PyTorch 2.4.1 不受支持）

请确保已安装 PyTorch。 请参阅 官方 PyTorch 安装指南 获取安装说明，并根据您的系统选择合适的选项。

通过 PyPI 安装

这是推荐的 MACE 安装方式。

pip install --upgrade pip
pip install mace-torch

注意： PyPI 上同名的包与此处无关。

从源码安装

git clone https://github.com/ACEsuit/mace.git
pip install ./mace

使用

训练

要训练一个 MACE 模型，可以使用 mace_run_train 脚本，该脚本通常位于 pip 安装二进制文件的默认路径下（或者你也可以直接运行 python3 <克隆目录路径>/mace/cli/run_train.py）。

mace_run_train \
    --name="MACE_model" \
    --train_file="train.xyz" \
    --valid_fraction=0.05 \
    --test_file="test.xyz" \
    --config_type_weights='{"Default":1.0}' \
    --E0s='{1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639}' \
    --model="MACE" \
    --hidden_irreps='128x0e + 128x1o' \
    --r_max=5.0 \
    --batch_size=10 \
    --max_num_epochs=1500 \
    --stage_two \
    --start_stage_two=1200 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --restart_latest \
    --device=cuda \

若要指定特定的验证集，可使用 --valid_file 参数。若需为验证集评估设置更大的批量大小，则可指定 --valid_batch_size。

要控制模型的规模，需要修改 --hidden_irreps 参数。对于大多数应用，推荐的默认模型规模是 --hidden_irreps='256x0e'（即 256 个不变消息）或 --hidden_irreps='128x0e + 128x1o'。如果模型精度不够，可以加入更高阶特征，例如 128x0e + 128x1o + 128x2e，或增加通道数至 256。此外，也可以通过 --num_channels=128 和 --max_L=1 参数来指定模型。

通常建议将孤立原子添加到训练集中，而不是像上述示例那样通过命令行输入它们的能量。要在训练集中标记孤立原子，只需在它们的信息字段中设置 config_type=IsolatedAtom 即可。

从头开始训练模型时，如果你不想使用或不知道孤立原子的能量，可以使用 --E0s="average" 选项，它会通过最小二乘法回归来估计原子能量。需要注意的是，使用拟合得到的 E0 值相当于拟合原子能量相对于平均值的偏差，而非拟合原子化能（后者是在使用孤立原子 E0 时的情况），这很可能导致分子动力学应用中的势能不够稳定。

在微调基础模型时，可以使用 --E0s="estimated" 选项，它会通过求解一个线性系统来优化修正基础模型在训练数据上的预测结果，从而估计原子参考能量。这种方法首先在所有训练构型上运行基础模型，计算预测误差（参考能量减去预测能量），然后通过最小二乘法求解最优的各元素 E0 修正值。总体而言，这种方式比“平均”选项更为理想。

如果启用了 --stage_two 关键字（之前称为 swa），则会在训练的最后约 20% 的 epoch 中提高损失函数中能量项的权重（从 --start_stage_two epoch 开始）。这一设置通常有助于降低能量误差。

可以通过 --default_dtype 关键字调整精度，默认为 float64，但使用 float32 可显著提升速度（通常在训练中可提速约两倍）。

--batch_size 和 --max_num_epochs 参数应根据训练集的大小进行调整。当训练数据量增加时，应增大批量大小；而当数据量减少时，则应减少 epoch 数。一个初步设置的经验法则可以是：将梯度更新次数设定为 200,000，其计算公式为 $\text{max-num-epochs}*\frac{\text{num-configs-training}}{\text{batch-size}}$。

代码能够处理具有异质标签的训练集，例如同时包含带有应力的体相结构和孤立分子的数据集。在这种情况下，为了让代码忽略分子上的应力，可以在分子的配置中添加 config_stress_weight = 0.0。

默认情况下，在训练结束时，结果文件夹中会生成一张图表，显示训练过程中损失和 RMSE 的变化趋势，以及模型在训练集、验证集和测试集上的预测散点图。可以通过 --plot False 来禁用此功能。若想在整个训练过程中跟踪这些指标（不包括对测试集的预测），可以启用针对训练集和验证集的定期绘图功能，方法是设置 --plot_frequency N，这样每 N 个 epoch 就会更新一次图表。

Apple Silicon GPU 加速

要使用 Apple Silicon GPU 加速，请确保安装最新版本的 PyTorch，并指定 --device=mps。

多 GPU 训练

进行多 GPU 训练时，可以使用 --distributed 标志。这将利用 PyTorch 的 DistributedDataParallel 模块在多个 GPU 上并行训练模型。对于大型数据集，建议结合在线数据加载方式（见下文）。示例 Slurm 脚本可在 mace/scripts/distributed_example.sbatch 中找到。

YAML 配置

支持使用 YAML 文件解析全部或部分参数。例如，要使用上述参数训练模型，可以创建一个名为 your_configs.yaml 的 YAML 文件，内容如下：

name: nacl
seed: 2024
train_file: train.xyz
stage_two: yes
start_stage_two: 1200
max_num_epochs: 1500
device: cpu
test_file: test.xyz
E0s:
  41: -1029.2809654211628
  38: -1484.1187695035828
  8: -2042.0330099956639
config_type_weights:
  Default: 1.0

然后在命令行中追加 --config="your_configs.yaml"。命令行中指定的任何参数都会覆盖 YAML 文件中的相应参数。

评估

要对 XYZ 文件中的体系使用你的 MACE 模型进行评估，可以运行 mace_eval_configs：

mace_eval_configs \
    --configs="your_configs.xyz" \
    --model="your_model.model" \
    --output="./your_output.xyz"

教程

你可以运行我们的 Colab 教程，快速入门 MACE。

我们还有更详细的 Colab 教程，包括：

• MACE 训练与评估简介
• MACE 主动学习与微调简介
• MACE 理论与代码（进阶）

使用 cuEquivariance 的 CUDA 加速

MACE 支持通过 cuEquivariance 库实现 CUDA 加速。要安装该库并使用加速功能，请参阅我们的文档：https://mace-docs.readthedocs.io/en/latest/guide/cuda_acceleration.html。

大型数据集的在线数据加载

如果你有一个可能无法完全放入 GPU 显存的大型数据集，建议在 CPU 上对数据进行预处理，并使用在线数据加载来训练模型。要对以 xyz 文件形式提供的数据集进行预处理，请运行 preprocess_data.py 脚本。以下是一个示例：

mkdir processed_data
python ./mace/scripts/preprocess_data.py \
    --train_file="/path/to/train_large.xyz" \
    --valid_fraction=0.05 \
    --test_file="/path/to/test_large.xyz" \
    --atomic_numbers="[1, 6, 7, 8, 9, 15, 16, 17, 35, 53]" \
    --r_max=4.5 \
    --h5_prefix="processed_data/" \
    --compute_statistics \
    --E0s="average" \
    --seed=123 \

要查看所有选项及其简要说明，请运行 python ./mace/scripts/preprocess_data.py --help 。该脚本将在 processed_data 文件夹中生成多个 HDF5 文件，可用于训练。其中将分别有一个用于训练的文件夹、一个用于验证的文件夹，以及针对测试集中每种 config_type 的单独文件夹。要训练模型，请按如下方式使用 run_train.py 脚本：

python ./mace/scripts/run_train.py \
    --name="MACE_on_big_data" \
    --num_workers=16 \
    --train_file="./processed_data/train.h5" \
    --valid_file="./processed_data/valid.h5" \
    --test_dir="./processed_data" \
    --statistics_file="./processed_data/statistics.json" \
    --model="ScaleShiftMACE" \
    --num_interactions=2 \
    --num_channels=128 \
    --max_L=1 \
    --correlation=3 \
    --batch_size=32 \
    --valid_batch_size=32 \
    --max_num_epochs=100 \
    --stage_two \
    --start_stage_two=60 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --error_table='PerAtomMAE' \
    --device=cuda \
    --seed=123 \

使用 Weights and Biases 进行实验跟踪

如果你想将 MACE 与 Weights and Biases 配合使用来记录你的实验，只需通过以下命令安装即可：

pip install ./mace[wandb]

然后指定必要的关键字参数（--wandb、--wandb_project、--wandb_entity、--wandb_name、--wandb_log_hypers）。

预训练的基础模型

我们提供了一系列适用于各种应用场景的预训练基础模型。这些模型可以直接用于推理，也可以作为在新数据集上进行微调的起点。 基础模型领域正在迅速发展。请访问 MACE-MP GitHub 仓库 和 MACE-OFF23 GitHub 仓库 查看最新发布的版本。

最新推荐的基础模型

模型名称	涉及元素数	训练数据集	理论水平	目标系统	模型规模	GitHub发布	备注	许可证
MACE-MP-0a	89	MPTrj	DFT (PBE+U)	材料	小, 中, 大	>=v0.3.6	基础模型的首次发布。	MIT
MACE-MP-0b3	89	MPTrj	DFT (PBE+U)	材料	中	>=v0.3.10	提高了高压稳定性及参考能。	MIT
MACE-MPA-0	89	MPTrj + sAlex	DFT (PBE+U)	材料	中-mpa-0	>=v0.3.10	提升了材料精度及高压稳定性。	MIT
MACE-OMAT-0	89	OMAT	DFT (PBE+U) VASP 54	材料	中-omat-0	>=v0.3.10		ASL
MACE-OFF23	10	SPICE v1	DFT (wB97M+D3)	有机化学	小, 中, 大	>=v0.3.6	首次发布，覆盖中性有机化学。	ASL
MACE-MATPES-PBE-0	89	MATPES-PBE	DFT (PBE)	材料	中	>=v0.3.10	无+U校正。	ASL
MACE-MATPES-r2SCAN-0	89	MATPES-r2SCAN	DFT (r2SCAN)	材料	中	>=v0.3.10	更适合材料体系的泛函。	ASL
MACE-OMOL-0	89	OMOL	DFT (wB97M-VV10)	分子/过渡金属/阳离子	大	>=v0.3.14	带电荷/自旋嵌入，分子精度极佳。	ASL
MACE-MH-0/1	89	OMAT/OMOL/OC20/MATPES	DFT (PBE/R2SCAN/wB97M-VV10)	无机晶体、分子和表面。更多信息	mh-0 mh-1	>=v0.3.14	在表面/体相/分子等跨领域任务上表现优异。	ASL

MACE-MP：材料项目力场

我们与材料项目（MP）合作，基于从 MP 放松轨迹中选取的 MPTrj 数据集 中的 160 万个体相晶体，训练了一个涵盖 89 种元素的通用 MACE 势能模型。这些模型已在 GitHub 上发布，地址为 https://github.com/ACEsuit/mace-foundations。如果您使用它们，请引用我们的论文 arXiv:2401.00096，其中还包含大量示例应用和基准测试。

[!CAUTION] MACE-MP 模型是基于 VASP 输出的 MPTrj 原始 DFT 能量训练而成，不能直接与 MP 的 DFT 能量或 CHGNet 的能量进行比较。MP 和 CHGNet 的能量对部分过渡金属氧化物、氟化物（采用 GGA/GGA+U 混合校正）以及 14 种阴离子物种进行了 MP2020Compatibility 校正。更多详情请参阅 MP 文档 和 MP2020Compatibility.yaml。

在 ASE 中的使用示例

from mace.calculators import mace_mp
from ase import build

atoms = build.molecule('H2O')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float32", device='cuda')
atoms.calc = calc
print(atoms.get_potential_energy())

MACE-OFF：可迁移有机力场

有一系列（小、中、大）可迁移有机力场。这些模型可用于模拟有机分子、晶体和分子液体，也可作为在新数据集上微调的起点。这些模型以 ASL 许可证 发布。模型已在 GitHub 上发布，地址为 https://github.com/ACEsuit/mace-off。如果您使用它们，请引用我们的论文 arXiv:2312.15211，其中也包含了详细的基准测试和示例应用。

在 ASE 中的使用示例

from mace.calculators import mace_off
from ase import build

atoms = build.molecule('H2O')
calc = mace_off(model="medium", device='cuda')
atoms.calc = calc
print(atoms.get_potential_energy())

MACE-Polar：静电基础模型

PolarMACE 检查点是针对分子化学的静电基础模型，基于 OMol25 数据集训练而成。有关使用方法、输出以及训练/微调细节，请参阅 PolarMACE 指南：

• https://mace-docs.readthedocs.io/en/latest/guide/polar_mace.html

微调基础模型

要微调其中一个 mace-mp-0 基础模型，您可以使用 mace_run_train 脚本，并添加额外参数 --foundation_model=model_type。例如，要在新数据集上微调小型模型，可以使用以下命令：

mace_run_train \
  --name="MACE" \
  --foundation_model="small" \
  --train_file="train.xyz" \
  --valid_fraction=0.05 \
  --test_file="test.xyz" \
  --energy_weight=1.0 \
  --forces_weight=1.0 \
  --E0s="average" \
  --lr=0.01 \
  --scaling="rms_forces_scaling" \
  --batch_size=2 \
  --max_num_epochs=6 \
  --ema \
  --ema_decay=0.99 \
  --amsgrad \
  --default_dtype="float32" \
  --device=cuda \
  --seed=3

其他选项包括“medium”和“large”，或者指定基础模型的路径。如果您想微调其他模型，则会从 --foundation_model=$path_model 提供的路径加载模型，所有超参数将自动提取。

缓存

默认情况下，自动下载的模型（如 mace_mp、mace_off 以及用于微调的数据）会存储在 ~/.cache/mace 目录下。可以通过设置环境变量 XDG_CACHE_HOME 来更改缓存路径。当该变量被设置时，新的缓存路径将扩展为 $XDG_CACHE_HOME/.cache/mace。

开发

该项目使用 pre-commit 工具，在每次提交代码时自动执行代码格式化和 lint 检查。我们还使用 black、isort、pylint 和 mypy。建议您通过安装 dev 包来设置开发环境：

pip install -e ".[dev]"
pre-commit install

第二行命令将初始化 pre-commit，以便在每次提交时自动运行代码检查。我们已设置了 CI 来检查这一点，但强烈建议您在提交（和推送）之前先运行这些命令，以避免意外提交不良代码。

我们欢迎基于 MIT 许可证 的拉取请求。请将许可证文本复制粘贴为注释放入您的拉取请求中。

参考文献

如果您使用此代码，请引用我们的论文：

@inproceedings{Batatia2022mace,
  title={{MACE}: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields},
  author={Ilyes Batatia and David Peter Kovacs and Gregor N. C. Simm and Christoph Ortner and Gabor Csanyi},
  booktitle={Advances in Neural Information Processing Systems},
  editor={Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho},
  year={2022},
  url={https://openreview.net/forum?id=YPpSngE-ZU}
}

@misc{Batatia2022Design,
  title = {The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials},
  author = {Batatia, Ilyes and Batzner, Simon and Kov{\'a}cs, D{\'a}vid P{\'e}ter and Musaelian, Albert and Simm, Gregor N. C. and Drautz, Ralf and Ortner, Christoph and Kozinsky, Boris and Cs{\'a}nyi, G{\'a}bor},
  year = {2022},
  number = {arXiv:2205.06643},
  eprint = {2205.06643},
  eprinttype = {arxiv},
  doi = {10.48550/arXiv.2205.06643},
  archiveprefix = {arXiv}
 }

联系方式

如有任何问题，请发送邮件至 ilyes.batatia@ens-paris-saclay.fr。

如遇 bug 或功能请求，请使用 GitHub Issues。

许可证

MACE 代码根据 MIT 许可证 发布并分发。（请注意，上述链接中的部分模型可能附带不同的许可证）。

MACE 快速上手指南

MACE (Machine Learning Atomic Cluster Expansion) 是一个基于高阶等变消息传递的快速、高精度机器学习原子间势函数工具，适用于材料科学和分子动力学模拟。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux, macOS (支持 Apple Silicon), Windows (需配置相应环境)
• Python 版本: >= 3.8 (若需配合 openMM 使用，建议使用 Python 3.9)
• PyTorch: >= 1.12
  • 注意: PyTorch 2.1 不支持 float64 训练，请使用 2.2 或更高版本。暂不支持 PyTorch 2.4.1。
  • 请务必先根据官方指引安装适合您 CUDA 版本的 PyTorch：PyTorch 官方安装指南
  • 国内用户建议: 可使用清华或中科大镜像源加速 PyTorch 安装。
• 硬件加速:
  • NVIDIA GPU: 需安装 CUDA Toolkit。
  • Apple Silicon: 需安装最新版 PyTorch 以支持 MPS 加速。

安装步骤

推荐通过 PyPI 进行安装，这是最简便的方式。

1. 升级 pip 并安装 MACE

pip install --upgrade pip
pip install mace-torch

注意: PyPI 上有一个同名的 MACE 包与本工具无关，请务必安装 mace-torch。

2. 从源码安装 (可选)

如果您需要开发或获取最新代码：

git clone https://github.com/ACEsuit/mace.git
pip install ./mace

基本使用

1. 模型训练 (Training)

使用 mace_run_train 脚本进行训练。以下是一个最小化的训练示例，假设您已有 train.xyz (训练集) 和 test.xyz (测试集)。

mace_run_train \
    --name="MACE_model" \
    --train_file="train.xyz" \
    --valid_fraction=0.05 \
    --test_file="test.xyz" \
    --config_type_weights='{"Default":1.0}' \
    --E0s="average" \
    --model="MACE" \
    --hidden_irreps='128x0e + 128x1o' \
    --r_max=5.0 \
    --batch_size=10 \
    --max_num_epochs=1500 \
    --stage_two \
    --start_stage_two=1200 \
    --ema \
    --device=cuda

关键参数说明：

• --E0s="average": 自动通过最小二乘法估算孤立原子能量（适合从头训练且未知原子能量的情况）。若微调基础模型，建议使用 --E0s="estimated"。
• --hidden_irreps: 控制模型大小。推荐默认值为 '256x0e' 或 '128x0e + 128x1o'。
• --device: 指定设备，NVIDIA GPU 用 cuda，Apple Silicon 用 mps，CPU 用 cpu。
• --stage_two: 启用后会在训练后期增加能量损失的权重，有助于降低能量误差。

YAML 配置方式 (推荐): 对于复杂参数，可创建 config.yaml 文件：

name: nacl
seed: 2024
train_file: train.xyz
max_num_epochs: 1500
device: cuda
E0s: "average"
hidden_irreps: '128x0e + 128x1o'

运行命令：

mace_run_train --config="config.yaml"

2. 模型评估 (Evaluation)

训练完成后，使用 mace_eval_configs 对新的结构文件进行能量和力的预测：

mace_eval_configs \
    --configs="your_configs.xyz" \
    --model="your_model.model" \
    --output="./your_output.xyz"

3. 使用预训练基础模型 (可选)

MACE 提供了针对材料项目 (Materials Project) 和有机分子的预训练模型。若在代码中通过 ASE 调用，示例如下：

from mace.calculators import mace_mp

# 加载 Materials Project 预训练模型
calc = mace_mp(model="medium", device="cuda")
atoms.set_calculator(calc)
energy = atoms.get_potential_energy()
forces = atoms.get_forces()

更多资源

• 在线教程: Google Colab 快速入门
• 完整文档: MACE Documentation
• 大型数据集训练: 若数据过大无法放入显存，请参考文档使用 preprocess_data.py 进行离线预处理并结合 HDF5 格式进行在线加载。