分子動力學模擬之SETTLE約束算法

• 2022 年 3 月 8 日
• 筆記
• Molecular Dynamics

技術背景

在上一篇文章中，我們討論了在分子動力學裏面使用LINCS約束算法及其在具備自動微分能力的Jax框架下的代碼實現。約束算法，在分子動力學模擬的過程中時常會使用到，用於固定一些既定的成鍵關係。例如LINCS算法一般用於固定分子體系中的鍵長關係，而本文將要提到的SETTLE算法，常用於固定一個構成三角形的體系，最常見的就是水分子體系。對於一個水分子而言，O-H鍵的鍵長在模擬的過程中可以固定，H-H的長度，或者我們更常見的作為一個H-O-H的夾角出現的參量，也需要固定。純粹從計算量來考慮的話，RATTLE約束算法需要迭代計算，LINCS算法需要求矩陣逆（雖然已經給出了截斷優化的算法），而SETTLE只涉及到坐標變換，顯然SETTLE在約束大規模的水盒子時，性能會更加優秀。

算法流程

類似於LINCS算法的，我們先引入一個核心的約束條件：

意思就是說，原子i和原子j之間的距離在分子動力學模擬的迭代過程中保持不變。\(d_{ij}\)可以是初始值，也可以是一個給定值。在滿足這個約束條件的同時，其實隱藏着另外一個約束條件：

換而言之，如果所有原子的運動方向都與其直接的鍵連方向垂直，那麼在模擬迭代的過程中鍵長距離就不會發生改變了。

約束前後

我們還是需要從一個簡單的三角形\(A_0B_0C_0\)模型出發來具體講解這個算法的流程：

假定三角形\(A_0B_0C_0\)為原始位置，三角形\(A_1B_1C_1\)為沒有加約束的坐標更新，而我們的目標是得到三角形\(A_3B_3C_3\)這個加了約束的三角形。這裡先解釋一下為什麼不是三角形\(A_2B_2C_2\)而是三角形\(A_3B_3C_3\)，因為這裡三角形的變換隻是涉及到加了約束條件和沒加約束條件的結果，但是實際上加約束條件是一個步驟比較複雜的過程，這裡寫成三角形\(A_3B_3C_3\)是為了方便跟後續的SETTLE約束所得到的結果下標對齊。

約束算法中的約束條件

在上一步的算法過程中，其實我們已經初步的把施加SETTLE約束算法的過程劃分成了兩個部分：先不加約束的演化，再考慮施加約束的演化。之所以能夠這麼劃分，是因為我們把約束條件考慮成一個約束作用力，這樣有兩個好處：一是約束力也變成連續可微的參量，二是矢量的作用是可以疊加和拆分的，這裡兩步的拆分，就是用到了其矢量作用力的效果。

將三角形初始所在的平面定義為\(\pi_0\)，在經過未加約束的偏移之後得到了三角形\(A_1B_1C_1\)，此時可能已經偏離了原始的平面\(\pi_0\)，這裡將三個頂點所在的位置都構造一個與\(\pi_0\)相互平行的平面，稱為\(\pi_A\)、\(\pi_B\)和\(\pi_C\)。這裡可能會有兩個需要着重注意的點，第一，我們之所以要定義\(\pi_A\)、\(\pi_B\)和\(\pi_C\)這三個平面，是因為約束力總是在徑向的，也就是說，不可能產生與平面\(\pi_0\)相垂直的約束力，因此約束力的作用只是讓三個頂點在與\(\pi_0\)平面相平行的平面上運動，也就是這裡的\(\pi_A\)、\(\pi_B\)和\(\pi_C\)三個平面。換而言之，三角形\(A_3B_3C_3\)的三個頂點必然在\(\pi_A\)、\(\pi_B\)和\(\pi_C\)三個平面內。第二，因為約束力是分子內部的作用力，也就是對重心的合力為0，不會導致整體分子的漂移，因此，在施加約束力前後，重心的位置不變。如果分別用\(D_1\)和\(D_3\)來標記三角形\(A_1B_1C_1\)和\(A_3B_3C_3\)的重心的話，那麼有\(D_1=D_3\)。並且，假如原始位置三角形\(A_0B_0C_0\)的重心\(d_0\)跟\(D_1\)和\(D_3\)如果是在同一個位置的話，那麼原始位置的三角形\(A_0B_0C_0\)就可以通過繞重心的旋轉跟三角形\(A_3B_3C_3\)重合。

