PME算法简单Python实现

技术背景

在前面的两篇博客中，我们分别介绍了Ewald算法求解静电势能和基于格点拉格朗日插值法的PME算法。在多种计算优化算法（Ewald求和、快速傅里叶变换、格点拉格朗日插值、截断近似）的加持下，使得我们不需要在实空间进行大量的迭代，也可以得到一个近似收敛的静电势能结果。相关的PME计算公式为：

\[\begin{align*} E&=E^S+E^L-E^{self}\\ &=\sum_{i,j\in \{Neigh\}}\frac{q_iq_j}{4\pi\epsilon_0|\mathbf{r}_j-\mathbf{r}_i|}Erfc\left(\frac{|\mathbf{r}_j-\mathbf{r}_i|}{\sqrt{2}\sigma}\right)\\ &+\frac{V}{2k_xk_yk_z\epsilon_0}\sum_{|\mathbf{k}|>0}\frac{e^{-\frac{\sigma^2 k^2}{2}}}{k^2}|F_{\mathbf{k}}(Q)(m_x,m_y,m_z)|^2\\ &-\frac{1}{4\pi\epsilon_0}\frac{1}{\sqrt{2\pi}\sigma}\sum_{i=0}^{N-1}q_i^2 \end{align*} \]

以下做一个Python版本的简单测试。

测试思路

为了对比PME算法的收敛性与实空间迭代的收敛性，我们先取一个一维的简单周期性点电荷体系，一方面通过对实空间盒子进行扩张，以得到一个收敛的趋势。另一方面通过PME算法直接截断计算静电势。然后对比两者的结果，按预期来说，PME的结果应该在大量的实空间迭代之后被近似到，也就是\(V_{PME}=V_{real}(N)\)。当然，因为不同的插值算法也有可能也会导致计算结果的差异性，所以这里使用了两种插值阶数来进行测试。

代码示例

import numpy as np
from scipy.special import erfc
from tqdm import trange
import matplotlib.pyplot as plt
np.random.seed(4)

num_charges=1000
crd = np.random.random(num_charges) * 10
q = np.random.random(num_charges)
q -= np.sum(q) / q.shape[0]
pbc_box = np.array([10.], np.float64)

def energy(crd, q, pbc_box, levels=0, epsilon=1):
    qij = np.einsum('i,j->ij', q, q)
    if levels == 0:
        dis = np.abs(crd[:, None] - crd[None])
        energy = np.triu(qij / (dis+1e-08) / 4 / np.pi / epsilon, k=1).sum()
    else:
        # left box
        crd1 = crd - levels * pbc_box
        dis = np.abs(crd[:, None] - crd1[None])
        energy = np.triu(qij / (dis+1e-08) / 4 / np.pi / epsilon, k=0).sum()
        # right box
        crd2 = crd + levels * pbc_box
        dis = np.abs(crd[:, None] - crd2[None])
        energy += np.triu(qij / (dis+1e-08) / 4 / np.pi / epsilon, k=0).sum()
    return energy

def pme4(crd, q, pbc_box, epsilon=1):
    sigma = (pbc_box[0] ** 2 / crd.shape[0]) ** (1/6) / np.sqrt(2*np.pi)
    qij = np.einsum('i,j->ij', q, q)
    dis = np.abs(crd[:, None] - crd[None])
    dis = np.where(dis<pbc_box-dis, dis, pbc_box-dis)
    coe = erfc(dis / np.sqrt(2) / sigma)
    real_energy = np.sum(coe * np.triu(qij / (dis+1e-08) / 4 / np.pi / epsilon, k=1))
    self_energy = np.sum(q ** 2) / np.sqrt(2 * np.pi) / sigma / 4 / np.pi / epsilon
    Q = np.zeros(10, dtype=np.float64)
    for idx in range(crd.shape[0]):
        x_floor = np.floor(crd[idx])
        x_shift = np.array([-1.5, -0.5, 0.5, 1.5], np.float32)
        x_idx = (np.floor(x_floor + x_shift) % 10).astype(np.int32)
        x = x_shift
        Q[x_idx[0]] += q[idx]*(-8*x[0]**3+12*x[0]**2+2*x[0]-3)/48
        Q[x_idx[1]] += q[idx]*(8*x[1]**3-4*x[1]**2-18*x[1]+9)/16
        Q[x_idx[2]] += q[idx]*(-8*x[2]**3-4*x[2]**2+18*x[2]+9)/16
        Q[x_idx[3]] += q[idx]*(8*x[3]**3+12*x[3]**2-2*x[3]-3)/48
    Q_ifft = np.fft.ifft(Q)
    sk = np.abs(Q_ifft) ** 2 * Q.shape[0]
    k = np.arange(Q.shape[0]) * 2 * np.pi / Q.shape[0]
    res_energy = np.sum(np.exp(-0.5*sigma**2*k[1:]**2)*sk[1:]/k[1:]**2)/2/epsilon/(2*np.pi/pbc_box[0])
    return real_energy - self_energy + res_energy

