aiqm/torchani

TorchANI 2.0 – PyTorch 中的神经网络原子间势能

TorchANI 是一个基于 PyTorch 构建的开源 Python 库，允许研究人员训练、评估和使用 ANI 风格的神经网络原子间势能 (NN-IP)。这些是机器学习模型，可以直接根据原子的坐标预测分子的势能和作用力，从而实现快速、准确的分子动力学 (MD) 和混合 ML/MM 模拟。

库的功能

• 模型库 – 提供即插即用的预训练 ANI 模型以及定义新模型的简洁 API。
• 训练工具 – 将 NN-IP 拟合到量子化学参考数据的工具。
• 快速推理 – 通过 C++/CUDA 扩展计算 ANI 描述符并在 GPU 上运行网络，实现近乎实时的 MD。
• 命令行界面 – ani 可执行文件提供快速工具（例如：能量/力评估、构建扩展）。
• 集成点 – 独立的 TorchANI-Amber 软件包允许您将势能引入 Amber 的经典 MD 引擎，用于 ML/MM 模拟。

适用人群

• 需要用于模拟的高性能、基于 ML 的力场的计算化学家或材料科学家。
• 正在开发新的 NN-IP 架构并希望拥有 PyTorch 原生基准的研究人员。
• 希望在无需编写底层 CUDA 代码的情况下，在 GPU 上运行 ML 驱动的分子动力学的任何人。

入门指南 (摘自 README)

1. 设置 Python 环境 (conda/mamba 或 virtualenv)。该库在 PyPI 上发布，但作者建议即使在 conda 环境中也使用 pip 进行安装，因为 conda 软件包已不再维护。
2. 安装兼容的 PyTorch (例如，针对 CUDA 13.0，使用 pip install torch==2.13 --index-url https://download.pytorch.org/whl/cu130)。TorchANI 2 需要 PyTorch ≥ 2.0。
3. 安装 TorchANI: pip install torchani。
4. 为描述符计算和快速推理构建可选的 C++/CUDA 扩展:

   ani build-extensions          # 为所有检测到的 GPU SM 构建
   # 或限制于特定 SM，例如：
   ani build-extensions --sm 8.0 --sm 8.9
   

5. （可选）如果您需要构建文档、运行测试或修改源码，请克隆仓库并通过提供的 environment.yaml 创建开发环境。

平台说明

• GPU 支持：适用于带有 CUDA 的 NVIDIA GPU；对于任何非平凡的工作负载，强烈建议使用。未测试 AMD/ROCm。
• macOS：无 CUDA；该库在 Apple 的 MPS 后端上未经测试。
• Conda 软件包：面向开发者提供；可以使用 recipe 目录在本地构建。

文档与帮助

• 完整的用户指南和 API 参考：https://aiqm.github.io/torchani/
• 从 TorchANI 1.x 迁移到 2.0 的迁移指南。
• 欢迎通过 GitHub issues 提交错误报告和功能请求。

引用 TorchANI

如果您在发表的著作中使用该库，请引用 README 中列出的两篇论文（原始 TorchANI 论文和 TorchANI 2.0 扩展论文）。

---

摘要 (TL;DR): TorchANI 2.0 是一个基于 PyTorch 的工具包，用于构建和运行分子模拟的神经网络势能，提供 GPU 加速推理、训练工具以及与 Amber 等经典 MD 软件包的集成。这是一个面向计算化学界的真正的 AI/ML 软件项目。