建立新的坐標系

基於\(d_0=D_1=D_3\)的假定，我們可以基於三角形\(A_0B_0C_0\)建立這樣一個新的坐標系\(X’Y’Z’\)，原點為\(d_0\)，\(X’-Y’\)平面與\(\pi_0\)平行，而\(Y’-Z’\)平面過\(A_0\)點：

如果從\(Z’\)軸向\(Z’\)軸負方向看（常規理解就是從上往下看的俯視圖），就是這樣的一個平面：

三維旋轉

在前面的章節中我們提到，通過旋轉操作，可以將初始的坐標旋轉到更施加約束之後的坐標完全重合的位置，因此我們先假設三個旋轉角，對原始坐標進行旋轉操作。最後再根據約束條件來計算對應的旋轉角，進而得到施加約束之後的新的坐標，也就是我們最終所需要的結果。在新的坐標系下，我們把三角形\(A_0B_0C_0\)重新標記為三角形\(a_0b_0c_0\)，接下來開始對三角形\(a_0b_0c_0\)的一系列三維旋轉操作：

在初始條件下，因為三角形跟\(X’-Y’\)處在同一個平面上，因此從\(Y’\)軸向\(Y’\)軸正方向看時，會看到一條直線，此時我們繞\(Y’\)軸旋轉一個\(\psi\)的角度，得到了三角形\(a_1b_1c_1\)。

按照同樣的操作，繞\(X’\)軸旋轉\(\phi\)以及繞\(Z’\)軸旋轉\(\theta\)的角度，經過三次的旋轉之後，得到一個新的三角形\(a_3b_3c_3\)。而此時的三角形\(a_3b_3c_3\)正是處於跟三角形\(A_3B_3C_3\)完全重合的狀態。因此，我們就可以根據約束條件計算出來三個歐拉角\(\psi\)、\(\phi\)和\(\theta\)，進而得到我們所需要的約束後的結果：三角形\(A_3B_3C_3\)。

算法解析表達式

關於具體解析表達式的推導，可以參考本文末尾的參考文章1中的附錄A，這裡我們僅展示一些已經推導出來的解析表達式的結果。首先假定未施加約束算法的三角形\(A_1B_1C_1\)在\(X’Y’Z’\)坐標系下的坐標分別為：\([X’_{A_1},Y’_{A_1},Z’_{A_1}]\)、\([X’_{B_1},Y’_{B_1},Z’_{B_1}]\)和\([X’_{C_1},Y’_{C_1},Z’_{C_1}]\)，以及三角形\(a_0b_0c_0\)在\(X’Y’Z’\)坐標系下的坐標分別為：

關於這個坐標數值，再回頭看下這個圖可能會更加清晰明了一些：

那麼我們最終可以得到的旋轉角為：

其中\(\alpha\)、\(\beta\)和\(\gamma\)的取值如下：

而最終得到的三角形\(a_3b_3c_3\)的坐標為：

在上述的公式中，我們根據未施加約束的三角形\(A_1B_1C_1\)在\(X’Y’Z’\)坐標軸下的新坐標，以及初始的三角形\(a_0b_0c_0\)在\(X’Y’Z’\)坐標軸下的新坐標，就可以計算出\(\phi，\psi，\theta\)這三個空間角。進而可以得到施加約束之後所得到的等效三角形\(a_3b_3c_3\)在\(X’Y’Z’\)坐標軸下的坐標，再經過兩個坐標之間的轉化之後，就可以得到我們所需要的施加約束條件之後的目標三角形\(A_3B_3C_3\)在原始的\(XYZ\)坐標系下的笛卡爾坐標。

三維空間坐標變換

在上一個章節中我們提到了還需要一個三維空間的坐標轉化，因為前後採取了兩個不同的坐標系。如果分別標記\(R_X,R_Y,R_Z\)為繞着\(X,Y,Z\)軸旋轉的矩陣，則相應的旋轉矩陣為：