def pme6(crd, q, pbc_box, epsilon=1):
    sigma = (pbc_box[0] ** 2 / crd.shape[0]) ** (1/6) / np.sqrt(2*np.pi)
    qij = np.einsum('i,j->ij', q, q)
    dis = np.abs(crd[:, None] - crd[None])
    dis = np.where(dis<pbc_box-dis, dis, pbc_box-dis)
    coe = erfc(dis / np.sqrt(2) / sigma)
    real_energy = np.sum(coe * np.triu(qij / (dis+1e-08) / 4 / np.pi / epsilon, k=1))
    self_energy = np.sum(q ** 2) / np.sqrt(2 * np.pi) / sigma / 4 / np.pi / epsilon
    Q = np.zeros(10, dtype=np.float64)
    for idx in range(crd.shape[0]):
        x_floor = np.floor(crd[idx])
        x_shift = np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5], np.float32)
        x_idx = (np.floor(x_floor + x_shift) % 10).astype(np.int32)
        x = x_shift
        Q[x_idx[0]] += -q[idx]*(x[0]**5-2.5*x[0]**4-1.5*x[0]**3+3.75*x[0]**2+0.3125*x[0]-0.78125)/120
        Q[x_idx[1]] += q[idx]*(x[1]**5-1.5*x[1]**4-6.5*x[1]**3+9.75*x[1]**2+1.5625*x[1]-2.34375)/24
        Q[x_idx[2]] += -q[idx]*(x[2]**5-0.5*x[2]**4-8.5*x[2]**3+4.25*x[2]**2+14.0625*x[2]-7.03125)/12
        Q[x_idx[3]] += q[idx]*(x[2]**5+0.5*x[2]**4-8.5*x[2]**3-4.25*x[2]**2+14.0625*x[2]+7.03125)/12
        Q[x_idx[4]] += -q[idx]*(x[1]**5+1.5*x[1]**4-6.5*x[1]**3-9.75*x[1]**2+1.5625*x[1]+2.34375)/24
        Q[x_idx[5]] += q[idx]*(x[0]**5+2.5*x[0]**4-1.5*x[0]**3-3.75*x[0]**2+0.3125*x[0]+0.78125)/120
    Q_ifft = np.fft.ifft(Q)
    sk = np.abs(Q_ifft) ** 2 * Q.shape[0]
    k = np.arange(Q.shape[0]) * 2 * np.pi / Q.shape[0]
    res_energy = np.sum(np.exp(-0.5*sigma**2*k[1:]**2)*sk[1:]/k[1:]**2)/2/epsilon/(2*np.pi/pbc_box[0])
    return real_energy - self_energy + res_energy

N = 100000
e = np.zeros(N)
for i in trange(N):
    e[i] = energy(crd, q, pbc_box, levels=i)
e = np.cumsum(e)
print (e[0], e[-1])

e2 = pme4(crd, q, pbc_box)
print (e2)

e3 = pme6(crd, q, pbc_box)
print (e3)

x = list(range(N))
plt.figure()
plt.xlabel('Box Layers')
plt.ylabel("Energy")
plt.plot(x, e, label='Normal')
plt.plot(x, np.ones_like(x) * e2, label='PME4')
plt.plot(x, np.ones_like(x) * e3, label='PME6')
plt.legend()
plt.savefig('energy.png')

运行输出为：

-6016.973545020364 -6008.267263384039
-6010.335037181866
-6009.780676419067

这个输出结果表示，不加任何额外的Box时，计算得到的电势能为-6016，在左右各加了10万个Box之后，得到的静电势能结果变为-6008。而PME计算使用4阶拉格朗日插值时一步就可以得到-6010，如果使用6阶的插值，一步就可以得到-6009。总体的不同Box Level下的静电势计算结果对比为：