Energy Slicing Theory

Potential Energy Slicing

This section contains the formulation used to implement the SlicedNonbondedForce class. For simplicity, the equations presented in here follow the conventional 1-based indexing, whereas the class API follows Python’s 0-based indexing.

Slicing the potential energy starts with partitioning all N particles of a system into n disjoint subsets or “colors”. This procedure results in \(n(n+1)/2\) slices, distinguished by order-invariant pairs of subset indices. Slice I, J comprises all pairs involving a particle in subset I and a particle in another (or the same) subset J.

Example: three particle subsets (●●●) form six energy slices

It is straightforward to slice pairwise Lennard-Jones interactions, as well as the direct-space, self-energy, and exclusion/exception parts of the Ewald summation. Therefore, the present section focuses on the reciprocal-space and background energy parts.

Sliced Background Energy

In the case of systems with non-zero net charge, it is necessary to add the energy of interaction with a neutralizing background charge. The background energy is given by [1]

\[E_{bg} = \frac{Q^2}{8\epsilon_0 V \alpha^2}.\]

where \(Q = \sum_{i=1}^N q_i\) is the total charge of the system. We can split the total charge into the sum of the charges in each subset,

\[Q = \sum_{I=1}^n Q_I,\]

where \(Q_I = \sum_{i \in I} q_i\) is the charge of subset I.

The goal is to express the background energy as

\[E_{bg} = \sum_{I=1}^n \sum_{J=I}^n \lambda^{elec}_{I,J} E^{bg}_{I,J},\]

where \(E^{bg}_{I,J}\) is the background contribution of Slice I, J. For this, we rewrite the square of the total charge as

\[Q^2 = \sum_{I=1}^n \sum_{J=1}^n Q_I Q_J.\]

Therefore, the background energy of a slice I, J is given by

\[E^{bg}_{I,J} = (2 - \delta_{I,J}) \frac{Q_I Q_J}{8\epsilon_0 V \alpha^2},\]

where \(\delta_{I,J}\) is the Kronecker delta. The prefactor \(2 - \delta_{I,J}\) allows us to run the summation over J only for \(J \geq I\).

Sliced Reciprocal-Space Energy

For a simulation box with edge matrix \(\mathbf L\) containing N particles under periodic boundary conditions, the standard reciprocal part of the electrostatic potential energy can be expressed as

\[E_{rec} = \frac{2\pi}{\epsilon_0 V} \sum_{\mathbf n \neq \mathbf 0} \frac{e^{-\frac{k^2}{4\alpha^2}}}{k^2} \sum_{i=1}^N \sum_{j=1}^N q_i q_j e^{\text{ⅈ} {\mathbf k}\cdot({\mathbf r}_i - {\mathbf r}_j)},\]

where \(\epsilon_0\) is the vacuum permittivity, \(V\) is the box volume, \(\alpha\) is the Ewald splitting parameter, \(\mathbf n \in \mathbb Z^3\) is an integer lattice vector, \(\mathbf k = 2\pi \mathbf L^{-1}{\mathbf n}\) is a reciprocal space wave vector, \(k = \|\mathbf k\|\) is the norm of \(\mathbf k\), and \(\text{ⅈ} = \sqrt{-1}\) is the imaginary unit. As the summations over indices i and j run for all particles, we can write

\[E_{rec} = \frac{2\pi}{\epsilon_0 V} \sum_{\mathbf n \neq \mathbf 0} \frac{e^{-\frac{k^2}{4\alpha^2}}}{k^2} \Bigg(\sum_{i=1}^N q_i e^{\text{ⅈ} {\mathbf k}\cdot{\mathbf r}_i} \Bigg) \Bigg(\sum_{j=1}^N q_j e^{-\text{ⅈ} {\mathbf k}\cdot{\mathbf r}_j}\Bigg).\]

The first term inside parentheses is the charge structure factor

\[S(\mathbf k) = \sum_i q_i e^{\text{ⅈ} {\mathbf k}\cdot{\mathbf r}_i}\]

and the second one is its complex conjugate \({\overline S}(\mathbf k)\), so that we can write

\[E_{rec} = \frac{2\pi}{\epsilon_0 V} \sum_{\mathbf n \neq \mathbf 0}\frac{e^{-\frac{k^2}{4\alpha^2}}}{k^2} |S(\mathbf k)|^2.\]

The goal is to express the reciprocal-space energy as

\[E_{rec} = \sum_{I=1}^n \sum_{J=I}^n \lambda^{elec}_{I,J} E^{rec}_{I,J},\]