我們也可以把這些變換看做是一個整體：\(T(\psi,\phi,\theta)=R_Z(\theta)R_X(\phi)R_Y(\psi)\)，這個變換不僅可以用於計算坐標系的變換下所有對應節點的坐標變換，還可以用於計算上一個步驟中提到的三角形\(a_3b_3c_3\)。但是上一個步驟中的三角形\(a_3b_3c_3\)的坐標已經給出，這裡我們為了得到坐標系的逆變換，因此不得不把坐標變換的完整形式列出來，則對應的\(T^{-1}(\psi,\phi,\theta)=R_Z(\theta)R_X(\phi)R_Y(\psi)\)就是其逆變換。了解清楚變換的形式之後，我們再回過頭來看這個\(XYZ\)坐標繫到\(X’Y’Z’\)坐標系的變換：

我們發現這裡其實不僅僅是包含有坐標軸的旋轉，還包含了坐標系原點的偏移，不過這個漂移倒是比較好處理，可以在後續的計算過程中點出即可。

坐標變換代碼實現

我們求解從原始的坐標\(XYZ\)到新坐標\(X’Y’Z’\)的代碼實現思路是這樣的：

1. 通過python代碼構建一個等腰三角形，通過旋轉使得其達到一個初始位置\(A_0B_0C_0\)，這個初始位置對應上述章節中的三角形\(A_0B_0C_0\)，之所以要通過三個角度的旋轉來實現，是為了同時演示這個三維旋轉的過程，對應的是如下代碼中的rotation函數；
2. 計算三角形\(A_0B_0C_0\)的質心center of mass，表示為\(M\)；
3. 因為三個點就可以固定一個平面，這裡我們選取的是\(A\)、\(B\)這兩個點以及\(\vec{BC}\otimes \vec{CA}\)這個向量（不能只取三角形所在平面上的點，否則無法正確求解方程，因為對應的參數矩陣秩為2），再假定一個旋轉矩陣\(R\)，聯立一個九元一次方程組，就可以解得旋轉矩陣的具體值。

通過這三個點聯立的方程組可以表示為：

相關的求解代碼如下所示：

# settle.py
from jax import numpy as np
from jax import vmap, jit

def rotation(psi,phi,theta,v):
    """ Module of rotation in 3 Euler angles. """
    RY = np.array([[np.cos(psi),0,-np.sin(psi)],
                   [0, 1, 0],
                   [np.sin(psi),0,np.cos(psi)]])
    RX = np.array([[1,0,0],
                   [0,np.cos(phi),-np.sin(phi)],
                   [0,np.sin(phi),np.cos(phi)]])
    RZ = np.array([[np.cos(theta),-np.sin(theta),0],
                   [np.sin(theta),np.cos(theta),0],
                   [0,0,1]])
    return np.dot(RZ,np.dot(RX,np.dot(RY,v)))

multi_rotation = jit(vmap(rotation,(None,None,None,0)))

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    # construct params
    ra = 0.5
    rb = 0.7
    rc = 1.2
    psi = 0.4
    phi = 0.5
    theta = 1.3
    # construct initial crd
    crd = np.array([[0, ra, 0],
                    [-rc, -rb, 0],
                    [rc, -rb, 0]])
    shift = np.array([0.1, 0.1, 0.1])
    crd = multi_rotation(psi,phi,theta,crd) + shift
    # get the center of mass
    com = np.average(crd,0)
    # 3 points are selected to solve the initial rotation matrix
    xyz = [0,0,0]
    xyz[0] = crd[0]-com
    xyz[1] = crd[1]-com
    cross = np.cross(crd[2]-crd[1],crd[0]-crd[2])
    cross /= np.linalg.norm(cross)
    xyz[2] = cross
    xyz = np.array(xyz)
    inv_xyz = np.linalg.inv(xyz)
    v0 = np.array([0,-rc,0])
    v1 = np.array([ra,-rb,0])
    v2 = np.array([0,0,1])
    # final rotation matrix is constructed by following
    Rot = np.array([np.dot(inv_xyz,v0),np.dot(inv_xyz,v1),np.dot(inv_xyz,v2)])
    print (Rot)
    # some test cases and results
    origin = crd[0]
    print(np.dot(Rot, origin-com))
    # [1.4901161e-08 5.0000000e-01 0.0000000e+00]
    origin = crd[1]
    print(np.dot(Rot, origin-com))
    # [-1.2000000e+00 -7.0000005e-01 -5.9604645e-08]
    origin = crd[2]
    print(np.dot(Rot, origin-com))
    # [ 1.2000000e+00  2.0000000e-01 -1.4901161e-08]
    origin = xyz[2]
    print(np.dot(Rot, origin))
    # [0.0000000e+00 2.9802322e-08 1.0000000e+00]

