分子动力学模拟之基于自动微分的LINCS约束

技术背景

在分子动力学模拟的过程中，考虑到运动过程实际上是遵守牛顿第二定律的。而牛顿第二定律告诉我们，粒子的动力学过程仅跟受到的力场有关系，但是在模拟的过程中，有一些参量我们是不希望他们被更新或者改变的，比如稳定的OH键的键长就是一个不需要高频更新的参量。这时就需要在一次不加约束的更新迭代之后（如Velocity-Verlet算法等），再施加一次约束算法，重新调整更新的坐标，使得规定的键长不会产生较大幅度的变更。

初始化坐标参数

为了实现LINCS这一算法，我们先初始化一组随机的坐标用于测试，比如我们测试一个10原子的体系：

# constrain.py
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
N = 10
crd = np.random.random((N, 3))

plt.figure()
plt.plot(crd[:,0], crd[:,1], 'o', color='black')
plt.savefig('initial.png')

初始化的体系效果如下，这是一个仅观测x-y平面的投影的结果（因为二维的投影在可视化上方便一些）：

坐标的更新

参考牛顿定律，我们也用随机的方法产生一组初始速度，用于定义原子体系下一步的运动，再定义一个时间步长，我们就可以获取到下一步的体系坐标：

# constrain.py
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
N = 10
crd = np.random.random((N, 3))
dt = 0.1
vel = np.random.random((N, 3))
new_crd = crd + vel * dt

plt.figure()
plt.plot(crd[:,0], crd[:,1], 'o', color='black')
plt.plot(new_crd[:,0], new_crd[:,1], 'o', color='red')
plt.savefig('move.png')

把旧的坐标和更新之后的坐标放到一起的可视化效果如下：

定义成键关系

因为LINCS约束是施加在键长这一相对参数上的，因此我们首先需要在测试的体系中定义一套成键的关系：

# constrain.py
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
N = 10
crd = np.random.random((N, 3))
dt = 0.1
vel = np.random.random((N, 3))
new_crd = crd + vel * dt

# Add bonds information
bonds = np.array([[0,1],[0,2],[0,4],[2,3],
                  [2,4],[3,8],[5,8],[4,6],
                  [6,7],[7,9]])

plt.figure()
plt.plot(crd[:,0], crd[:,1], 'o', color='black')
plt.plot(new_crd[:,0], new_crd[:,1], 'o', color='red')
for bond in bonds:
    plt.plot(crd[bond][:,0], crd[bond][:,1], color='green')
    plt.plot(new_crd[bond][:, 0], new_crd[bond][:, 1], color='purple')
plt.savefig('move.png')

然后我们把成键关系也在可视化的结果中展现出来，得到这样一张图：

LINCS算法

接下来我们就讲到本文最核心的LINCS算法，其大致流程可以分为如下图（图片来自于参考链接1与LINCS原始文章）所示的3个步骤：

大致描述就是：先按照无约束的条件进行更新，这一点事实上我们在上一个章节中通过速度来更新坐标已经实现了这一操作。然后将更新后的成键在旧的成键上进行投影。最后对新的成键执行一个变换，即可得到保持原有键长的新的体系坐标。我们先看下相关的代码实现和结果，感兴趣的童鞋可以再往后阅读代码实现的思路和原理。

# constrain.py
import numpy as np
from jax import numpy as jnp
from jax import grad, jit, vmap
import matplotlib.pyplot as plt

# Initialization
np.random.seed(0)
N = 10
Dimension = 3
crd = np.random.random((N, Dimension))
# Mass diag
M = np.random.random(N)
Mi = np.identity(N) * M
Mii = np.identity(N) * (M ** (-1))
dt = 0.1
vel = np.random.random((N, Dimension))
new_crd = crd + vel * dt

# Add bonds information
bonds = np.array([[0,1],[0,2],[0,4],[2,3],
                  [2,4],[3,8],[5,8],[4,6],
                  [6,7],[7,9]])
