AI²BMD：人工智能驱动的从头计算生物分子动力学模拟

目录

• 概述
• 快速入门
• 数据集
• 系统要求
• 高级设置
• 相关研究
• 引用
• 许可证
• 免责声明
• 联系方式

概述

AI²BMD 是一款能够以从头计算精度高效模拟蛋白质分子动力学的程序。本仓库包含该模拟程序、数据集以及与 AI²BMD 相关的公开资料。AI²BMD 的主要内容已发表在《Nature》期刊上（https://www.nature.com/articles/s41586-024-08127-z）。

以下动画展示了 AI²BMD 的工作原理。

https://github.com/user-attachments/assets/912a3e5a-c465-4dc7-8c2d-9f7807cac2a7

快速入门

AI²BMD 的源代码托管于此仓库中。 我们已将源代码和运行时库打包成 Docker 镜像，并提供一个 Python 启动脚本，以简化安装流程。 要运行该模拟程序，您无需克隆此仓库。只需下载 scripts/ai2bmd 并执行即可（需 Python >=3.7 和 Docker 环境）。

我们可以按照如下步骤运行一次分子动力学模拟：

# 如果已设置好启动脚本，可跳过以下两行
wget 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/scripts/ai2bmd'
chmod +x ai2bmd
# 下载 Chignolin 蛋白质结构数据文件
wget 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/examples/chig.pdb'
# 下载预处理并溶剂化的 Chignolin 蛋白质结构数据文件
wget --directory-prefix=chig_preprocessed 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/examples/chig_preprocessed/chig-preeq.pdb'
wget --directory-prefix=chig_preprocessed 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/examples/chig_preprocessed/chig-preeq-nowat.pdb'
# 从容器注册表拉取 Docker 镜像
docker pull ghcr.io/microsoft/ai2bmd:latest
# 启动程序，所有模拟参数均采用默认值
# 如果用户未加入 docker 组，可能需要使用 sudo 执行以下命令
./ai2bmd --prot-file chig.pdb --preprocess-dir chig_preprocessed --preeq-steps 0 --sim-steps 1000 --record-per-steps 1

这里我们以一个非常简单的蛋白质 Chignolin 作为示例。 程序将使用默认参数进行模拟。

结果将保存在一个名为 Logs-chig 的新目录中。该目录包含模拟轨迹文件：

• chig-traj.traj：ASE 二进制格式的完整轨迹文件。

注意：目前，AI²BMD 支持具有中性末端封端（ACE 和 NME）、单链且由标准氨基酸组成的蛋白质的 MD 模拟。

数据集

蛋白质单元数据集

蛋白质单元数据集涵盖了约 2000 万个在 DFT 水平下计算得到的二肽构象。可通过以下命令下载：

# 如果已设置好启动脚本，可跳过以下两行
wget 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/scripts/ai2bmd'
chmod +x ai2bmd
# 如果用户未加入 docker 组，可能需要使用 sudo 执行以下命令
./ai2bmd --download-training-data

下载完成后，当前工作目录中将生成多个 numpy 数据文件 (*.npz)。

AIMD-Chig 数据集

AIMD-Chig 数据集包含 166 个原子的 Chignolin 蛋白质的 200 万个构象，以及它们对应的势能和原子力，这些数据均由密度泛函理论 (DFT) 在 M06-2X/6-31G* 水平下计算得出。

• 阅读文章 [AIMD-Chig：利用从头计算分子动力学探索 166 原子蛋白质 Chignolin 的构象空间]（https://www.nature.com/articles/s41597-023-02465-9）。

• 查看报道 [首个达到从头计算精度的蛋白质完整构象分子动力学数据集及其背后的新型计算技术]（https://bioengineeringcommunity.nature.com/posts/aimd-chig-exploring-the-conformational-space-of-proteins-at-dft-level）。

• 获取数据集 [AIMD-Chig]（https://figshare.com/articles/dataset/_strong_AIMD-Chig_exploring_the_conformational_space_of_166-atom_protein_strong_em_strong_Chignolin_strong_em_strong_with_strong_em_strong_ab_initio_strong_em_strong_molecular_dynamics_strong_/22786730）。

系统要求

硬件要求

AI²BMD 程序可在 x86-64 架构的 GNU/Linux 系统上运行。 我们建议使用具备以下配置的机器：

• CPU：8 核及以上
• 内存：32 GB 及以上
• GPU：支持 CUDA 的 GPU，显存 8 GB 及以上

