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下的静电势计算结果对比为:

这个结果表示,如果我们使用6阶插值的PME算法,单步的计算就可以得到实空间迭代10000个Box的近似结果。

总结概要

基于前面几篇博客关于PME算法的理论推导,本文给出了一个简单版本的Python代码实现,并且对比了PME算法相比于实空间迭代算法的优越性。从结果上来看,一维的静电势能计算中,PME单步得到的计算结果非常接近于实空间迭代1万个Box的近似结果。

版权声明

本文首发链接为:https://www.cnblogs.com/dechinphy/p/pme-python.html

作者ID:DechinPhy

更多原著文章:https://www.cnblogs.com/dechinphy/

请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html

posted @ 2024-10-31 10:52  DECHIN  阅读(146)  评论(0编辑  收藏  举报