# Bond length
di = np.linalg.norm(crd[bonds[:,0]] - crd[bonds[:,1]], axis=1)

# Automatic differentiation
def B(new_crd, bond, crd):
    return jnp.linalg.norm(new_crd[bond[0]]-new_crd[bond[1]]) -\
           jnp.linalg.norm(crd[bond[0]]-crd[bond[1]])
B_grad = grad(B, argnums=(0,))
B_vmap = jit(vmap(B_grad,(None,0,None)))
B_value = B_vmap(new_crd, bonds, crd)[0]

# LINCS
ccrd = new_crd.copy()
tmp0 = jnp.einsum('ij,kjl->kil', Mii, B_value)
tmp1 = jnp.einsum('jil,kil->jk', B_value, tmp0)
tmp2 = np.linalg.inv(tmp1)
tmp3 = jnp.einsum('ijk,jk->i', B_value, new_crd)-di
tmp4 = jnp.einsum('ij,j->i', tmp2, tmp3)
tmp5 = jnp.einsum('ijk,i->jk', B_value, tmp4)
tmp6 = jnp.einsum('ij,jk->ik', Mii, tmp5)
ccrd -= tmp6

# Draw
plt.subplot(211)
plt.plot(crd[:,0], crd[:,1], 'o', color='black')
plt.plot(new_crd[:,0], new_crd[:,1], 'o', color='blue')
plt.plot(ccrd[:,0], ccrd[:,1], 'o', color='red')
for bond in bonds:
    plt.plot(crd[bond][:,0], crd[bond][:,1], color='black')
    plt.plot(new_crd[bond][:,0], new_crd[bond][:,1], color='blue')
    plt.plot(ccrd[bond][:, 0], ccrd[bond][:, 1], color='red')

plt.subplot(212)
di = np.linalg.norm(crd[bonds[:,0]] - crd[bonds[:,1]], axis=1)
diuc = np.linalg.norm(new_crd[bonds[:,0]] - new_crd[bonds[:,1]], axis=1)
dic = np.linalg.norm(ccrd[bonds[:,0]] - ccrd[bonds[:,1]], axis=1)
plt.plot(di, color='black')
plt.plot(diuc, color='blue')
plt.plot(dic, '+', color='red')
plt.savefig('move.png')

执行输出的结果如下图所示：

在这个结果中我们可以看到第二个图中红色的十字就是施加LINCS约束之后的结果，很显然的距离原始的键长更近。需要额外提醒的是，第一张图中的成键实际上是三维的成键，所以视觉上的大小差异不是真是的键长大小差异，具体差异数值还是以第二张图中展示的为准。

LINCS算法原理以及代码实现思路

首先我们提到了分子的动力学模拟过程还是遵守牛顿第二定律，也就是：

\[\frac{d^2\textbf r}{dt^2}=\textbf M^{-1}\textbf f \]

其中\(\textbf r\)是一个\(N\times 3\)的三维坐标体系，这里\(N\)是体系的原子数，\(\textbf M\)是一个\(N\times N\)的对角矩阵，每一个对角元代表一个原子的质量。事实上在计算过程中更加经常用到的是\(\textbf M\)的逆矩阵，又由于\(\textbf M\)是一个对角矩阵，因此\(\textbf M^{-1}\)实际上就是每个对角元为对应原子质量的倒数这样的一个对角矩阵。\(\textbf f\)是跟\(\textbf r\)维度相同的体系作用力。

LINCS约束的方程可以表述为K个方程：

\[g_i(\textbf r)=|\textbf r_{i1}-\textbf r_{i2}|-d_i=0\ \ \ \ i=1,...,K \]

其中K的大小在这里代表了成键的对数，简单理解就是保证每一对更新后的键的键长的大小与正常的键长大小保持一致，比方说固定了一个OH基中O和H的相对距离。施加该约束的过程可以表述为拉格朗日乘子法：

\[-\textbf M\frac{d^2\textbf r}{dt^2}=\frac{\partial}{\partial\textbf r}(\textbf V-\lambda \cdot g) \]

