机器学习势函数（MLPF）入门：用DeePMD-kit加速亿级原子模拟

点击 “AladdinEdu，同学们用得起的【H卡】算力平台”，注册即送-H卡级别算力，80G大显存，按量计费，灵活弹性，顶级配置，学生更享专属优惠。

引言：从传统分子模拟到机器学习势函数的革命

分子动力学模拟是计算化学和材料科学中不可或缺的工具，它允许研究人员在原子尺度上研究材料的性质和行为。然而，传统分子动力学模拟面临着巨大的计算挑战：基于第一性原理的精确方法（如密度泛函理论DFT）计算成本极高，只能处理数百个原子的小系统；而经验势函数方法虽然可以处理更大体系，但精度有限且泛化能力差。

机器学习势函数（Machine Learning Potential Function, MLPF）的出现彻底改变了这一局面。通过机器学习方法从量子力学计算数据中学习势能面，MLPF既保持了接近第一性原理的精度，又能实现经典分子动力学的计算效率。在众多MLPF实现中，DeePMD-kit（Deep Potential Molecular Dynamics）凭借其出色的性能和可扩展性，成为了当前最热门的分子模拟加速方案。

本文将带你全面了解DeePMD-kit的工作流程，并通过一个完整的实例演示如何训练和部署机器学习势函数，最终实现亿级原子规模的分子动力学模拟。

1. 机器学习势函数与DeePMD-kit基础

1.1 机器学习势函数的核心思想

机器学习势函数的基本思想是用一个经过训练的神经网络来替代传统的经验势函数或量子力学计算。这个神经网络接收原子坐标和类型作为输入，输出系统的势能和原子力。

关键优势：

高精度：通过学习量子力学数据，保持接近第一性原理的精度
高效率：计算成本与系统大小呈线性关系，可处理超大体系
可迁移性：良好的泛化能力，适用于不同条件下的模拟

1.2 DeePMD-kit框架概述

DeePMD-kit是由深度势能团队开发的开源软件包，包含以下核心组件：

DeePMD：深度学习势函数模型
DP-GEN：自动生成训练数据的主动学习框架
DP Train：模型训练工具
DP Test：模型测试和验证工具
LAMMPS/DeePMD：支持DeePMD的分子动力学模拟引擎

1.3 环境准备与安装

1.3.1 通过Conda安装（推荐）

# 创建conda环境
conda create -n deepmd python=3.8
conda activate deepmd# 安装DeePMD-kit
conda install deepmd-kit=2.1.5=lammps -c conda-forge# 验证安装
dp -h

1.3.2 通过Docker安装

# 拉取官方镜像
docker pull ghcr.io/deepmodeling/deepmd-kit:2.1.5-cuda11.6# 运行容器
docker run -it --gpus all -v $(pwd):/workspace ghcr.io/deepmodeling/deepmd-kit:2.1.5-cuda11.6

1.3.3 源码编译安装

# 安装依赖
conda install tensorflow=2.8.0
conda install cmake make gcc# 克隆源码
git clone https://github.com/deepmodeling/deepmd-kit.git
cd deepmd-kit# 编译安装
mkdir build && cd build
cmake -DTENSORFLOW_ROOT=$CONDA_PREFIX ..
make -j8
make install

2. DeePMD-kit工作流程详解

2.1 完整工作流程概述

DeePMD-kit的典型工作流程包含四个主要阶段：

数据准备：通过第一性原理计算生成训练数据
模型训练：使用深度学习训练势函数模型
模型测试：验证模型的准确性和泛化能力
分子模拟：使用训练好的模型进行大规模MD模拟

2.2 数据准备阶段

2.2.1 第一性原理数据生成

首先我们需要用量子力学方法生成训练数据。以水分子为例：