该程序已在以下 GPU 上测试通过：

• A100
• V100
• RTX A6000
• Titan RTX

软件要求

该程序已在以下系统上测试通过：

• 操作系统：Ubuntu 20.04，Docker：27.1
• 操作系统：ArchLinux，Docker：26.1

高级设置

环境配置

为方便实用，运行时库及依赖项已被打包至 Docker 镜像中。在启动 Docker 镜像之前，您需要安装 Docker 软件（详情请参阅 https://docs.docker.com/engine/install/），并将当前用户添加到 docker 组中，具体操作如下：

sudo groupadd docker
sudo usermod -aG docker $USER
newgrp docker

蛋白质文件准备

AI²BMD 的输入文件应为 .pdb 格式。如果 .pdb 文件中缺少氢原子，需添加氢原子。随后，应在蛋白质的 N 端用 ACE（乙酰基）封端，在 C 端用 NME（N-甲基）封端。这些步骤可使用 PyMOL 软件通过以下命令高效完成。

from pymol import cmd
pymol.finish_launching()
cmd.load("your_protein.pdb","molecule")
cmd.h_add("molecule") # 添加氢原子

cmd.wizard("mutagenesis")
cmd.get_wizard().set_n_cap("acet")
selection = "/%s//%s/%s" % (molecule, chain, resi) # 选择 N 端
cmd.get_wizard().do_select(selection)
cmd.get_wizard().apply()

cmd.get_wizard().set_c_cap("nmet")
selection = "/%s//%s/%s" % (molecule, chain, resi) # 选择 N 端
cmd.get_wizard().do_select(selection)
cmd.get_wizard().apply()

cmd.set_wizard()

接下来，可以使用 AmberTools 的 pdb4amber 工具调整 .pdb 文件中的原子名称，特别是确保 ACE 和 NME 的命名符合 ai2bmd 的要求。ACE 和 NME 的原子名称应符合以下规范：

• ACE：C、O、CH3、H1、H2、H3
• NME：N、CH3、H、HH31、HH32、HH33

pdb4amber -i your_protein.pdb -o processed_your_protein.pdb

此外，请确认蛋白质链中不存在 TER 分隔符，并且残基编号应从 1 开始，不得有间隙。

完成上述步骤后，您的 .pdb 文件应类似于以下格式：

ATOM      1  H1  ACE     1      10.845   8.614   5.964  1.00  0.00           H
ATOM      2  CH3 ACE     1      10.143   9.373   5.620  1.00  0.00           C
ATOM      3  H2  ACE     1       9.425   9.446   6.437  1.00  0.00           H
ATOM      4  H3  ACE     1       9.643   9.085   4.695  1.00  0.00           H
ATOM      5  C   ACE     1      10.805  10.740   5.408  1.00  0.00           C
ATOM      6  O   ACE     1      10.682  11.417   4.442  1.00  0.00           O
...
ATOM    170  N   NME    12       9.499   8.258  10.367  1.00  0.00           N
ATOM    171  H   NME    12       9.393   8.028  11.345  1.00  0.00           H
ATOM    172  CH3 NME    12       8.845   7.223   9.569  1.00  0.00           C
ATOM    173 HH31 NME    12       7.842   6.990   9.925  1.00  0.00           H
ATOM    174 HH32 NME    12       8.798   7.589   8.543  1.00  0.00           H
ATOM    175 HH33 NME    12       9.418   6.305   9.435  1.00  0.00           H
END

您也可以参考 examples 文件夹中的蛋白质文件。需要注意的是，目前机器学习势能对含有二硫键的蛋白质支持尚不完善，我们将在近期进行更新。

预处理

在预处理阶段，系统会先构建溶剂化体系，然后进行能量最小化和若干预平衡阶段。目前，AI²MD 通过参数 preprocess_method 提供两种预处理方法。

如果您选择 FF19SB 方法，系统将依次经历溶剂化、能量最小化、加热以及多个预平衡阶段。为了利用多核 CPU 和 GPU 加速预处理过程，您需要获取 AMBER 软件包，并修改 src/AIMD/preprocess.py 中的相关命令。

如果您选择 AMOEBA 方法，系统将仅进行溶剂化和能量最小化阶段。我们强烈建议您进行预平衡模拟，以使模拟系统充分松弛。

模拟

AI²BMD 通过参数 mode 提供两种生产模拟模式。默认的 fragment 模式表示蛋白质会被分割成二肽，然后在每一步模拟中由机器学习势能进行计算。