上述代碼中所得到的Rot這個矩陣，就是我們所需的將\(XYZ\)坐標系轉移到\(X’Y’Z’\)坐標系的旋轉矩陣，當然，需要注意的是在做坐標系映射的過程中，除了考慮坐標系的旋轉，還需要考慮坐標系的平移，在這裡就是重心的偏移量，這個偏移量在原始的文章中是沒有使用到的，但是我們實際的模擬過程中是肯定會遇到的。只是說因為約束力的合力是0，因此在從三角形\(A_1B_1C_1\)到三角形\(A_3B_3C_3\)的過程是不需要考慮重心偏移的，但是我們在第一步從三角形\(A_0B_0C_0\)到三角形\(A_1B_1C_1\)或者是三角形\(a_0b_0c_0\)的過程是必然會面臨的。同時，在上述代碼的結尾處我們給出了四個驗證的案例，這與我們所預期的結果是相符合的，坐標值從\(XYZ\)坐標系變換成了以\(\pi_0\)為\(X’-Y’\)平面且\(Y’-Z’\)平面過\(a_0\)點的坐標繫上。

需要特別提及的是，上述代碼中所使用到的JAX框架支持了vmap這種便捷矢量化計算的操作，因此在rotation函數中只實現了一個旋轉矩陣對一個向量的操作，再通過vmap將其擴展到了對多個矢量，也就是多個點空間旋轉操作上，變成了multi_rotation函數，這樣的操作也更加符合我們對多個原子坐標的定義形式。

SETTLE代碼實現

在前面的章節中我們已經基本完成了SETTLE約束算法的介紹，這裡我們先總結梳理一遍實現SETTLE的步驟，再看下代碼實現以及相關的效果：

1. 獲取\(XYZ\)坐標系下三角形\(A_0B_0C_0\)的坐標以及三角形\(A_1B_1C_1\)的坐標；
2. 根據三角形\(A_0B_0C_0\)的坐標計算得\(r_a,r_b,r_c\)的值以及質心\(M\)的坐標；
3. 根據三角形\(A_0B_0C_0\)的坐標以及\(r_a,r_b,r_c\)的值和質心\(M\)的坐標，計算\(XYZ\)坐標繫到\(X’Y’Z’\)坐標系的變換矩陣\(Rot\)及其逆變換\(Rot^{-1}\)；
4. 用\(Rot\)和質心坐標計算三角形\(A_1B_1C_1\)在坐標系\(X’Y’Z’\)下的坐標；
5. 根據三角形\(A_1B_1C_1\)在坐標系\(X’Y’Z’\)下的坐標計算得三個旋轉角\(\phi,\psi,\theta\)；
6. 根據\(\phi,\psi,\theta\)和\(r_a,r_b,r_c\)計算三角形\(A_3B_3C_3\)，也就是我們的目標三角形，在\(X’Y’Z’\)坐標系下的坐標；
7. 使用\(Rot^{-1}\)將三角形\(A_3B_3C_3\)在坐標系\(X’Y’Z’\)下的坐標變換回坐標系\(XYZ\)下的坐標，至此就完成了SETTLE算法的實現。

相關代碼實現如下所示：

# settle.py
from jax import numpy as np
from jax import vmap, jit