# generate_qm_data.py
import ase
from ase import Atoms
from ase.calculators.psi4 import Psi4
import numpy as np# 创建水分子体系
def create_water_cluster(n_molecules=32):"""创建水分子团簇"""# 单个水分子坐标water = Atoms('H2O',positions=[(0, 0, 0),(0.757, 0.586, 0),(-0.757, 0.586, 0)])# 复制创建团簇cluster = Atoms()for i in range(n_molecules):pos = np.random.rand(3) * 10  # 随机位置water_copy = water.copy()water_copy.positions += poscluster += water_copyreturn cluster# 生成训练结构
structures = []
for i in range(100):  # 生成100个结构cluster = create_water_cluster(32)structures.append(cluster)# 使用Psi4进行DFT计算
calculator = Psi4(method='B3LYP', basis='6-31G*', num_threads=4, memory='4GB')# 计算每个结构的能量和力
qm_data = []
for i, atoms in enumerate(structures):print(f"计算结构 {i+1}/{len(structures)}")atoms.set_calculator(calculator)try:energy = atoms.get_potential_energy()forces = atoms.get_forces()qm_data.append({'coordinates': atoms.positions.copy(),'energy': energy,'forces': forces.copy(),'cell': atoms.cell.array.copy(),'atom_types': atoms.get_chemical_symbols()})except Exception as e:print(f"结构 {i} 计算失败: {e}")# 保存数据
np.save('qm_training_data.npy', qm_data)

2.2.2 数据格式转换

将量子力学数据转换为DeePMD格式：

# convert_to_deepmd.py
import numpy as np
from deepmd.utils.data import DataSystem# 加载QM数据
qm_data = np.load('qm_training_data.npy', allow_pickle=True)# 创建DeePMD格式的数据集
def create_deepmd_dataset(qm_data, output_dir):"""创建DeePMD格式的训练数据集"""import osos.makedirs(output_dir, exist_ok=True)# 创建类型文件atom_types = list(set(qm_data[0]['atom_types']))type_map = {symbol: i for i, symbol in enumerate(atom_types)}with open(os.path.join(output_dir, 'type_map.raw'), 'w') as f:for symbol in atom_types:f.write(f'{symbol}\n')# 保存每个结构的数据for i, data in enumerate(qm_data):struct_dir = os.path.join(output_dir, f'struct_{i:03d}')os.makedirs(struct_dir, exist_ok=True)# 坐标np.savetxt(os.path.join(struct_dir, 'coord.raw'), data['coordinates'])# 能量np.savetxt(os.path.join(struct_dir, 'energy.raw'), [data['energy']])# 力np.savetxt(os.path.join(struct_dir, 'force.raw'), data['forces'])# 晶胞np.savetxt(os.path.join(struct_dir, 'box.raw'), data['cell'].reshape(1, -1))# 原子类型type_indices = [type_map[symbol] for symbol in data['atom_types']]np.savetxt(os.path.join(struct_dir, 'type.raw'), type_indices, fmt='%d')# 转换数据
create_deepmd_dataset(qm_data, 'deepmd_training_data')

2.3 模型训练阶段

2.3.1 配置文件准备

创建训练配置文件input.json：

{"model": {"type_map": ["O", "H"],"descriptor": {"type": "se_e2_a","sel": [60, 120],"rcut": 6.0,"rcut_smth": 5.0,"neuron": [25, 50, 100],"axis_neuron": 16,"resnet_dt": true},"fitting_net": {"neuron": [240, 240, 240],"resnet_dt": true,"precision": "float64"}},"learning_rate": {"type": "exp","start_lr": 0.001,"decay_steps": 5000,"stop_lr": 1e-8},"loss": {"start_pref_e": 0.02,"limit_pref_e": 1,"start_pref_f": 1000,"limit_pref_f": 1,"start_pref_v": 0,"limit_pref_v": 0},"training": {"training_data": {"systems": ["deepmd_training_data"],"batch_size": 1,"atomic": false},"validation_data": {"systems": ["deepmd_validation_data"],"batch_size": 1,"atomic": false,"numb_btch": 1},"numb_steps": 1000000,"seed": 1,"disp_file": "lcurve.out","disp_freq": 1000,"save_freq": 10000}
}

2.3.2 开始模型训练

# 使用dp train进行训练
dp train input.json# 训练过程中可以监控损失曲线
tail -f lcurve.out# 冻结模型用于推理
dp freeze -o frozen_model.pb

2.3.3 训练过程监控