AI²BMD 也支持用户自行训练机器学习势能，并在不进行碎片化的情况下进行模拟。visnet 模式表示蛋白质的整体势能和原子力将由 ViSNet 模型直接计算，无需进行碎片化。使用此模式时，您需要自行使用分子数据训练 ViSNet 模型，将其上传至 src/ViSNet，并为参数 ckpt-type 设置相应的值。这样，您就可以使用 AI²BMD 模拟程序来模拟除蛋白质以外的任何分子。如需自行训练 ViSNet 模型，请参阅分支 ViSNet，其中包含源代码、模型训练说明及其他技术细节。

要执行包括预处理在内的完整 AI²BMD 模拟，可参考以下命令。

wget 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/scripts/ai2bmd'
chmod +x ai2bmd
# 如果用户未配置 docker 组，可能需要使用 sudo 运行以下命令
./ai2bmd --prot-file path/to/target-protein.pdb --sim-steps nnn  ...
#        '-------- 必填参数 ---------' '-- 可选参数 --'
#
# 常见可选参数：
#
# 【模拟目录映射选项】
#   --base-dir path/to/base-dir    模拟运行目录（默认为当前目录）
#   --log-dir  path/to/log-dir     日志及结果保存目录（默认为 base-dir/Logs-蛋白质名）
#   --src-dir  path/to/src-dir     将本地 src 目录挂载到仓库的 src/ 目录位置（默认不启用）
#
# 【模拟参数选项】
#   --sim-steps nnn                模拟步数
#   --temp-k nnn                   模拟温度（单位：开尔文）
#   --timestep nnn                 模拟时间步长（单位：飞秒）
#   --preeq-steps nnn              每个约束条件下的预平衡模拟步数
#   --max-cyc nnn                  预处理中能量最小化的最大循环次数
#   --preprocess-method [method]   预处理方法
#   --mode [mode]                  模拟过程中是否进行碎片化
#   --record-per-steps nnn         轨迹保存频率
#
# 【性能优化选项】
#   --device-strategy [strategy]   计算设备分配策略
#       excess-compute                 保留最后一块 GPU 用于非键合及溶剂相关计算
#       small-molecule                 最大化资源用于模型推理
#       large-molecule                 提升大分子模拟性能
#   --chunk-size nnn               每批处理的原子数（减少内存消耗）
#
# 【其他启动选项】
#   --software-update              若指定，则在运行前更新 Docker 镜像中的程序
#   --download-training-data       若指定，则下载 AI2BMD 训练数据并解压到工作目录，忽略其他所有选项。
#   --gpus                         指定传递给程序的 GPU 设备，可选：
#                                  all:        将所有可用 GPU 传递给程序。

#                                  none:       禁用 GPU 直通。
#                                  i[,j,k...]  直通部分 GPU。示例：--gpus 0,1

后处理分析

模拟轨迹的格式为 ASE 格式的 .traj 文件。若需将其转换为 .dcd 格式以便进行可视化，可先安装 MDAnalysis，并以 src/utils/traj2dcd.py 脚本作为参考，执行以下命令：

python traj2dcd.py --input xxx.pdb --output xxx.dcd --pdb xxx.pdb --num-atoms nnn --stride nnn

# 参数说明：
# --input         输入轨迹文件的名称
# --output        输出轨迹文件的名称
# --pdb           与输入轨迹对应的参考 PDB 文件
# --num-atoms     蛋白质或整个溶剂化体系的原子数
# --stride        输出轨迹的频率

故障排除

由于对质子跳跃行为的建模不足或模拟体系设置不当，模拟可能会崩溃。质子跳跃现象在大型生物分子和长时间模拟中尤为常见。鉴于机器学习势能训练过程中已出现过若干质子跳跃案例，模型可能遇到“分布外”情况，从而给出错误的原子作用力，最终导致模拟崩溃。我们将持续更新 AI²BMD 势能，提升其预测能力；同时也强烈建议大家贡献数据集，以进一步微调模型。

为避免或缓解模拟崩溃，我们提供以下建议：1) 在正式生产模拟运行前充分弛豫模拟体系；2) 延长预平衡模拟的时间；3) 在正式生产模拟中增加施加约束条件的模拟时长（通过参数 preeq-steps）；4) 从崩溃前几步处重新启动模拟；5) 扩大溶剂盒子的尺寸；6) 调整其他模拟体系设置。

