作者 | 杨慧丹 审稿 | 周珍冉
今天给大家介绍2020年12月22日在arXiv上发表的预印本“TorchMD: A deep learning framework for molecular simulations”，该文章的通讯作者是西班牙庞培法布拉大学的Gianni De Fabritiis教授。文章中提出了”TorchMD“，一个混合经典和机器学习势的分子模拟的框架。文章中所有的力计算，包括键、角、二面体、Lennard-Jones和库仑相互作用，都用PyTorch数组和运算表示。此外，TorchMD 还支持学习和模拟神经网络势。使用了标准的Amber全原子模拟进行验证，学习从头算势，执行端到端训练，最后学习和模拟了一个蛋白质折叠的粗粒度模型。
1
简介
经典分子动力学（MD）是一种计算密集型技术，能够对分子过程进行定量研究。MD已经有了显著的发展，并成功地应用于构象变化、折叠、结合、渗透和许多其他问题，然而它也有明显的局限性：需要广泛和仔细优化的原子间势表（称为力场），以及难以达到几个重要生理过程的时间尺度。随着深度神经网络（DNN）架构的出现，机器学习（ML）势变得特别有吸引力。DNNs提供了一个非常有前景的方法，它在从代价更高的方法获得的大规模数据库上进行训练之后，在MD模拟中去嵌入快速而准确的势能函数。
作者介绍的TorchMD利用ML库PyTorch从头构建的分子动力学代码。TorchMD通过将MD中的键合和非键合力术语扩展到任意复杂的DNN上，实现了机器学习势的快速成型和集成。TorchMD有两个关键点:其一，使用PyTorch编写，容易集成其他ML-PyTorch模型；其二，提供了执行端到端可微模拟的能力，在它所有的参数上都是可微的。本文介绍了TorchMD的功能，重点介绍了支持的功能形式和数据驱动DNN势的有效拟合策略。
2
方法
2.1TorchMD模拟
TorchMD提供NVT集成模拟，包括郎之万控温器。由麦克斯韦-玻耳兹曼分布导出起始原子速度，用velocity-Verlet算法做积分，长程静电近似使用反应场方法。TorchMD也支持周期系统的模拟，它可以使用L-BFGS算法进行最小化，并且可以在PyTorch支持的任何设备（CPU、GPU、TPU）上运行进行模拟。不同于专门的MD代码的是，它不是为速度而设计的。TorchMD使用与经典MD代码（如ACEMD）一致的化学单位。
2.2分析势函数
TorchMD支持通过parmed读取AMBER力场参数，它还实现了自己易于读取的基于YAML的力场格式以更快地进行原型设计和开发。
图1.模拟水箱的力场YAML文件示例
TorchMD实现了AMBER势的函数形式。它提供了所有基本的AMBER术语：谐波键、角、扭转和非键范德华（VdW）和静电能。它们的势函数的实现如下：
除此之外，TorchMD还允许使用PyTorch编写的任何其他外部势Vext 。
因此，总电势计算:
由于PyTorch提供了自动微分，因此无需根据力计算分析梯度，只需一次单独的autograd PyTorch调用系统的总能量就可以获得力。然而出于性能考虑对所有势的分析项都采用了分析梯度。
2.3训练机器学习势
TorchMD提供了一个在PyTorch中训练神经网络势的完全可用的代码：TorchMD-Net (github.com/torchmd/ torchmd-net)。目前作者使用的是基于SchNet的模型，采用了SchNetPack中的特征化和原子层，但完全重写了训练和推理部分。为了允许在多个GPU上进行训练而使用了PyTorch lightning框架对网络进行训练。对于QM9数据集，作者使用能量的均方误差的标准损失函数来训练模型。对于粗粒度模型，作者使用自下而上的“力匹配”方法进行训练，重点是从原子模拟中再现系统的热力学。
3
结果
作者提供了一些应用示例来演示TorchMD的功能。
(1) 一组典型的MD用例（水盒、小肽、蛋白质和配体）主要用于评估速度和能量守恒。
表1.不同系统上50000步1 fs/timestep的性能比较
(2)在QM9这个包含134k个小分子构象和能量的数据集上验证代码。
图2.QM9数据集的学习曲线
(3) 通过从短MD轨道恢复力场参数证明了TorchMD的端到端可微性。
图3.训练过程中的训练损耗和部分原子电荷的演变
(4) 使用在全原子MD模拟数据上训练的NNP，提出了一种微蛋白Chignolin的粗粒度模拟，也描述了如何产生一个基于神经网络的chignolin粗粒度模型。
chignolin的两个粗粒度模型：一个仅基于α-碳原子（CA），另一个基于α-碳原子和β-碳原子（CACB）。
图4.参照组全原子MD模拟(左)和两个粗粒度模型CA(中)和CACB(右)的二维自由能面。参照组MD模拟图显示了表面上三个能量最小值的位置，对应于折叠状态(红点)、展开构象(蓝点)和错误折叠状态(橙点)。
4
结论
本文中作者演示了TorchMD，并展示了几种可能的应用，从Amber全原子模拟，到参数的端到端学习，最后是蛋白质折叠的粗粒度神经网络势。TorchMD参数的端到端可微性是诸如Open Force Field Initiative等项目可能利用的一个特性。此外，为了提高速度，研究小组计划在OpenMM和ACEMD中促进机器学习势的集成，同时作者相信TorchMD可以发挥重要作用，通过促进ML和MD领域之间的实验，加快model-train-evaluate原型开发的周期，以及促进在分子模拟中采用基于数据的方法。
代码
https://github.com/torchmd
参考资料
arXiv:2012.12106 [physics.chem-ph]
(or arXiv:2012.12106v1 [physics.chem-ph] for this version)
https://arxiv.org/abs/2012.12106