def rotation(psi,phi,theta,v):
    """ Module of rotation in 3 Euler angles. """
    RY = np.array([[np.cos(psi),0,-np.sin(psi)],
                   [0, 1, 0],
                   [np.sin(psi),0,np.cos(psi)]])
    RX = np.array([[1,0,0],
                   [0,np.cos(phi),-np.sin(phi)],
                   [0,np.sin(phi),np.cos(phi)]])
    RZ = np.array([[np.cos(theta),-np.sin(theta),0],
                   [np.sin(theta),np.cos(theta),0],
                   [0,0,1]])
    return np.dot(RZ,np.dot(RX,np.dot(RY,v)))

multi_rotation = jit(vmap(rotation,(None,None,None,0)))

def get_rot(crd):
    """ Get the coordinates transform matrix. """
    # get the center of mass
    com = np.average(crd, 0)
    rc = np.linalg.norm(crd[2]-crd[1])/2
    ra = np.linalg.norm(crd[0]-com)
    rb = np.sqrt(np.linalg.norm(crd[2]-crd[0])**2-rc**2)-ra
    # 3 points are selected to solve the initial rotation matrix
    xyz = [0, 0, 0]
    xyz[0] = crd[0] - com
    xyz[1] = crd[1] - com
    cross = np.cross(crd[2] - crd[1], crd[0] - crd[2])
    cross /= np.linalg.norm(cross)
    xyz[2] = cross
    xyz = np.array(xyz)
    inv_xyz = np.linalg.inv(xyz)
    v0 = np.array([0, -rc, 0])
    v1 = np.array([ra, -rb, 0])
    v2 = np.array([0, 0, 1])
    # final rotation matrix is constructed by following
    Rot = np.array([np.dot(inv_xyz, v0), np.dot(inv_xyz, v1), np.dot(inv_xyz, v2)])
    inv_Rot = np.linalg.inv(Rot)
    return Rot, inv_Rot

def xyzto(Rot, crd, com):
    """ Apply the coordinates transform matrix. """
    return np.dot(Rot, crd-com)

multi_xyzto = jit(vmap(xyzto,(None,0,None)))

def toxyz(Rot, crd, com):
    """ Apply the inverse of transform matrix. """
    return np.dot(Rot, crd-com)

multi_toxyz = jit(vmap(toxyz,(None,0,None)))

def get_circumference(crd):
    """ Get the circumference of all triangles. """
    return np.linalg.norm(crd[0]-crd[1])+np.linalg.norm(crd[0]-crd[2])+np.linalg.norm(crd[1]-crd[2])

jit_get_circumference = jit(get_circumference)

def get_angles(crd_0, crd_t0, crd_t1):
    """ Get the rotation angle psi, phi and theta. """
    com = np.average(crd_0, 0)
    rc = np.linalg.norm(crd_0[2] - crd_0[1]) / 2
    ra = np.linalg.norm(crd_0[0] - com)
    rb = np.sqrt(np.linalg.norm(crd_0[2] - crd_0[0]) ** 2 - rc ** 2) - ra
    phi = np.arcsin(crd_t1[0][2]/ra)
    psi = np.arcsin((crd_t1[1][2]-crd_t1[2][2])/2/rc/np.cos(phi))
    alpha = -rc*np.cos(psi)*(crd_t0[1][0]-crd_t0[2][0])+(-rb*np.cos(phi)-rc*np.sin(psi)*np.sin(phi))*(crd_t0[1][1]-crd_t0[0][1])+ \
            (-rb*np.cos(phi)+rc*np.sin(psi)*np.sin(phi))*(crd_t0[2][1]-crd_t0[0][1])
    beta = -rc*np.cos(psi)*(crd_t0[2][1]-crd_t0[1][1])+(-rb*np.cos(phi)-rc*np.sin(psi)*np.sin(phi))*(crd_t0[1][0]-crd_t0[0][0])+ \
           (-rb*np.cos(phi)+rc*np.sin(psi)*np.sin(phi))*(crd_t0[2][0]-crd_t0[0][0])
    gamma = crd_t1[1][1]*(crd_t0[1][0]-crd_t0[0][0])-crd_t1[1][0]*(crd_t0[1][1]-crd_t0[0][1])+\
        crd_t1[2][1]*(crd_t0[2][0]-crd_t0[0][0])-crd_t1[2][0]*(crd_t0[2][1]-crd_t0[0][1])
    sin_part = gamma/np.sqrt(alpha**2+beta**2)
    theta = np.arcsin(sin_part)-np.arctan(beta/alpha)
    return phi, psi, theta