除了直接进行模拟之外，我们也鼓励用户利用 AI²BMD 对现有模拟轨迹进行重加权，并据此计算蛋白质的相关性质。

相关研究

模型架构

ViSNet

ViSNet（Vector-Scalar interactive graph neural Network）是一种等变的、基于几何信息增强的分子图神经网络，能够显著缓解计算成本与充分利用几何信息之间的矛盾。

• ViSNet 已发表在《Nature Communications》上，论文题目为《通过等变向量-标量交互消息传递增强分子的几何表征》（链接）。

• ViSNet 同时入选了《Nature Communications》“编辑精选”栏目，分别在“人工智能与机器学习”（链接）和“生物技术与方法”（链接）两个领域。

• ViSNet 在【首届全球AI药物研发大赛】中荣获冠军，并在【OGB-LSC @ NeurIPS 2022 PCQM4Mv2赛道】中跻身获奖者行列！

• 欢迎访问分支 ViSNet，获取源代码、模型训练说明及更多技术细节。

Geoformer

Geoformer（Geometric Transformer）是一种新颖的几何Transformer模型，用于高效建模分子结构，以实现多种分子性质的预测。Geoformer 引入了一种新的位置编码方法——原子间位置编码（IPE），用于在Transformer架构中参数化原子环境。通过引入IPE，Geoformer能够在基于Transformer的架构中捕捉到超越两两距离的宝贵几何信息。Geoformer可以被视为ViSNet的一种Transformer变体。

• Geoformer 已发表于NeurIPS 2023。
• 请阅读Geoformer的论文《带有原子间位置编码的几何Transformer》（链接）。
• 欢迎访问分支 Geoformer，获取源代码、模型训练说明及更多技术细节。

面向MLFF的细粒度力场指标

近年来，机器学习力场（MLFF）作为一种经济高效的替代方案，逐渐取代从头算分子动力学（MD）模拟而受到广泛关注。然而，尽管MLFF在测试集上的误差较小，但在实际的MD模拟过程中，其泛化能力和鲁棒性往往存在明显不足。

为解决这些问题，我们提出了一种综合性的评估框架，通过全局力场指标以及从元素和构象角度出发的细粒度指标，系统地衡量MLFF在每个原子和每种分子构象上的表现。此外，在模型训练过程中引入这些力场指标，可以进一步提升MLFF的性能和MD模拟的稳定性。具体方法包括：将这些力场指标作为损失函数来训练MLFF模型；通过对原始数据集中的样本进行重加权来进行微调；以及结合额外的未探索数据继续训练。

• 请阅读封面文章《利用细粒度力场指标改进用于分子动力学模拟的机器学习力场》（链接）。

马尔可夫状态模型的随机延迟时间参数化

马尔可夫状态模型（MSM）在研究蛋白质构象动力学中起着关键作用。通常，人们会使用一个固定延迟时间的滑动计数窗口来采样子轨迹，以便进行转移计数和构建MSM。然而，采用固定延迟时间采样的子轨迹在不同延迟时间选择下可能表现不佳，这不仅需要丰富的先验经验，还可能导致估计结果不够稳健。

为此，我们提出了一种基于泊松过程的新型随机方法，用于生成扰动式的延迟时间来采样子轨迹，并以此构建马尔可夫链。针对双势阱系统、WW结构域、BPTI以及SARS-CoV-2的RBD–ACE2复合物等复杂体系的全面评估表明，我们的算法能够在不破坏马尔可夫性质的前提下，显著提高所构建MSM的稳健性和准确性。此外，该算法的优势在复杂生物过程中的慢动态模式中尤为突出。

• 请阅读封面文章《蛋白质动力学马尔可夫状态模型的随机延迟时间参数化》（链接）。
• 另外，您还可以参考一项应用案例：通过研究Omicron变异株的Spike-ACE2复合物结构，揭示其高传染性机制——《关于SARS-CoV-2 Omicron RBD-ACE2相互作用的结构见解》（链接）。

引用

