torch-sim程序: 为机器学习势模型量身打造的下一代原子模拟引擎

摘要：TorchSim 是一个以PyTorch 为核心重写原子模拟基础模块的创新引擎，它可以显著加速当今流行的机器学习势模型，如 MACE、Fairchem 和 SevenNet，相比传统工具如 ASE，在高性能 GPU 上可达百倍加速。

近日，Radical-AI 团队正式发布了全新的开源原子模拟引擎 —— TorchSim，它为机器学习势能（MLIP）时代的模拟需求提供了革命性的解决方案。

TorchSim 是一个以 PyTorch 为核心重写原子模拟基础模块的创新引擎，它可以显著加速当今流行的机器学习势模型，如 MACE、Fairchem 和 SevenNet，相比传统工具如 ASE，在高性能 GPU 上可达百倍加速。

只需不到 40 行代码，即可在 GPU 上运行批量的分子动力学模拟，并记录轨迹：

importtorch
importtorch_simasts
frommace.calculators.foundations_modelsimportmace_mp
fromtorch_sim.modelsimportMaceModel
fromase.buildimportbulk

device = torch.device("cuda")
mace = mace_mp(model="small", return_raw_model=True)
mace_model = MaceModel(model=mace, device=device)

cu_atoms = bulk("Cu""fcc", a=3.58).repeat((222))
many_cu_atoms = [cu_atoms] *50
trajectory_files = [f"Cu_traj_{i}"foriinrange(50)]

final_state = ts.integrate(
system=many_cu_atoms,
model=mace_model,
n_steps=50
timestep=0.002
temperature=1000
integrator=ts.nvt_langevin,
trajectory_reporter=dict(filenames=trajectory_files, state_frequency=10),
)

你可以通过 PyPI 或源码安装：

pip install torch-sim-atomistic

或从 GitHub 获取最新版本：

在 NVIDIA H100 (80GB) 上，TorchSim 可同时运行数千个原子系统，单原子模拟时间远优于 ASE，是构建大规模、高精度材料模拟平台的理想选择。