jit_get_angles = jit(get_angles)

def get_d3(crd_0, psi, phi, theta):
    """ Calculate the new coordinates by 3 given angles. """
    com = np.average(crd_0, 0)
    rc = np.linalg.norm(crd_0[2] - crd_0[1]) / 2
    ra = np.linalg.norm(crd_0[0] - com)
    rb = np.sqrt(np.linalg.norm(crd_0[2] - crd_0[0]) ** 2 - rc ** 2) - ra
    return np.array([[-ra*np.cos(phi)*np.sin(theta), ra*np.cos(phi)*np.cos(theta), ra*np.sin(phi)],
                     [-rc*np.cos(psi)*np.cos(theta)+rb*np.sin(theta)*np.cos(phi)+rc*np.sin(theta)*np.sin(psi)*np.sin(phi),
                      -rc*np.cos(psi)*np.sin(theta)-rb*np.cos(theta)*np.cos(phi)-rc*np.cos(theta)*np.sin(psi)*np.sin(phi),
                      -rb*np.sin(phi)+rc*np.sin(psi)*np.cos(phi)],
                     [rc*np.cos(psi)*np.cos(theta)+rb*np.sin(theta)*np.cos(phi)-rc*np.sin(theta)*np.sin(psi)*np.sin(phi),
                      rc*np.cos(psi)*np.sin(theta)-rb*np.cos(theta)*np.cos(phi)+rc*np.cos(theta)*np.sin(psi)*np.sin(phi),
                      -rb*np.sin(phi)-rc*np.sin(psi)*np.cos(phi)]])

jit_get_d3 = jit(get_d3)

if __name__ == '__main__':
    from mpl_toolkits.mplot3d import Axes3D
    import matplotlib.pyplot as plt
    import numpy as onp
    onp.random.seed(0)
    # construct params
    ra = 1.0
    rb = 0.5
    rc = 1.2
    psi = 0.4
    phi = 0.5
    theta = 1.3
    # construct initial crd
    crd = np.array([[0, ra, 0],
                    [-rc, -rb, 0],
                    [rc, -rb, 0]])
    shift = np.array([0.1, 0.1, 0.1])
    # get the initial crd
    crd_0 = multi_rotation(psi,phi,theta,crd) + shift
    vel = np.array(onp.random.random(crd_0.shape)-0.5)
    dt = 1
    # get the unconstraint crd
    crd_1 = crd_0 + vel * dt
    com_0 = np.average(crd_0, 0)
    com_1 = np.average(crd_1, 0)
    # get the coordinate transform matrix and correspond inverse operation
    rot, inv_rot = get_rot(crd_0)
    crd_t0 = multi_xyzto(rot, crd_0, com_0)
    com_t0 = np.average(crd_t0, 0)
    crd_t1 = multi_xyzto(rot, crd_1, com_1)+com_1
    com_t1 = np.average(crd_t1, 0)
    print ('crd_t1:\n', crd_t1)
    # crd_t1:
    # [[0.11285806  1.1888411   0.22201033]
    #  [-1.3182535 - 0.35559598  0.3994387]
    # [1.5366794 - 0.00262779
    # 0.3908713]]
    phi, psi, theta = jit_get_angles(crd_0, crd_t0, crd_t1-com_t1)
    crd_t3 = jit_get_d3(crd_t0,psi,phi,theta)+com_t1
    com_t3 = np.average(crd_t3, 0)
    crd_3 = multi_toxyz(inv_rot, crd_t3, com_t3) + com_1
    print ('crd_t3:\n', crd_t3)
    # crd_t3:
    # [[0.01470824  1.2655654   0.22201033]
    #  [-1.0361676 - 0.3326143   0.39943868]
    # [1.3527434 - 0.10233352
    # 0.39087126]]
    print(jit_get_circumference(crd_0))
    # 6.2418747
    print(jit_get_circumference(crd_t0))
    # 6.2418737
    print(jit_get_circumference(crd_1))
    # 6.853938
    print(jit_get_circumference(crd_t1))
    # 6.8539376
    print(jit_get_circumference(crd_t3))
    # 6.2418737
    print(jit_get_circumference(crd_3))
    # 6.241874

    # Plotting
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    x_0 = np.append(crd_0[:,0],crd_0[0][0])
    y_0 = np.append(crd_0[:,1],crd_0[0][1])
    z_0 = np.append(crd_0[:,2],crd_0[0][2])
    ax.plot(x_0, y_0, z_0, color='black')
    x_1 = np.append(crd_1[:, 0], crd_1[0][0])
    y_1 = np.append(crd_1[:, 1], crd_1[0][1])
    z_1 = np.append(crd_1[:, 2], crd_1[0][2])
    ax.plot(x_1, y_1, z_1, color='blue')
    x_3 = np.append(crd_3[:, 0], crd_3[0][0])
    y_3 = np.append(crd_3[:, 1], crd_3[0][1])
    z_3 = np.append(crd_3[:, 2], crd_3[0][2])
    ax.plot(x_3, y_3, z_3, color='red')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.show()