# monitor_training.py
import numpy as np
import matplotlib.pyplot as pltdef plot_training_curve(log_file='lcurve.out'):"""绘制训练曲线"""data = np.loadtxt(log_file)steps = data[:, 0]loss = data[:, 1]lr = data[:, 2]rmse_e = data[:, 3]rmse_f = data[:, 4]rmse_v = data[:, 5]fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))# 损失曲线ax1.semilogy(steps, loss)ax1.set_xlabel('Steps')ax1.set_ylabel('Loss')ax1.set_title('Training Loss')# 学习率ax2.semilogy(steps, lr)ax2.set_xlabel('Steps')ax2.set_ylabel('Learning Rate')ax2.set_title('Learning Rate Schedule')# 能量RMSEax3.plot(steps, rmse_e)ax3.set_xlabel('Steps')ax3.set_ylabel('RMSE Energy (eV)')ax3.set_title('Energy RMSE')# 力RMSEax4.plot(steps, rmse_f)ax4.set_xlabel('Steps')ax4.set_ylabel('RMSE Force (eV/Å)')ax4.set_title('Force RMSE')plt.tight_layout()plt.savefig('training_curves.png', dpi=300, bbox_inches='tight')plt.show()# 绘制训练曲线
plot_training_curve()

2.4 模型测试与验证

2.4.1 模型准确性测试

# 使用dp test测试模型
dp test -m frozen_model.pb -s deepmd_test_data -n 1000 -d results

2.4.2 验证脚本

# validate_model.py
import numpy as np
import matplotlib.pyplot as plt
from deepmd.infer import DeepPotdef validate_model(model_path, test_data):"""验证模型准确性"""# 加载模型dp = DeepPot(model_path)# 加载测试数据test_energies = []pred_energies = []test_forces = []pred_forces = []for i, data_dir in enumerate(test_data):# 加载测试数据coords = np.loadtxt(f'{data_dir}/coord.raw')energy = np.loadtxt(f'{data_dir}/energy.raw')forces = np.loadtxt(f'{data_dir}/force.raw')cell = np.loadtxt(f'{data_dir}/box.raw').reshape(3, 3)atom_types = np.loadtxt(f'{data_dir}/type.raw', dtype=int)# 模型预测pred_e, pred_f, _ = dp.eval(coords, cell, atom_types)test_energies.append(energy)pred_energies.append(pred_e[0])test_forces.append(forces.flatten())pred_forces.append(pred_f.flatten())# 计算误差test_energies = np.array(test_energies)pred_energies = np.array(pred_energies)test_forces = np.concatenate(test_forces)pred_forces = np.concatenate(pred_forces)energy_mae = np.mean(np.abs(test_energies - pred_energies))force_mae = np.mean(np.abs(test_forces - pred_forces))energy_rmse = np.sqrt(np.mean((test_energies - pred_energies)**2))force_rmse = np.sqrt(np.mean((test_forces - pred_forces)**2))print(f"能量MAE: {energy_mae:.4f} eV")print(f"能量RMSE: {energy_rmse:.4f} eV")print(f"力MAE: {force_mae:.4f} eV/Å")print(f"力RMSE: {force_rmse:.4f} eV/Å")# 绘制散点图fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))# 能量散点图ax1.scatter(test_energies, pred_energies, alpha=0.6)ax1.plot([min(test_energies), max(test_energies)], [min(test_energies), max(test_energies)], 'r--')ax1.set_xlabel('DFT Energy (eV)')ax1.set_ylabel('DeePMD Energy (eV)')ax1.set_title('Energy Comparison')# 力散点图ax2.scatter(test_forces, pred_forces, alpha=0.1, s=1)force_range = [min(min(test_forces), min(pred_forces)), max(max(test_forces), max(pred_forces))]ax2.plot(force_range, force_range, 'r--')ax2.set_xlabel('DFT Force (eV/Å)')ax2.set_ylabel('DeePMD Force (eV/Å)')ax2.set_title('Force Comparison')plt.tight_layout()plt.savefig('validation_results.png', dpi=300, bbox_inches='tight')plt.show()return energy_mae, force_mae# 验证模型
test_dirs = [f'deepmd_test_data/struct_{i:03d}' for i in range(50)]
validate_model('frozen_model.pb', test_dirs)

3. 大规模分子动力学模拟

3.1 LAMMPS集成

3.1.1 准备LAMMPS输入文件

创建lammps.in输入文件：

