技術背景
在分子動力學模擬的過程中，考慮到運動過程實際上是遵守牛頓第二定律的。而牛頓第二定律告訴我們，粒子的動力學過程僅跟受到的力場有關系，但是在模擬的過程中，有一些參量我們是不希望他們被更新或者改變的，比如穩定的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算法原理以及代碼實現思路
首先我們提到了分子的動力學模擬過程還是遵守牛頓第二定律，也就是：
其中\(\textbf r\)是一個\(N\times 3\)的三維坐標體系，這里\(N\)是體系的原子數，\(\textbf M\)是一個\(N\times N\)的對角矩陣，每一個對角元代表一個原子的質量。事實上在計算過程中更加經常用到的是\(\textbf M\)的逆矩陣，又由於\(\textbf M\)是一個對角矩陣，因此\(\textbf M^{-1}\)實際上就是每個對角元為對應原子質量的倒數這樣的一個對角矩陣。\(\textbf f\)是跟\(\textbf r\)維度相同的體系作用力。
LINCS約束的方程可以表述為K個方程：
其中K的大小在這里代表了成鍵的對數，簡單理解就是保證每一對更新后的鍵的鍵長的大小與正常的鍵長大小保持一致，比方說固定了一個OH基中O和H的相對距離。施加該約束的過程可以表述為拉格朗日乘子法：
其中非勢能項可以定位為\(B^{T}\lambda\)，其中\(B\)定義為：
由於這個形式涉及到了微分，不過由於自動微分這項技術的誕生，使得我們不需要自己再去手動的計算這個微分項，只需要把\(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也是使用了最終的這個表達式來進行計算或者優化）為：
而更新完坐標之后，對應的速度也需要得到校正，這里以leap-frog算法簡單說明一下其速度的更新方法：
由於速度的計算方法較為簡單，這里我們主要分析下坐標更新的代碼實現流程，以及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這一算法變的不再困難。
版權聲明
本文首發鏈接為：https://www.cnblogs.com/dechinphy/p/lincs.html
作者ID：DechinPhy
更多原著文章請參考：https://www.cnblogs.com/dechinphy/
打賞專用鏈接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
騰訊雲專欄同步：https://cloud.tencent.com/developer/column/91958