NEP 操作演示

NEP 操作演示

这里，我们以 MatPL 源码根目录/example/HfO2/nep_demo 为例（HfO2 训练集来源），演示 NEP 模型的训练、测试、lammps 模拟以及其他功能。当前案例目录结构如下所示：

nep_demo/
├── nep_kokkos_lmps/
│   ├── kkin.lmp
│   ├── lmp.config
│   ├── nep_to_lmps.txt
│   ├── rungpu_mcloud.job
│   └── rungpu_station.job
├── nep_lmps/
│   ├── in.lmp-cpu-25
│   ├── in.lmp-kk
│   ├── lmp.config
│   ├── nep_to_lmps.txt
│   ├── runcpu_mcloud.job
│   ├── runcpu_station.job
│   ├── rungpu_mcloud.job
│   └── rungpu_station.job
├── nep_lmps_deviation/
│   ├── in.lmp-kk
│   ├── lmp.config
│   ├── nep0.txt
│   ├── nep1.txt
│   ├── nep2.txt
│   ├── nep3.txt
│   ├── rungpu_mcloud.job
│   └── rungpu_station.job
├── nep_test.json
├── nep_train.json
├── train_mcloud.job
└── train_station.job

• nep_train.json 和 nep_test.json 分别是 NEP 模型的训练和测试参数文件，其中的 ../pwdata/ 指向 HfO2/pwdata 训练数据目录。
• train_mcloud.job 和 train_station.job 分别是 mcloud 和自建集群的 Slurm 训练任务示例。
• nep_lmps 同时提供 CPU 和 NEP KOKKOS GPU 模拟的输入文件与作业脚本。
• nep_kokkos_lmps 提供 NEP KOKKOS GPU 模拟示例。
• nep_lmps_deviation 提供四模型偏差计算示例，用于主动学习流程。
• 各 lammps 案例目录中，lmp.config 为初始结构，nep_to_lmps.txt 或 nep*.txt 为 NEP 力场文件，*.lmp* 为 lammps 输入文件，*.job 为 Slurm 作业脚本。

train 训练

在 nep_demo 目录下使用如下命令即可开始训练：

MatPL train nep_train.json
# 或根据运行环境修改 Slurm 脚本后提交训练任务
sbatch train_mcloud.job
# 自建集群使用 sbatch train_station.job

输入文件解释

nep_train.json 中的内容如下所示，关于 NEP 的参数解释，请参考 NEP 参数手册：

{
    "model_type": "NEP",
    "atom_type": [
        8, 72
    ],
    "optimizer": {
        "optimizer": "ADAM",
        "epochs": 30, 
        "batch_size": 1,
        "print_freq": 10,
        "train_energy": true,
        "train_force": true,
        "train_virial": true
    },

    "format": "pwmlff/npy",
    "train_data": [
        "../pwdata/init_000_50/", "../pwdata/init_002_50/", 
        "../pwdata/init_004_50/", "../pwdata/init_006_50/", 
        "../pwdata/init_008_50/", "../pwdata/init_010_50/", 
        "../pwdata/init_012_50/", "../pwdata/init_014_50/", 
        "../pwdata/init_016_50/", "../pwdata/init_018_50/", 
        "../pwdata/init_020_20/", "../pwdata/init_022_20/", 
        "../pwdata/init_024_20/", "../pwdata/init_026_20/", 
        "../pwdata/init_001_50/", "../pwdata/init_003_50/", 
        "../pwdata/init_005_50/", "../pwdata/init_007_50/", 
        "../pwdata/init_009_50/", "../pwdata/init_011_50/", 
        "../pwdata/init_013_50/", "../pwdata/init_015_30/", 
        "../pwdata/init_017_50/", "../pwdata/init_019_50/", 
        "../pwdata/init_021_20/", "../pwdata/init_023_20/", 
        "../pwdata/init_025_20/", "../pwdata/init_027_20/"
    ],
    "valid_data":[
        "../pwdata/init_000_50/", "../pwdata/init_004_50/", 
        "../pwdata/init_008_50/"       
    ]
}

训练结束后的力场文件目录请参考 model_record 详解

多节点多卡训练

多节点多卡训练的目录结构与上面相同，案例请参考 MatPL 源码根目录/example/parallelnep 为例（HfO2 训练集来源）。