# LAMMPS输入文件 - 水分子模拟units metal
atom_style atomic
timestep 0.0005# 读取初始结构
read_data water_system.data# DeePMD势函数
pair_style deepmd frozen_model.pb
pair_coeff * *# 温度控制
velocity all create 300.0 12345
fix nvt all nvt temp 300.0 300.0 0.1# 输出设置
thermo_style custom step temp pe ke etotal press
thermo 1000
dump traj all atom 10000 water_trajectory.xyz# 运行模拟
run 1000000

3.1.2 准备初始结构

# create_water_system.py
import numpy as npdef create_large_water_system(n_molecules=100000, output_file='water_system.data'):"""创建大规模水分子系统"""# 单个水分子坐标water_coords = np.array([[0.0, 0.0, 0.0],          # 氧原子[0.757, 0.586, 0.0],      # 氢原子1[-0.757, 0.586, 0.0]      # 氢原子2])# 计算系统尺寸box_size = (n_molecules * 30.0) ** (1/3)  # 估算盒子大小n_per_side = int(np.ceil(n_molecules ** (1/3)))# 生成原子坐标coordinates = []atom_types = []count = 0for i in range(n_per_side):for j in range(n_per_side):for k in range(n_per_side):if count >= n_molecules:break# 随机位置pos = np.array([i, j, k]) * (box_size / n_per_side)pos += np.random.rand(3) * 0.1  # 小随机偏移# 添加水分子for atom_idx in range(3):coord = water_coords[atom_idx] + poscoordinates.append(coord)# 原子类型：1=氧, 2=氢if atom_idx == 0:atom_types.append(1)else:atom_types.append(2)count += 1coordinates = np.array(coordinates)# 写入LAMMPS数据文件with open(output_file, 'w') as f:f.write(f"LAMMPS data file for water system\n\n")f.write(f"{len(coordinates)} atoms\n")f.write("2 atom types\n\n")f.write(f"0.0 {box_size} xlo xhi\n")f.write(f"0.0 {box_size} ylo yhi\n")f.write(f"0.0 {box_size} zlo zhi\n\n")f.write("Masses\n\n")f.write("1 15.9994\n")  # 氧原子f.write("2 1.00794\n")  # 氢原子f.write("\n")f.write("Atoms\n\n")for i, (coord, atom_type) in enumerate(zip(coordinates, atom_types), 1):f.write(f"{i} {atom_type} {coord[0]} {coord[1]} {coord[2]}\n")print(f"创建了包含 {n_molecules} 个水分子 ({len(coordinates)} 个原子) 的系统")return box_size# 创建包含10万个水分子的系统
box_size = create_large_water_system(100000)

3.2 运行大规模模拟

3.2.1 使用LAMMPS运行模拟

# 使用MPI并行运行LAMMPS
mpirun -np 64 lmp -in lammps.in -log simulation.log

3.2.2 性能优化配置

对于亿级原子模拟，需要优化运行配置：

# 使用GPU加速
mpirun -np 8 lmp -sf gpu -pk gpu 4 -in lammps.in# 使用多节点并行
mpirun -hostfile hostfile -np 256 lmp -in lammps.in

3.3 结果分析与可视化

3.3.1 轨迹分析

# analyze_trajectory.py
import numpy as np
import matplotlib.pyplot as plt
from ase.io import readdef analyze_water_trajectory(traj_file='water_trajectory.xyz'):"""分析水分子轨迹"""# 读取轨迹trajectory = read(traj_file, index=':')# 提取分析数据times = []temperatures = []energies = []pressures = []for i, atoms in enumerate(trajectory):# 从注释行提取热力学数据comment = atoms.info.get('comment', '')if comment:parts = comment.split()if len(parts) >= 7:times.append(i * 0.0005)  # 时间 (ps)temperatures.append(float(parts[2]))  # 温度 (K)energies.append(float(parts[3]))  # 势能 (eV)pressures.append(float(parts[6]))  # 压力 (bar)# 绘制热力学性质fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))# 温度ax1.plot(times, temperatures)ax1.set_xlabel('Time (ps)')ax1.set_ylabel('Temperature (K)')ax1.set_title('Temperature Evolution')# 能量ax2.plot(times, energies)ax2.set_xlabel('Time (ps)')ax2.set_ylabel('Potential Energy (eV)')ax2.set_title('Energy Evolution')# 压力ax3.plot(times, pressures)ax3.set_xlabel('Time (ps)')ax3.set_ylabel('Pressure (bar)')ax3.set_title('Pressure Evolution')# 径向分布函数if len(trajectory) > 100:from ase.ga.utilities import get_rdfrdf, distances = get_rdf(trajectory[-100:], 200, (0, 0))ax4.plot(distances, rdf)ax4.set_xlabel('Distance (Å)')ax4.set_ylabel('g(r)')ax4.set_title('Oxygen-Oxygen RDF')plt.tight_layout()plt.savefig('simulation_analysis.png', dpi=300, bbox_inches='tight')plt.show()return times, temperatures, energies, pressures# 分析轨迹
analyze_water_trajectory()

