PolarMACE水分子训练实操

PolarMACE在MACE局部模型基础上加入可学习的电荷与长程静电项。本篇用水分子的小数据集走一遍准备数据、训练、读取模型的流程，采用单极子版本入门；不把它与只直接预测偶极的EnergyDipolesMACE混为一谈。

1 下载并先读已有结果#

下载数据、小模型、实际输出、输入与绘图脚本。

python3 -m zipfile -e polarmace-water-practical-wave966-*.zip .
cd polarmace-water-practical
bash read.sh
bash plot.sh

读取结果在analysis-read/，图片在figures/。默认只整理资料包数据，不启动新训练或模型计算。

环境为MACE0.3.16、PyTorch2.8.0/CUDA12.6、e3nn0.4.4；配套graph_longrange固定为0.4.0的匹配源码版本，使用一张V10016G训练。本篇不修改公共MACE安装。

2 训练数据是什么#

标签来自已经完成的CP2K水分子电子结构与响应计算。16个几何按10训练、3验证、3独立测试分开；主要改变两条等长O–H键，是很窄的教学数据分布。

head -n 5 data/train.xyz

一个完整的训练帧如下：

3
Properties=species:S:1:pos:R:3:REF_forces:R:3:REF_charges:R:1 Lattice="12 0 0 0 12 0 0 0 12" pbc="F F F" config_type=Water geometry_id=0 REF_energy=-468.38766342125734 REF_total_charge=0 REF_dipole="_JSON [-2.7673130784442603e-10,8.243074364301676e-10,0.3819058895206621]" REF_polarizability="_JSON [[1.0786487884333646,1.4015309827257072e-09,-2.4598661781821466e-11],[1.416794007807199e-09,0.6903833607483519,1.773030049515023e-09],[1.2284512419957828e-10,-3.704617738226125e-11,0.9640221789568078]]"
O 6 6 6 -2.0820549489051102e-08 1.411458604197278e-08 -3.0108564675420522 0
H 6.6943328138049996 6 6.6035743066709998 2.8573106136726256 -3.0504348365842329e-08 1.5015386946226901 0
H 5.3056671861950004 6 6.6035743066709998 -2.8573105931038003 1.5830851085430049e-08 1.5015387049071027 0

REF_energy单位eV，REF_forces单位eV/Å，本资料包的REF_dipole采用e·Å。这里的REF_charges零值是占位项，不是Mulliken或Bader原子电荷标签；本例损失监督能量、力和总偶极，并约束整个分子总电荷为零。

数据单位必须先统一再训练。1 e·Å约为4.8032 Debye；不要只凭训练日志中固定的偶极单位文字替代自己输入标签的单位约定。

3 准备个人依赖与任务目录#

bash setup.sh
python3 prepare.py

setup.sh把匹配的MIT许可graph_longrange安装到解压目录的.long-range，需要能访问源码地址；不升级公共环境。这个例子固定匹配版本，换成其他版本时应先检查API兼容。

prepare.py生成work-new。已有该目录时保留结果，换一份新解压目录继续。

4 看完整训练配置#

cat train.yml

train_file: ../data/train.xyz
valid_file: ../data/valid.xyz
energy_key: REF_energy
forces_key: REF_forces
dipole_key: REF_dipole
polarizability_key: REF_polarizability
charges_key: REF_charges
total_charge_key: REF_total_charge
E0s: average
batch_size: 5
valid_batch_size: 3
max_num_epochs: 200
default_dtype: float64
seed: 96629
ema: false
ema_decay: 0.99
swa: false
num_workers: 0
device: cuda
log_dir: logs
model_dir: models
checkpoints_dir: checkpoints
results_dir: results
lr: 0.003
weight_decay: 5.0e-10
config_type_weights: '{''Water'':1.0}'
r_max: 4.0
hidden_irreps: 16x0e+16x1o+16x2e
MLP_irreps: 16x0e
num_interactions: 2
correlation: 2
max_ell: 2
max_L: 2
clip_grad: 10.0
patience: 200
scheduler_patience: 30
eval_interval: 1
save_cpu: true
compute_stress: false
name: water-PolarMACE
model: PolarMACE
loss: energy_forces_dipole
error_table: EnergyDipoleRMSE
dipole_weight: 100.0
polarizability_weight: 100.0
energy_weight: 1.0
forces_weight: 10.0
num_radial_basis: 8
num_cutoff_basis: 5
atomic_multipoles_max_l: 0
field_feature_max_l: 0
field_feature_widths: '[1.0]'
num_recursion_steps: 1
include_electrostatic_self_interaction: false
field_si: false
return_electrostatic_potentials: true

几个关键设置：

• model: PolarMACE与loss: energy_forces_dipole同时学习能量、力与偶极。
• atomic_multipoles_max_l: 0只使用单极子电荷，本篇没有训练高阶原子多极矩。
• field_feature_max_l: 0和一次递归给出一个简单的长程特征设置，不代表所有体系的最优参数。
• MLP_irreps: 16x0e为该读出层使用标量；隐藏层仍包含矢量和张量表示。
• max_num_epochs: 200是本例训练轮数；最终使用验证集选定的保存模型。

5 提交训练，再读取新模型#

cd work-new
sbatch ../submit.sh

任务结束后：

ls models
tail -n 12 native.log
cd ..
bash predict.sh

bash predict.sh提交一个GPU读取任务，读取自己新训练的work-new/models/water-PolarMACE.model，结果写入prediction-new.json，其中包含每个测试分子的能量、完整受力、有效电荷、偶极和静电能。

若只想在申请的GPU上读取资料包附带模型，安装个人依赖后可执行：

bash predict.sh reference

它写入prediction-reference.json，不训练新模型。已有预测文件请保留，使用新的工作目录避免覆盖。

6 本例实际精度#

这些结果说明这个小模型在所给数据范围内可用，不代表已覆盖断键、液态水、离子或陌生构型。单个孤立分子也不能用来证明凝聚态长程收敛或性能优势。

下一篇电荷、偶极与受力读取介绍如何理解这些输出，以及如何检查力和几何变换。

长程库的原始方法与接口见graph_electrostatics项目，模型训练入口见MACE官方文档。