该目录下提供了单节点单卡 1node-1g-run.job 、单节点多卡 1node-4g-run.job 、多节点多卡 2node-8g-run.job 三种启动脚本供参考，该脚本适用于 mcloud 用户。 对于在线安装用户，MatPL-2026.3 的环境加载请参考文件env.sh。

多节点多卡训练启动时要求提供主机节点的地址以及可用端口，建议通过如下shell 命令自动获取

MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# 动态分配空闲端口
function get_free_port() {
    python -c 'import socket; s = socket.socket(socket.AF_INET, socket.SOCK_STREAM); s.bind(("", 0)); print(s.getsockname()[1]); s.close()'
}
MASTER_PORT=$(get_free_port)

export MASTER_ADDR=$MASTER_ADDR
export MASTER_PORT=$MASTER_PORT

echo "addrs: $MASTER_ADDR"
echo "port:  $MASTER_PORT"
echo "tasks: $SLURM_NTASKS"

srun MATPL train train.json 

警告

注意，NEP 多卡训练只支持使用 ADAM 优化器，不支持 LKF 或 GKF 优化器。

test 测试

test 命令支持来自 MatPL nep_model.ckpt 力场文件，以及在 lammps 或 GPUMD 中使用的 nep5.txt、nep4.txt 格式文件。

MatPL test nep_test.json

test.json 中的内容如下所示，参数解释请参考 参数手册

{
    "model_type": "NEP",
    "format": "pwmlff/npy",
    "model_load_file": "./model_record/nep_model.ckpt",
    "test_data": [
        "../init_000_50", "../init_004_50", "../init_008_50", 
        "../init_012_50", "../init_016_50", "../init_020_20", 
        "../init_024_20", "../init_001_50", "../init_005_50", 
        "../init_009_50", "../init_013_50", "../init_017_50", 
        "../init_021_20", "../init_025_20", "../init_002_50", 
        "../init_006_50", "../init_010_50", "../init_014_50", 
        "../init_018_50", "../init_022_20", "../init_026_20", 
        "../init_003_50", "../init_007_50", "../init_011_50", 
        "../init_015_30", "../init_019_50", "../init_023_20", 
        "../init_027_20"
    ]
}

测试结束后的力场文件目录请参考 test_result 详解

infer 推理单结构

infer 命令支持来自MatPL nep_model.ckpt 力场文件、GPUMD 的 nep4.txt 文件、 lammps 和 GPUMD 中通用的nep5.txt 格式文件。

MatPL infer nep_model.ckpt atom.config pwmat/config
MatPL infer gpumd_nep.txt 0.lammpstrj lammps/dump Hf O
# Hf O 为 lammps/dump格式的结构中的元素名称，Hf为结构中1号元素类型，O为元素中2号元素类型

推理成功后，将在窗口输出推理的总能、每原子能量、每原子受力和维里

totxt 转ckpt训练文件为nep5.txt

用于把 MatPL 训练的 nep_model.ckpt 文件转换为 txt 格式的nep5.txt 文件，该文件可用于 GPUMD 或 lammps-MatPL 中做分子动力学模拟。

MatPL totxt nep_model.ckpt

执行成功将在执行该命令的所在目录生成名称为nep5.txt文件

lammps MD

HfO2 案例分别提供了单模型、NEP KOKKOS 和多模型偏差计算目录。下面按“准备力场—选择案例—修改输入—启动模拟”的顺序说明。

1. 准备力场文件

正常训练结束后，model_record 目录中会生成 nep5.txt，可直接用于 lammps。如果只有 nep_model.ckpt，可使用前文介绍的 totxt 命令转换：

MatPL totxt nep_model.ckpt

将生成的 nep5.txt 复制到相应案例目录，并按输入脚本中的文件名设置为 nep_to_lmps.txt。lammps-MatPL 同样支持 GPUMD 的 NEP4 和 NEP5 力场文件。

2. 选择 lammps 案例

案例目录	用途	主要输入文件
nep_lmps	CPU 和 NEP KOKKOS GPU 模拟	in.lmp-cpu-25、in.lmp-kk
nep_kokkos_lmps	NEP KOKKOS GPU 模拟	kkin.lmp
nep_lmps_deviation	四模型偏差计算	in.lmp-kk、nep0.txt—nep3.txt