3.3.2 性能统计

# performance_analysis.py
import re
import matplotlib.pyplot as pltdef analyze_performance(log_file='simulation.log'):"""分析模拟性能"""# 读取日志文件with open(log_file, 'r') as f:log_content = f.read()# 提取性能数据steps = []times = []performances = []# 匹配性能信息pattern = r'Step Temp.*\n.*\n.*\n.*\n.*\n.*\n.*\n.*\n.*\n.*\n.*\n.*Loop time of (\d+\.\d+) on (\d+) procs'matches = re.findall(pattern, log_content)for match in matches:loop_time = float(match[0])n_procs = int(match[1])performance = 1.0 / loop_time  # 步/秒performances.append(performance)# 绘制性能曲线plt.figure(figsize=(10, 6))plt.plot(performances)plt.xlabel('Output Step')plt.ylabel('Performance (steps/second)')plt.title('MD Simulation Performance')plt.grid(True)plt.savefig('performance_analysis.png', dpi=300, bbox_inches='tight')plt.show()avg_performance = np.mean(performances)print(f"平均性能: {avg_performance:.2f} 步/秒")print(f"相当于 {avg_performance * 3600:.0f} 步/小时")return performances# 分析性能
performances = analyze_performance()

4. 高级主题与最佳实践

4.1 DP-GEN主动学习

对于更复杂的系统，可以使用DP-GEN进行主动学习：

# 安装DP-GEN
conda install dp-gen -c conda-forge# 准备DP-GEN配置文件
cat > param.json << EOF
{"type_map": ["O", "H"],"mass_map": [15.999, 1.008],"init_data_sys": ["deepmd_training_data"],"sys_configs": [["water_system.data", "frozen_model.pb"]],"numb_models": 4,"traj_freq": 10,"temps": [300.0],"pressures": [1.0],"explore_stride": 1.0,"model_devi_f_trust_lo": 0.10,"model_devi_f_trust_hi": 0.30,"model_devi_v_trust_lo": 0.002,"model_devi_v_trust_hi": 0.010
}
EOF# 运行DP-GEN
dpgen run param.json machine.json

4.2 多组件系统建模

对于包含多种元素的应用：

{"model": {"type_map": ["Cu", "O", "H"],"descriptor": {"type": "se_e2_a","sel": [100, 60, 120],"rcut": 6.0,"rcut_smth": 5.0}}
}

4.3 性能优化技巧

混合精度训练：使用float16精度减少内存使用
梯度累积：处理大批次大小
分布式训练：多GPU或多节点训练
模型压缩：减少模型大小以提升推理速度

5. 实际应用案例

5.1 金属材料模拟

# metal_simulation.py
def setup_metal_simulation():"""设置金属材料模拟"""# 创建面心立方铜晶体from ase.build import fcc111, add_adsorbatefrom ase import Atoms# 创建基底slab = fcc111('Cu', size=(4, 4, 4), vacuum=10.0)# 添加吸附物oxygen = Atoms('O', positions=[(0, 0, 0)])add_adsorbate(slab, oxygen, height=2.0, position='fcc')return slab# 生成训练数据并训练金属势函数

5.2 界面反应研究

# interface_study.py
def study_interface_reactions():"""研究界面反应"""# 创建金属-水界面系统from ase.build import fcc100, add_adsorbate, moleculefrom ase import Atoms# 金属基底metal_slab = fcc100('Cu', size=(6, 6, 4), vacuum=15.0)# 添加水分子层for i in range(3):for j in range(3):water = molecule('H2O')water.rotate(180, 'x')add_adsorbate(metal_slab, water, height=3.0, position=(i*2.5, j*2.5))return metal_slab