ANKH: A Generalized $\mathcal{O}(N)$ Interpolated Ewald Strategy for Molecular Dynamics Simulations

Fast Multipole Methods [20, 21, 18] (FMM) is a family of fast divide and conquer strategies to compute $N$-body problems in linear or linearithmic complexity. Important efforts have been made in the literature to incorporate 3D FMM (i.e. for the Coulomb potential) in molecular dynamics simulation codes [37, 25, 27]. Such a FMM scheme is based on a separable spherical harmonic expansion of the Coulomb potential $\frac{1}{|\mathbf{x}-\mathbf{y}|}$ writing as:

where $\mathbf{x}=(r_{\mathbf{x}},\theta_{\mathbf{x}},\phi_{\mathbf{x}})$, $\mathbf{y}=(r_{\mathbf{y}},\theta_{\mathbf{y}},\phi_{\mathbf{y}})$ in spherical coordinates, $r_{<}:=min\{r_{\mathbf{x}},r_{\mathbf{y}}\}$, $r_{>}:=max\{r_{\mathbf{x}},r_{\mathbf{y}}\}$ and $Y^{m}_{l}$ refers to the spherical harmonic of order $(l,m)$. Restricting to subset of particles $\mathbf{x}\in t\subset\mathbb{R}^{3},\mathbf{y}\in s\subset\mathbb{R}^{3}$ such that $\frac{radius(t)+radius(s)}{distance(t,s)}\leq\nu$, $\nu>0$ sufficiently small, one may approximate Eq. (2) by truncating this series to a finite order $P$. This allows to derive a fast summation scheme for any $\mathbf{x}\in X$, for problems consisting in computing $\phi$ such that

where $q(\mathbf{y})$ is a charge associated to the particle $\mathbf{y}$. Such a problem in Eq. (3) is referred to as a $N$-body problem and such $t$ and $s$ are named well-separated target ($t$) and source ($s$) cells, where a cell denotes in this paper a cubical subset of $\mathbb{R}^{3}$. These cells are obtained in practice through a recursive splitting of the computational box $B$ into smaller boxes. Assuming that $B$ is a cube, one may divide it into $8$ other cubes and repeat this procedure until the induced cubes (the cells) are sufficiently small. A tree whose nodes are cells and such that the daughters of a given cell are the non-empty cubes coming from its decomposition is named an octree. The root cell is considered to belong at octree level $0$ and any non-root cell belongs at level $E+1$, where $E$ denotes its mother level. In the case of perfect octrees (i.e. with maximal amount of daughter per non-leaf node, that is $8$ daughters), we consider that two cells at the same tree level with strictly positive distance (i.e. non-neighbor cells) are well-separated.

The sum over $\mathbf{y}\in Y_{|s}$ (the restriction of $Y$ to the particles in $s$) can be switched, for any $\mathbf{x}\in X_{|t}$ to the terms depending only on $\mathbf{y}$, due to the separability in the variables $\mathbf{x}$ and $\mathbf{y}$ of Eq. (2). Expression in Eq. (2) results [21] in a three-step summation scheme expressing Eq. (3) in the form

or equivalently under matrix form $U^{*}[t]C[t,s]V[s]q_{|s}$ where $q_{|s}$ refers to the vector of charges of particles in $s$. The result $\mathcal{M}_{s}$ of the product $V[s]q_{|s}$ is named the multipole expansion of source cell $s$. The matrix $V[s]$ involved in this product is named Particle-To-Multipole operator (shortly denoted by P2M) since it maps information associated to particles (i.e. here the charges associated to atoms) to the multipole expansion of the cell in which they belong. This multipole expansion is transformed into a local one $\mathcal{L}_{t}$ by means of product by the Multipole-To-Local operator (shortly denoted by M2L) $C[t,s]$. Similarly, the product by $U^{*}$ is referred to as the Local-To-Particle operator (L2P). Nevertheless, the efficiency of the FMM algorithm is also guaranteed by two additional operators based on the idea that multipole expansions of a given non leaf cell can be approximating through combination of its daughter’s multipole expansions by means of a Multipole-To-Multipole operator (M2M) and local expansions in a given non root cell can be approximated from its mother local expansion thanks to a Local-To-Local operator (L2L). Finally, two leaf-cells that are not well-separated still need to interact, but without expansion. This is done by applying the Particle-To-Particle operator (P2P), i.e. nothing else than a direct computation involving only the particles of these two cells. The relation between all these operators according to the tree structure are depicted in Fig. 1.