其中非势能项可以定位为\(B^{T}\lambda\)，其中\(B\)定义为：

\[B_i=\frac{\partial g_i}{\partial r_i} \]

由于这个形式涉及到了微分，不过由于自动微分这项技术的诞生，使得我们不需要自己再去手动的计算这个微分项，只需要把\(g_i\)的形式给定，就可以在Jax中非常方便的计算其导数，并且有别于数值微分，自动微分兼具了高性能与高精度。而另外一点是向量化的操作，在Numba和Jax中分别支持了CPU上和GPU上的向量化操作，我们只需要写一条计算的方法，就可以把这个计算公式扩展到对更高维的数据进行处理，在Jax中这一功能接口为vmap。举个例子说，我们只需要写好计算\(B_i\)的过程，就可以直接用vmap推广到求整个的\(B\)。思路大体上就是如此，具体的过程可以参考上一章节中的源代码。

需要注意的是，这是一个0项，即一阶导数\(\frac{dg}{dt}\)和二阶导数\(\frac{d^2g}{dt^2}\)都是0的项，再结合leap-frog坐标更新算法，可以得到最终的坐标更新表达式（具体的推导过程还是建议看下原始文章，很多平台比如Gromacs也是使用了最终的这个表达式来进行计算或者优化）为：

\[\textbf r_{n+1}=\textbf r_{n+1}^{unc}-\textbf M^{-1}B_n(B_n\textbf M^{-1}B_n^T)^{-1}(B_n\textbf r_{n+1}^{unc}-\textbf d) \]

而更新完坐标之后，对应的速度也需要得到校正，这里以leap-frog算法简单说明一下其速度的更新方法：

\[v_{n+\frac{1}{2}}=\frac{r_{n+1}-r_n}{\Delta t} \]

由于速度的计算方法较为简单，这里我们主要分析下坐标更新的代码实现流程，以及Python的实现过程中有可能遇到的一些坑。

注意事项一

\(\textbf r_{n+1}\)是基于\(\textbf r_{n+1}^{unc}\)来进行调整的，但是如果一开始直接使用：

r=r_unc

来初始化的话，会导致r_unc被覆盖，要知道r_unc还是会被频繁调用的，所以我们初始化的时候最好加上一个copy的操作。

注意事项二

矩阵乘法是从右往左来计算的，而Python中默认的矩阵乘法是从左往右的，因此最好不要直接使用Python中的乘号来直接计算多个矩阵的乘法，替代方案是手写numpy的multiply或者dot等函数配置参数。

注意事项三

在原始的论文中很多地方用到了求转置矩阵的操作，而面对高维矩阵的时候一定要指明操作所对应的轴，在本文的代码实现中，我们是使用了爱因斯坦求和的操作，这个操作在numpy和jax中都有接口支持。

注意事项四

在原始的论文中，为了避免对矩阵进行求逆，使用了一些展开和截断的近似计算的技术。但是对于体系规模不大的场景，其实直接使用numpy或者jax中的求逆函数，速度也不会很慢，本文旨在算法的实现，这里就直接使用了jax的求逆函数。

注意事项五

在jax中的一些函数返回的结果是一个tuple的形式，这是使用vmap和jit技术经常会遇到的情况，虽然并不是很难处理，只需要在得到的结果上取一个0的index即可，但是在实际计算的过程中还是需要注意。

总结

具体的代码实现，都在上一个章节中完整的展示了出来，这一章节只是介绍了LINCS算法的形式以及实现LINCS算法的一些思路，更加详细的推导，还是建议看下原始论文。

总结概要

本文通过完整的案例及其算法实现的过程，介绍了LINCS（Linear Constraint Solver）这一分子动力学模拟过程常用的约束算法。得益于Jax这一框架的便用性及其对numpy的强大支持、对GPU计算的优化、还有自动微分与向量化运算等技术的实现，使得我们实现LINCS这一算法变的不再困难。