(#：共同第一作者；*：通讯作者）

Tong Wang#*, Xinheng He#, Mingyu Li#, Yatao Li#, Ran Bi, Yusong Wang, Chaoran Cheng, Xiangzhen Shen, Jiawei Meng, He Zhang, Haiguang Liu, Zun Wang, Shaoning Li, Bin Shao*, Tie-Yan Liu. 利用AI²BMD对蛋白质分子动力学进行从头计算表征。Nature 2024。

Yusong Wang#, Tong Wang#*, Shaoning Li#, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao*, Tie-Yan Liu. 通过等变向量-标量交互消息传递增强分子的几何表征，Nature Communications, 15.1 (2024): 313。

Tong Wang#*, Xinheng He#, Mingyu Li#, Bin Shao*, Tie-Yan Liu. AIMD-Chig：利用从头算分子动力学探索166个原子的Chignolin蛋白构象空间，Scientific Data 10, 549 (2023)。

Yusong Wang#, Shaoning Li#, Tong Wang*, Bin Shao, Nanning Zheng, Tie-Yan Liu. 带有原子间位置编码的几何Transformer。NeurIPS 2023。

Zun Wang#, Hongfei Wu#, Lixin Sun, Xinheng He, Zhirong Liu, Bin Shao, Tong Wang*, Tie-Yan Liu. 利用细粒度力场指标改进用于分子动力学模拟的机器学习力场，The Journal of Chemical Physics，第159卷，第3期，封面文章。

Shiqi Gong#, Xinheng He#, Qi Meng, Zhiming Ma, Bin Shao*, Tong Wang*, Tie-Yan Liu. 蛋白质动力学马尔可夫状态模型的随机延迟时间参数化，The Journal of Physical Chemistry B 2022年第126卷第46期，封面文章，2022年。

许可证

版权所有© 微软公司。保留所有权利。

根据MIT许可证授权。

免责声明

AI²BMD 是一项研究项目。它并非微软官方支持的产品。

联系方式

如有任何问题、建议或技术支持，请联系 Tong Wang（项目负责人）和 Ran Bi。

AI2BMD 快速上手指南

AI2BMD 是一款基于人工智能的从头算（ab initio）生物分子动力学模拟程序，能够以量子力学精度高效模拟蛋白质动力学。本指南将帮助您快速在本地环境中部署并运行该工具。

环境准备

系统要求

• 操作系统: x86-64 GNU/Linux (推荐 Ubuntu 20.04 或 ArchLinux)
• CPU: 8 核及以上
• 内存: 32 GB 及以上
• GPU: 支持 CUDA 的显卡，显存 8 GB 及以上
  • 已测试型号: A100, V100, RTX A6000, Titan RTX

前置依赖

• Python: 版本 >= 3.7
• Docker: 已安装并配置好用户权限
  • 若未配置 Docker 用户组，请执行以下命令：

    sudo groupadd docker
    sudo usermod -aG docker $USER
    newgrp docker
    

安装步骤

AI2BMD 采用 Docker 容器化部署，无需克隆整个代码仓库，只需下载启动脚本即可。

1. 下载启动脚本并赋予执行权限：

   wget 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/scripts/ai2bmd'
   chmod +x ai2bmd
   

2. 拉取 Docker 镜像： 首次运行时脚本会自动拉取，也可手动预拉取：

   docker pull ghcr.io/microsoft/ai2bmd:latest
   

基本使用

以下示例演示如何对一个小蛋白 Chignolin 进行简单的动力学模拟。

1. 准备数据文件

下载示例蛋白结构文件及预处理文件：

# 下载原始蛋白结构
wget 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/examples/chig.pdb'

# 下载预处理后的溶剂化结构文件
mkdir -p chig_preprocessed
wget --directory-prefix=chig_preprocessed 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/examples/chig_preprocessed/chig-preeq.pdb'
wget --directory-prefix=chig_preprocessed 'https://raw.githubusercontent.com/microsoft/AI2BMD/main/examples/chig_preprocessed/chig-preeq-nowat.pdb'

2. 运行模拟

执行以下命令启动模拟（使用默认参数，模拟 1000 步）：

./ai2bmd --prot-file chig.pdb --preprocess-dir chig_preprocessed --preeq-steps 0 --sim-steps 1000 --record-per-steps 1

注意：如果当前用户未加入 docker 组，可能需要在命令前加 sudo。

3. 查看结果

模拟完成后，结果将保存在 Logs-chig 目录中：

• chig-traj.traj: ASE 二进制格式的全轨迹文件。

输入文件注意事项

目前 AI2BMD 仅支持以下类型的蛋白质模拟：

• 单链蛋白
• 标准氨基酸
• 末端需封闭：N 端为 ACE (乙酰基)，C 端为 NME (N-甲基)
• 暂不支持二硫键（后续版本将更新支持）

若需处理自己的 PDB 文件，请确保补充氢原子并按上述要求添加末端封闭基团（可参考 PyMOL 或 AmberTools 进行处理）。