In the case of quasi-uniform particle distributions (i.e. with approximately the same density of particles in all the simulation box), octrees can be considered as perfect. In this situation, the FMM algorithm is easy to provide (and usually referred to as the uniform FMM algorithm in the literature), according to the definition of the different FMM operators. This is summarized in Alg. 1, where we used the notation $V[s,s^{\prime}]$ to denote the M2M operator between the source cell $s$ and one of its daughter $s^{\prime}$; $U^{*}[t^{\prime},t]$ to denote the L2L operator between the target cell $t$ and one of its daughter $t^{\prime}$; for any target cell $t$, $\Lambda(t)$ the interaction list of $t$, that is the set of source cells $s$ that are well-separated from $t$ and whose ancestors in the octree are not well-separated from any ancestor of $t$; $\mathscr{T}$ refers to the octree; $\mathcal{L}(\mathscr{T})$ refers to the set of leaves of $\mathscr{T}$ and $K(\mathbf{x},\mathbf{y})=\frac{1}{|\mathbf{x}-\mathbf{y}|}$.

Alg. 1 results in $\mathcal{O}(N)$ complexity [21] with prefactor depending on the user required precision. Notations of Alg. 1 directly provide a way of approximating the influence $\phi_{t^{\prime},s^{\prime}}$ of $s^{\prime}$ a daughter of $s$ on $t^{\prime}$ and daughter of $t$, provided that $t$ and $s$ are well-separated, under matrix form:

This last can be interpreted as low-rank matrix approximation when seeing Eq. (3) as a matrix-vector product with matrix entries corresponding to evaluation of the Coulomb potential on pairs of target and source cells. Notice that the interaction between equally located particles is not well-defined, and we consider in this article that they are just equal to $0$.

In terms of parallel implementation, 3D FMM as shown its ability to scale in distributed memory context. Combined with its linear complexity, this makes this method particularly attractive for MD simulations. However, few limitations still remain: the presentation we made only takes into accounts charges and considered non-periodic case. This last can be solved by means of periodic FMM [11].

Actually, a similar scheme can be derived for any kernel $K(\mathbf{x},\mathbf{y})$ (not only the Coulomb one) that satisfies particular assumption, namely the asymptotically smooth behavior [9]. There exist different approaches to obtain a fast separable formula in a kernel-independent way [46, 18]. Among them, we are especially interested into the interpolation-based FMM (IBFMM) because of its flexibility, its well documented error bounds as well as the particular form of the approximated expansion of $K$ it provides. Indeed, denoting by $S_{k}[t]$ (resp. $S_{l}[s]$) a multivariate Lagrange polynomial associated to the interpolation node $\mathbf{x}_{k}^{t}$ (resp. $\mathbf{y}_{l}^{s}$) in $t$ (resp. $s$), a Lagrange interpolation of $K$ in $t\times s$ writes:

for any $(\mathbf{x},\mathbf{y})\in t\times s$. Eq. (6) clearly exhibits a separable expansion of $K$ in the $\mathbf{x}$ and $\mathbf{y}$ variables. Hence, the restriction of Eq. (3) to particles in $t\times s$ (i.e. to $X_{|t}$ and $Y_{|s}$) and replacing the Coulomb kernel by a generic one $K$) can be approximated by:

by analogy with Eq. (4). This identifies most of the FMM operators. Moreover, interpreting multipole expansions generated this way as charges associated to interpolation nodes (now seen as particles), one can obtain a new $N$-body problem with the same kernel $K$. Hence, this procedure can be repeated all along the tree structure in order to approximate multipole expansions in the non-leaf cells (which is the analogous of the M2M operator). In the same way, evaluating interpolant at a non-leaf cell on interpolation nodes in its daughters allows to derive a analogous of the L2L operator. To be more precise, we assume that the Lagrange interpolation applied on cell $c$ uses the interpolation grid $\Xi_{c}=\{\mathbf{y}^{s}_{0},...,\mathbf{y}^{s}_{P-1}\}$ with $P$ interpolation nodes over $c$. In practice, $\Xi_{c}$ is chosen as product of one-dimensional interpolation nodes (such as Chebyshev nodes [18]). Hence, for any source cell $s$, the IBFMM M2M operator between $s$ and $s^{\prime}$, with $s^{\prime}$ a daughter of $s$, is defined as

Using Eq. (6), we can identify that the IBFMM L2P operator on leaf cell $c$ is the adjoint of the IBFMM P2M operator on this $c$. Hence, for any non-leaf cell $c$, the IBFMM L2L operator between $c$ and a daughter $c^{\prime}$ of $c$ is the adjoint of the IBFMM M2M operator between $c$ and $c^{\prime}$, that is:

As for the classical FMM algorithm, the IBFMM reaches the linear complexity with respect to the number of particles in the system. In the remainder of this article, we consider that $K$ is translationally and rotationally invariant, which is always verified in our applications. A schematic view of IBFMM is provided in Fig. 2.

On one side, the parallelization of PME algorithm in distributed memory can be a practical bottleneck due to the scalability of the FFT [4]. On the other side, there exists fast scalable methods in distributed memory for the evaluation of $N$-body problems involving asymptotically smooth kernels, such as (IB)FMM (see Sect. 2.2). We thus propose a new alternative to PME, based on a very simple idea. Indeed, the kernel

is asymptotically smooth [5], as well as all its derivatives. For sufficiently large $p$, thanks to the absolute convergence of the terms in Eq. (10), we have

with $Z_{p}:=\big{\{}(z_{0},z_{1},z_{2})\in\mathbb{Z}^{3}\hskip 2.84544pt|\hskip 2.84544pt|z_{k}|\leq p,\hskip 2.84544ptk\in[\![0,2]\!]\big{\}}$ and $p$ to be fixed by user.

Our approach is thus very simple: we consider small $\xi$ (typically $\xi=0.01$), so that $\mathcal{E}_{rec}\approx 0$ and can be neglected for our targeted accuracies. In this case, the counterpart relies in the computation of $\mathcal{\tilde{E}}_{real}$ that requires a relatively large amount of periodic images (i.e. a large $p$) to reach the numerical convergence. Since $\mathcal{E}_{self}$ involves $\mathcal{O}(N)$ flops to be computed (and is embarrassingly parallel), only this $\mathcal{\tilde{E}}_{real}$ has a corresponding intensive computational cost. Hopefully, Eq. (12) allows the computation of $\tilde{\mathcal{E}}_{real}$ to be handled through fast methods for $N$-body problems with kernel $\mathfrak{D}_{\mathbf{x}}\mathfrak{D}_{\mathbf{y}}H(\mathbf{x},\mathbf{y})$.

In this section, we consider that $p=0$. Extension to other $p$’s will be done in Sect. 3.3. Let $\mathcal{\tilde{E}}_{real}(t,s):=\sum_{\mathbf{x}\in t}\sum_{\mathbf{y}\in s}\mathfrak{D}_{\mathbf{x}}\mathfrak{D}_{\mathbf{y}}H(\mathbf{x},\mathbf{y})$, where $t,s\in\mathcal{L}(\mathscr{T})$. We can now divide $\mathcal{\tilde{E}}_{real}$ into two parts:

where $dist(t,s)$ denotes the distance between $t$ and $s$. Hence, by increasing the octree depth, one decreases the number of interactions computed inside $\mathcal{N}_{real}$ (near field) while increasing the number of interactions in $\mathcal{F}_{real}$ (far field).

Each $(t,s)\in\mathcal{L}(\mathscr{T})$ such that $dist(t,s)\neq 0$ actually is a well-separated pair of cells. Hence, by means of multivariate polynomial interpolation as in Eq. (6) on $t$ and $s$, one may accurately approximate $\mathcal{\tilde{E}}_{real}(t,s)$ as:

These $\mathcal{Q}_{c}(\mathbf{u})$, with $\mathbf{u}$ any interpolation node in $c$, can be interpreted as modified charges associated to this interpolation node, which can itself be seen as a particle. Hence, one can perform an IBFMM in which the multipole/local expansions are computed using the modified charges while the near field (i.e. $\mathcal{N}_{real}$) is evaluated using the explicit formula for the complete interaction between particles with respect to the kernel $H$ (i.e. involving charges, dipoles and quadrupoles here). In other words, this method (that we named ANKH) consists in the following steps:

1. 1.

   Compute the tree decomposition $\mathscr{T}$ of the simulation box $B$,

2. 2.

   For each leaf $c\in\mathcal{L}(\mathscr{T})$, compute the modified charges in $c$ (using the same interpolation grid up to a translation to the center of $c$),

3. 3.

   Compute $\mathcal{N}_{real}$ using direct computation,

4. 4.

   Compute $\mathcal{F}_{real}$ solving the $N$-body problem defined on interpolation nodes and modified charges,

5. 5.

   Compute $\mathcal{E}_{self}$ using Eq. (10),

6. 6.

   Return $\mathcal{N}_{real}+\mathcal{F}_{real}+\mathcal{E}_{self}$.

The $N$-body problem mentioned in step 4 can be written this way:

where $H_{t,s}=\begin{cases}H\text{ if }dist(t,s)>0,\\ 0\text{ otherwise}\end{cases}$. According to Eq. (14), the quantity into parenthesise can be efficiently computed in linear time by means of IBFMM applied on kernel $H$ and in which we simply avoid the P2P operators. We name this approach ANKH-FMM.

In the case of quasi-uniform particle distribution (actually the kind of particle distributions we are targeting), a perfect octree with depth $~{}log(N)$ generates leaves with $\mathcal{O}(1)$ particles. Hence, in this situation, step 3 of the ANKH method costs $\mathcal{O}(N)$ flops. For such perfect trees, the tree construction (step 1) also is $\mathcal{O}(N)$ and computation of $\mathcal{E}_{self}$ is, of course, also linear in complexity (step 5). Since it is obvious that step 6 is $\mathcal{O}(1)$, the overall complexity is $\mathcal{O}(N)$.

Two lasts questions remain: what about the periodicity and the computation of modified charges? The first of these questions is the purpose of Sect. 3.3 and the second one is the purpose of Sect. 3.4.

Various formulations of the FMM provide efficient ways of dealing with periodicity [35, 11, 32]. We consider a slightly different approach in our method, well-suited for IBFMM. One idea is to perform a classical IBFMM on the computational box as well as on the direct $3^{3}-1$ adjacent images of it. Then, all other distant images are well-separated from the computational box and their interaction can be approximated using only the multipole expansion $\mathcal{M}_{B}$ of this box. More precisely, denoting $C[B,I]$ the M2L matrix between the computational box $B$ and its image $I$, one wants to fastly compute:

The remaining task is thus to sum the M2L matrices between images of B and B itself efficiently. This can be done by means of interpolation over equispaced grid [6, 13].