這個代碼實現的流程，可以通過理解main函數中的順序來進行解讀：首先用隨機數構建前後兩個三角形\(A_0B_0C_0\)和\(A_1B_1C_1\)用於測試SETTLE算法。然後使用get_rot函數來得到從坐標\(XYZ\)到\(X’Y’Z’\)的映射旋轉關係。但是這裡需要注意的是，這個函數得到的旋轉矩陣只有旋轉給定矢量的功能，其中重心偏移是需要自己手動加上的。有了這一層的映射關係關係之後，就可以計算得到\(X’Y’Z’\)坐標系下所對應的兩個三角形，在代碼中就是crd_t0和crd_t1。根據新坐標系下的兩個三角形的坐標，可以計算出來\(\psi,\phi,\theta\)這三個角度，進而計算出來我們所需要的\(X’Y’Z’\)坐標系下的三角形\(a_3b_3c_3\)，也就是代碼中的crd_t3。此時我們通過計算所得到的三角形的周長，我們可以發現中間未加約束的周長變化較大，但是再施加約束之後，周長與原始三角形的周長一致。在最後，我畫了幾個三維圖以供示意：

其中黑色的是原始的三角形，藍色的是未施加約束條件的偏移，其中重心也發生了較為明顯的變化，而紅色的三角形對應的是施加約束後的三角形。還可以從另外一個角度來查看施加約束前後的兩個三角形的平面關係：

需要特別注意的是，獲取坐標變換的矩陣只能針對矢量進行旋轉，也就是這裡針對重心的旋轉。而從未施加約束的三角形\(A_1B_1C_1\)到施加約束的三角形\(A_3B_3C_3\)重心是不會發生改變的，因此在施加\(Rot^{-1}\)的時候，最後需要加上三角形\(A_1B_1C_1\)在\(XYZ\)坐標軸下的重心，才是三角形\(a_3b_3c_3\)在\(XYZ\)坐標軸下的真正位置。

總結概要

繼上一篇文章介紹了分子動力學模擬中常用的LINCS約束算法之後，本文再介紹一種SETTLE約束算法，及其基於Jax的實現方案。LINCS約束算法相對來說比較通用，更適合於成鍵關係比較複雜的通用的體系，而SETTLE算法更加適用於三原子軸對稱體系，比如水分子。SETTLE算法結合velocity-verlet算法，可以確保一個分子只進行整體的旋轉運動，互相之間的距離又保持不變。比較關鍵的是，SETTLE算法所依賴的參數較少，也不需要微分，因此在性能上十分有優勢。

版權聲明

本文首發鏈接為：//www.cnblogs.com/dechinphy/p/settle.html

作者ID：DechinPhy

更多原著文章請參考：//www.cnblogs.com/dechinphy/

打賞專用鏈接：//www.cnblogs.com/dechinphy/gallery/image/379634.html

騰訊雲專欄同步：//cloud.tencent.com/developer/column/91958

參考文章

1. SETTLE: An Analytical Version of the SHAKE and RATTLE Algorithm for Rigid Water Models. Shuichi Miyamoto and Peter A.Kollman.

*註：本文所有的算法流程示意圖，均來自於參考文章1