TorchSim/torch-sim

# TorchSim [](https://github.com/torchsim/torch-sim/actions/workflows/test.yml) [](https://codecov.io/gh/torchsim/torch-sim) [](https://python.org/downloads) [](https://pypi.org/project/torch-sim-atomistic) [][zenodo] TorchSim 是一款面向 MLIP 时代的下一代开源原子模拟引擎。通过使用 Pytorch 重写原子模拟的核心原语，它能够将流行的机器学习势函数加速数个数量级。 * 自动 batch 处理与 GPU 内存管理，可实现显著的模拟加速 * 支持 MACE、Fairchem、SevenNet、ORB、MatterSim、metatomic 和 Nequix MLIP 模型 * 支持经典的 Lennard-Jones、Morse 和 soft-sphere 势函数 * 分子动力学积分方案，如 NVE、NVT Langevin 和 NPT Langevin * 使用梯度下降和 FIRE 对原子位置和晶胞进行弛豫 * 交换 Monte Carlo 和混合交换 Monte Carlo 算法 * 可扩展的二进制轨迹写入格式，支持任意属性 * 为新用户提供简单直观的高阶 API * 与 ASE、Pymatgen 和 Phonopy 集成 * 以及更多功能：可微模拟、弹性属性、自定义 workflow... ## 快速开始 这里快速演示了 TorchSim 的许多核心功能： 原生支持 GPU、MLIP 模型、ASE 集成、简单的 API、 自动 batch 和轨迹报告，所有这些都在 40 行代码以内。 ### 运行批量 MD ``` import torch import torch_sim as ts # 在 gpus 上原生运行 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 轻松地从 mace-mp 加载模型 from mace.calculators.foundations_models import mace_mp from torch_sim.models.mace import MaceModel mace = mace_mp(model="small", return_raw_model=True) mace_model = MaceModel(model=mace, device=device) from ase.build import bulk cu_atoms = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2)) many_cu_atoms = [cu_atoms] * 50 trajectory_files = [f"Cu_traj_{i}.h5md" for i in range(len(many_cu_atoms))] # 通过 batching 同时运行它们 final_state = ts.integrate( system=many_cu_atoms, model=mace_model, n_steps=50, timestep=0.002, temperature=1000, integrator=ts.Integrator.nvt_langevin, trajectory_reporter=dict(filenames=trajectory_files, state_frequency=10), ) final_atoms_list = final_state.to_atoms() # 从 trajectory file 中提取最终能量 final_energies = [] for filename in trajectory_files: with ts.TorchSimTrajectory(filename) as traj: final_energies.append(traj.get_array("potential_energy")[-1]) print(final_energies) ``` ### 运行批量弛豫 然后使用 FIRE 弛豫这些结构只需要几行代码。 ``` # relax 所有高温状态 relaxed_state = ts.optimize( system=final_state, model=mace_model, optimizer=ts.Optimizer.fire, autobatcher=True, init_kwargs=dict(cell_filter=ts.CellFilter.frechet), ) print(relaxed_state.energy) ``` ## 加速 与使用热门 MLIP 的 ASE 相比，TorchSim 实现了高达 100 倍的加速。 此图比较了 ASE 和 `torch_sim` 的单原子耗时。单原子耗时定义为 原子数 / 总时间。虽然 ASE 只能运行单个 `n_atoms`（在 $x$ 轴上）的系统， 但 `torch_sim` 可以运行内存允许范围内的尽可能多的系统。在一张 H100 80 GB 显卡上， 内存能够容纳的最大原子数对于 [EGIP](https://github.com/FAIR-Chem/fairchem) 约为 8,000， 对于 [MACE-MPA-0](https://github.com/ACEsuit/mace) 约为 10,000，对于 [Mattersim V1 1M](https://github.com/microsoft/mattersim) 约为 22,000， 对于 [SevenNet](https://github.com/MDIL-SNU/SevenNet) 约为 2,500，对于 [PET-MAD](https://github.com/lab-cosmo/pet-mad) 约为 9000。 该指标通过同时捕获速度和内存使用情况来描述模型性能。 ## 安装 ### PyPI 安装 ``` pip install torch-sim-atomistic ``` ### 从源码安装 ``` git clone https://github.com/TorchSim/torch-sim cd torch-sim pip install . ``` ## 示例 要了解 TorchSim 的工作原理，请从文档中的[综合教程](https://torchsim.github.io/torch-sim/user/overview.html)开始。 ## 核心模块 TorchSim 的包结构在 [API 参考](https://torchsim.github.io/torch-sim/reference/index.html)文档中进行了总结，并在下方以矩形树图的形式展示。  ## 许可证 TorchSim 采用 [MIT 许可证](LICENSE)发布。 ## 引用 如果您在研究中使用了 TorchSim，请引用我们的[出版物](https://iopscience.iop.org/article/10.1088/3050-287X/ae1799)。 ``` @article{cohen2025torchsim, title={TorchSim: An efficient atomistic simulation engine in PyTorch}, author={Cohen, Orion and Riebesell, Janosh and Goodall, Rhys and Kolluru, Adeesh and Falletta, Stefano and Krause, Joseph and Colindres, Jorge and Ceder, Gerbrand and Gangan, Abhijeet S}, journal={AI for Science}, volume={1}, number={2}, pages={025003}, year={2025}, publisher={IOP Publishing}, doi={10.1088/3050-287X/ae1799} } ``` ## 致谢 我们旨在为 TorchSim 所基于的数十年工作致以所有的 [duecredit](https://github.com/duecredit/duecredit)。通过运行 `DUECREDIT_ENABLE=yes uv run --with-editable . --extra docs --extra test python -m duecredit <(printf 'import pytest\nraise SystemExit(pytest.main(["-q"]))\n')`，可以获取该包的自动参考文献列表。此列表尚不完整，我们欢迎提交 PR 来帮助改善我们的引用覆盖率。 要收集特定教程运行（例如 autobatching）的引用，请使用： ``` DUECREDIT_ENABLE=yes uv run --with-editable . --extra docs --extra test python -m duecredit examples/tutorials/autobatching_tutorial.py ```

标签：GPGPU计算, PyTorch, 凭据扫描, 分子动力学, 原子模拟, 机器学习原子间势, 科学与计算, 逆向工具