Indeed, such tool allow to express $C[B,I]$ as a product $\chi^{*}\mathbb{F}^{*}D[B,I]\mathbb{F}\chi$, where $\chi$ is a zero-padding matrix that increases the size of the inner matrix by slightly less than $8$, $\mathbb{F}$ denotes the discrete Fourier matrix (that can be applied through FFT) and $D[B,I]$ is a diagonal matrix that can be computed in linearithmic time with respect to the size of $C[B,I]$.

More precisely, let $\Xi_{B}:=\{\Xi_{B}(i,j,k):=(-r_{B}+2r_{B}i/(L-1),-r_{B}+2r_{B}j/(L-1),-r_{B}+2r_{B}k/(L-1))\hskip 0.28436pt|\hskip 0.28436pti,j,k\in[\![0,L-1]\!]\}$ be a equispaced grid over $B$ with $L$ nodes in each direction (i.e. with $L^{3}$ nodes in total) and let $\Xi_{I}$ be its translation from the origin to the center of the image $I$ of $B$. The matrix $C[B,I]$, in this case, has a block-Toeplitz structure [6, 13]. Any such block-Toeplitz matrix can be embedded into a circulant one of size $(2L-1)^{3}$ whose first row is composed of the $(2L-1)^{3}$ different entries of $C[B,I]$. $\chi\in\{0,1\}^{(2L-1)^{3}\times L^{3}}$ and $D[B,I]\in\mathbb{C}^{(2L-1)^{3}\times(2L-1)^{3}}$ are then defined as follows:

where $R[B,I]$ is the first row of the circulant matrix formed by embedding of the Toeplitz matrix $C[B,I]$, i.e.

Since $\mathbb{F}$ is practically performed through FFT, computation of $D[B,I]$ results in $\mathcal{O}(L^{3}log(L))$ flops.

We can thus reformulate Eq. (15) as

As a sum of diagonal matrix, $D_{B}$ clearly is diagonal. We can switch the sum in Eq. (18) to the vector $R[B,I]$ thanks to Eq. (16), resulting into

Since the image $I$ of $B$ is a translation of $B$ by $\mathbf{t}=2r_{B}I$, we have that $\Xi_{I}(\hat{i},\hat{j},\hat{k})=\Xi_{B}(\hat{i},\hat{j},\hat{k})+\mathbf{t}$, so that, thanks to Eq. (17) and the translational invariance of $H$,

Notice that $-I$ also lies into the set of images of $B$ so that one can eliminate the minus sign in Eq. (20). Also, $\{\mathbf{z}^{(\mathbf{0})}_{i,j,k}\hskip 0.28436pt|\hskip 0.28436pti,j,k\in[\![0,(2L-1)]\!]\}$ can be interpreted as a particle distribution over $B$ and its direct neighbors. Moreover, applying the sum over $I$ of Eq. (18) to Eq. (20) gives

which can be interpreted as a $N$-body problem with all charges set to $1$ over two particle distributions: the set $P_{target}$ of $(2L-1)^{3}$ different $\mathbf{z}^{(\mathbf{0})}_{ijk}$’s and the set $P_{source}$ of $\mathbf{t}$’s. For large $p$, using a hierarchical method for such problem, one may compute this quantities very efficiently since the first of these particle distribution is already well-separated from each $\mathbf{t}$. In addition, both $H$, $P_{target}$ and $P_{source}$ are invariant under the action of the rotation group that preserves the cube. This can be exploited to speed-up the computation [13]. See Fig. 3 for a schematic representation.

A critical point here (that actually motivates this entire subsection) is that $D_{B}$ does not depend on the particle distribution nor its modified charges, but only relies on the size of the simulation box $B$. Hence, the computation of $D_{B}$ can be done only once inside a precomputation step, not at each IBFMM call. This has a particular interest since in daylife molecular dynamics applications, such IBFMM calls on the same $B$ would be numerous, but needing only the single precomputation of $D_{B}$ with our method.