where \(E^{rec}_{I,J}\) is the reciprocal-space contribution of Slice I, J.

The structure factor of a subset I can be computed as

\[S_I(\mathbf k) = \sum_{i \in I} q_i e^{\text{ⅈ} {\mathbf k}\cdot{\mathbf r}_i}.\]

Since we have disjoint subsets, this makes \(S(\mathbf k) = \sum_{I=1}^n S_I(\mathbf k)\). Substituting into \(S(\mathbf k) {\overline S}(\mathbf k)\) and expanding the product, all imaginary terms cancel out, as expected. The reciprocal-space energy then becomes

\[E_{rec} = \frac{2\pi}{\epsilon_0 V} \sum_{\mathbf n \neq \mathbf 0} \frac{e^{-\frac{k^2}{4\alpha^2}}}{k^2} \sum_{I=1}^n \sum_{J=1}^n S_I(\mathbf k) \cdot S_J(\mathbf k),\]

where \(x \cdot y = \text{Re}(x)\text{Re}(y) + \text{Im}(x)\text{Im}(y)\) for two complex numbers x and y. It is now clear that the energy of a slice I, J should be calculated by

\[E^{rec}_{I,J} = (2-\delta_{I,J}) \frac{2\pi}{\epsilon_0 V} \sum_{\mathbf n \neq \mathbf 0} \frac{e^{-\frac{k^2}{4\alpha^2}}}{k^2} S_I(\mathbf k) \cdot S_J(\mathbf k).\]

Sliced PME Implementation

Each subset-specific structure factor \(S(\mathbf k)\) can be calculated using the smooth Particle Mesh Ewald method [2]. This requires n FFT calculations for evaluating all slices. In the case of forces, n additional inverse FFT calculations are necessary. Therefore, in a single processor and for a fixed number of particles, the computational effort scales approximately linearly with the number of subsets (not the number of slices). In a GPU, as with the OpenCL and CUDA instances of openmm.Platform, we take advantage the capability of either package vkFFT or cuFFT to speed-up parallel FFT calculations via batched transforms.

Performance

The tables below compare SlicedNonbondedForce with the openmm.NonbondedForce it replaces in two systems of about 50,000 atoms, using PME with a 1 nm cutoff on the CUDA platform in mixed precision (NVIDIA RTX 4090, OpenMM 8.6). With n subsets, the interactions between subset 0 and each other subset are scaled by two parameters, one for Coulomb and one for Lennard-Jones, as one would do to decouple a ligand from its environment. A step is a LangevinMiddleIntegrator step of 2 fs, which computes only forces. An evaluation is a call to getState with getForces=True and getEnergy=True. Columns marked with derivatives request the derivatives with respect to all 2(n-1) parameters and, for evaluations, pass getParameterDerivatives=True. Times are in milliseconds and can be reproduced with devtools/benchmarks/benchmark.py.

The first system is a box of 16,814 TIP3P water molecules (50,442 atoms, 8 nm) split into n contiguous blocks of whole molecules.

Force	Subsets	Derivatives	Step	Step with derivatives	Evaluation	Evaluation with derivatives
NonbondedForce	–	–	0.26	–	0.57	–
SlicedNonbondedForce	1	–	0.29	–	0.63	–
SlicedNonbondedForce	2	2	0.36	0.36	0.70	1.40
SlicedNonbondedForce	4	6	0.65	0.65	0.97	2.27
SlicedNonbondedForce	8	14	1.06	1.14	1.45	4.33

The second system is the TYK2 kinase (Amber ff14SB) bound to ligand ejm_31 (OpenFF 2.3) from the OpenFF protein-ligand benchmark, in TIP3P water with 0.15 M NaCl (51,405 atoms, 8.15 nm cubic box). It is split into ligand (subset 0, 32 atoms), protein (4,652 atoms) and solvent.

Force	Subsets	Derivatives	Step	Step with derivatives	Evaluation	Evaluation with derivatives
NonbondedForce	–	–	0.27	–	0.61	–
SlicedNonbondedForce	3	4	0.41	0.41	0.74	1.66

The cost of a step grows sublinearly with the number of subsets: going from 1 to 8 subsets makes a step about 3.6 times as expensive. Since derivatives are only computed along with the energy, requesting them barely affects steps (8% with 14 derivatives). The cost of an evaluation with derivatives grows with their number, not with the number of slices: a parameter that scales several slices adds a single derivative. Derivatives are best requested only for the parameters whose energies are needed, and only in the steps where they are reported.