機器學習勢函數（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