In terms of complexity, at runtime, $D[B,I]$ is already computed, so the cost of handling the periodicity this way is the one of two FFTs on a small grid of size $(2L-1)^{3}$. Since $L$ is a (small) constant, these FFTs cost $\mathcal{O}(L^{3}log(L))=\mathcal{O}(1)$ flops.

The idea of switching derivatives to the interpolation polynomials has already found applications in the literature [30]. In our case, things are slightly more difficult because we consider differential operators instead of simple derivatives. To do so, we consider a Lagrange interpolation over product of 1D Chebyshev grids in each leaf cell as well as all along the tree structure with a fixed order $L$ in each direction. Notice that this interpolation is different from the one of Sect. 3.3, that performs over a equispaced grid and only for the root. However, this does not impact the method at all since the root is only involved in interactions between non-adjacent images of $B$ (see Fig. 3) and because the definition of the M2M/L2L operators is valid no matter the interpolation grid type is.

The choice of multivariate Lagrange interpolation over products of 1D Chebyshev nodes is motivated by the numerical instability that arises on large interpolation order and at fixed arithmetic precision when dealing with equispaced grids. Moreover, this choice of interpolation is known to be particularly good in the context of multilevel summation schemes.

As presented in Sect. 3.2, in the ANKH approach, the differential operators apply on interpolation polynomials directly. Any Lagrange polynomial over a product grid is the product of the Lagrange polynomials in each one-dimensional grids. Assuming for the sake of clarity (One only has to introduce a scaling function in order to recover the general case [30, 6]) that the one-dimensional grids are defined on $[-1,1]$, we thus have the following decomposition of such polynomial $S$:

where $S_{iL^{2}+jL+k}$ can be either $S_{k}[t]$ or $S_{l}[s]$ following the notations of Eq. (6), $S_{i}$’s are one-dimensional Lagrange polynomials over Chebyshev (i.e. Lagrange-Chebyshev polynomials) nodes and $\mathbf{x}=(x_{0},x_{1},x_{2})\in\mathbb{R}^{3}$. Explicit forms [18] are known for these polynomials $S_{i}$:

where $T_{m}$ is the $m^{th}$ Chebyshev polynomial and $r_{k}:=cos\left(\frac{2k-1}{2L}\right)$ is the $k^{th}$ Chebyshev node. $T_{m}(x)$ can be computed using the recurrence relation:

Problem is that explicit derivation of Eq. (23) leads to

which is numerically unstable near the bounds of $[-1,1]$. This can here be solved by using Chebyshev polynomials of the second kind in the derivation, but the same kind of problem appears when trying to derive $S_{k}$ once again (which is needed to take into account quadrupoles). We thus opted for another way of deriving these polynomials: there exists [3] a upper triangular matrix $\mathcal{D}$ such that:

This can trivially be generalized to any derivation order, giving:

This last form does not suffer from numerical instability near the bounds of $[-1,1]$. To exploit it in ANKH, we start by reformulating the set of Lagrange-Chebyshev polynomials of Eq. (23) in matrix form

where $\mathbb{H}$ does not depend on $x$ and be be precomputed only once. We can now combine Eq. (27) and Eq. (28) in order to obtain our general formula:

To extend this technique to the three-dimensional case, a direct application of Eq. (22) provides expressions of the multivariate polynomials.

We summarize using $C$++-like pseudocode in Alg. 2 the steps needed to compute the modified charges for multipoles up to quadrupoles (trivial generalizations can be extrapolated from this pseudo-code) using numerical differentiation through derivation of Chebyshev polynomials. Notice that the different $\mathbf{T}$’s computed in Alg. 2 could be stacked in practice inside a matrix with $9$ columns before applying the product by $\mathbb{H}$, which allows to benefit from matrix-matrix products instead of matrix-vector ones (i.e. BLAS3 instead of BLAS2), that should better perform.