每个目录都分别提供了 *_mcloud.job 和 *_station.job，请根据运行环境选择并修改模块名称、软件路径、分区和 GPU 资源参数。

3. 设置 lammps 输入文件

NEP KOKKOS GPU 接口使用 half 近邻表并开启 Newton 通信。HfO2 单模型案例的核心设置如下：

package kokkos neigh half comm device
newton on

pair_style   matpl/nep/kk nep_to_lmps.txt
pair_coeff   * * 72 8

pair_coeff * * 后的元素按 data 文件中的原子类型顺序排列。可使用元素名称，也可使用原子序数；本案例中 1 号类型为 Hf（72），2 号类型为 O（8）。

多模型案例使用第一个模型进行 MD，其他模型参与偏差计算：

pair_style   matpl/nep/kk nep0.txt nep1.txt nep2.txt nep3.txt \
             out_freq ${DUMP_FREQ} out_file model_devi.out
pair_coeff   * * 72 8

out_freq 用于设置偏差输出频率，out_file 用于设置输出文件名。完整 NPT 控制参数请直接参考案例中的 in.lmp-kk、in.lmp-cpu-25 或 kkin.lmp，无需重复复制输入脚本。

4. 启动 lammps 模拟

在案例目录中，可直接提交与运行环境对应的 Slurm 脚本，例如：

sbatch rungpu_mcloud.job
# 或
sbatch rungpu_station.job

也可在加载 lammps 环境后直接运行：

# 单 GPU，nep_kokkos_lmps/kkin.lmp
mpirun -np 1 --bind-to numa lmp -k on g 1 -sf kk -pk kokkos -in kkin.lmp

# 单节点 4 GPU，nep_lmps/in.lmp-kk
mpirun -np 4 --bind-to numa lmp -k on g 4 -sf kk -pk kokkos -in in.lmp-kk

# 2 节点，每节点 4 GPU
mpirun -np 8 --bind-to numa --map-by ppr:4:node \
    lmp -k on g 4 -sf kk -pk kokkos -in in.lmp-kk

如果使用 MatPL Pro 闭源版本，运行前还需要设置 NEP_GPU_LIB_PATH 和 NEP_LICENSE_PATH。多节点运行时，所有节点必须能访问共享库、License、输入文件和力场文件。

MatPL 热流计算

接口区别：开放源码版本使用 matpl/heatflux；MatPL Pro 闭源版本使用 KOKKOS GPU 热流接口 matpl/heatflux/kk。请根据已安装的 lammps-MatPL 接口版本选择对应命令。

开放源码版本的热流 compute 设置为：

compute      flux all matpl/heatflux

MatPL Pro 闭源版本的 matpl/heatflux/kk 直接在 GPU 上计算 MatPL 热流，无需使用传统的 ke/atom + pe/atom + centroid/stress/atom + heat/flux 后处理链：

package kokkos neigh half comm device
newton on

pair_style   matpl/nep/kk nep.txt
pair_coeff   * * C

compute      flux all matpl/heatflux/kk
fix          fluxout all matpl/heatflux/ave/kk 10 100 1000 flux file compute_HeatFlux.out

compute matpl/heatflux/kk 输出包含 6 个分量的全局向量，依次为：

1. Jx：x 方向总热流。
2. Jy：y 方向总热流。
3. Jz：z 方向总热流。
4. Jconv,x：x 方向对流热流。
5. Jconv,y：y 方向对流热流。
6. Jconv,z：z 方向对流热流。

因此，virial 热流贡献可由总热流减去对流热流得到，即 (1-4, 2-5, 3-6)。fix matpl/heatflux/ave/kk 用于对热流向量进行时间平均，并将结果写入 compute_HeatFlux.out。

热流案例

• 开放源码版本 Heat_flux 案例
• MatPL Pro 闭源版本 Heat_flux 案例

ASE 接口

NEP 模型提供了 ase 接口，使用方式如下脚本例子所示gitee 或 github。

from src.ase.calculate import MatPL_calculator
calc = MatPL(model_file='nep_model.ckpt or nep.txt')
atoms = ..... # create ase.atoms.Atoms
atoms.calc = calc # or atoms.set_calculator(calc)
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
stress = atoms.get_stress()

注意，在使用本ase接口时确保已经导入了MatPL的环境变量。