Electrostatic Interactions#

Electrostatic interactions arise from Coulombic forces between charged particles. In periodic systems, the \(1/r\) potential decays slowly, requiring special techniques to handle the conditionally convergent lattice sum. ALCHEMI Toolkit-Ops provides GPU-accelerated implementations of Ewald summation, Particle Mesh Ewald (PME), and Damped Shifted Force (DSF) electrostatics via NVIDIA Warp, with PyTorch and JAX autograd support for machine learning applications (Ewald and PME support full position/charge/cell autograd; DSF provides charge gradients via autograd and computes forces/virials analytically).

Tip

For periodic systems, start with ewald_summation() (PyTorch) / ewald_summation() (JAX) or particle_mesh_ewald() (PyTorch) / particle_mesh_ewald() (JAX). For non-periodic systems or large-scale simulations, consider approximate dsf_coulomb() (PyTorch only) which provides \(O(N)\) scaling with smooth force continuity at the cutoff.

Overview of Available Methods#

ALCHEMI Toolkit-Ops provides electrostatics modules for point charges:

Method	Scaling	Best For
Ewald Summation	\(O(N^2)\)	Small/medium systems (<5000 atoms)
Particle Mesh Ewald	\(O(N \log N)\)	Large periodic systems
Damped Shifted Force (DSF)	\(O(N)\)	Large systems, non-periodic
Direct Coulomb	\(O(N^2)\)	Non-periodic or as real-space component
Ewald Multipole	\(O(N^2)\)	Multipolar systems, small/medium
PME Multipole	\(O(N \log N)\)	Multipolar systems, large

All methods support:

Single-system and batched calculations
Periodic boundary conditions
Automatic differentiation (see per-method details below)
Both neighbor list (COO) and neighbor matrix formats

Method	Position gradients	Charge gradients	Cell gradients
Ewald / PME	Autograd	Autograd	Autograd
Direct Coulomb	Autograd	Autograd	Autograd
DSF	Analytical forces	Analytical (straight-through)	Analytical virial (PBC)

Data Formats#

Tensor Specifications#

The table below lists out the general syntax and expected shapes for tensors used in the electrostatics code. When possible to do so, we encourage developers and users to align their variable naming to what is shown here for ease of debugging and consistency.

Tensor	Shape	Dtype	Description
`positions`	`(N, 3)`	`float64/float32`	Atomic coordinates
`charges`	`(N,)`	`float64`	Atomic partial charges
`cell`	`(1, 3, 3)` or `(B, 3, 3)`	`float64`	Unit cell lattice vectors (rows)
`pbc`	`(1, 3)` or `(B, 3)`	`bool`	Periodic boundary conditions per axis
`batch_idx`	`(N,)`	`int32`	System index for each atom (batched only)
`alpha`	`float` or `(B,)` tensor	`float64`	Ewald splitting parameter

Output Data Types#

Energies are always computed and returned in float64 for numerical stability during accumulation. Forces, virial, and charge gradients match the input precision – float32 when positions are float32, float64 when positions are float64.

Neighbor Representations#

The electrostatics functions accept neighbors in two formats:

Neighbor List (COO): Shape (2, num_pairs) where row 0 contains source indices and row 1 contains target indices. Each pair is listed once. Provide with neighbor_list and neighbor_shifts arguments.

Neighbor Matrix: Shape (N, max_neighbors) where each row contains neighbor indices for that atom, padded with fill_value. Provide with neighbor_matrix and neighbor_matrix_shifts arguments.

Tip

See the neighbor list documentation for API usage and performance considerations when deciding between COO and matrix representations.

Ewald Summation#

Mathematical Background#

The Ewald method splits the slowly-converging Coulomb sum into four components:

\[E_{\text{total}} = E_{\text{real}} + E_{\text{reciprocal}} - E_{\text{self}} - E_{\text{background}}\]

Real-Space (Short-Range):

\[E_{\text{real}} = \frac{1}{2} \sum_{i \neq j} q_i q_j \frac{\text{erfc}(\alpha r_{ij})}{ r_{ij}}\]

The complementary error function \(\text{erfc}(\alpha r)\) rapidly damps interactions beyond approximately \(r \approx 3/\alpha\), confining contributions to a local neighborhood.

Reciprocal-Space (Long-Range):

\[E_{\text{reciprocal}} = \frac{1}{2V} \sum_{\mathbf{k} \neq 0} \frac{4\pi}{k^2} \exp\left(-\frac{k^2}{4\alpha^2}\right) |S(\mathbf{k})|^2\]

where the structure factor is:

\[S(\mathbf{k}) = \sum_j q_j \exp(i\mathbf{k} \cdot \mathbf{r}_j)\]

Self-Energy Correction:

\[E_{\text{self}} = \frac{\alpha}{\sqrt{\pi}} \sum_i q_i^2\]

Removes the spurious self-interaction introduced by the Gaussian charge distribution.

Background Correction (for non-neutral systems):

\[E_{\text{background}} = \frac{\pi}{2\alpha^2 V} Q_{\text{total}}^2\]

Usage Examples#

Explicit Parameters#

PyTorch

from nvalchemiops.torch.interactions.electrostatics import ewald_summation

energies, forces = ewald_summation(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=0.3,        # Ewald splitting parameter
    k_cutoff=8.0,     # Reciprocal-space cutoff in inverse length
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.electrostatics import ewald_summation

energies, forces = ewald_summation(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=0.3,        # Ewald splitting parameter
    k_cutoff=8.0,     # Reciprocal-space cutoff in inverse length
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

Automatic Parameter Estimation#

When alpha or k_cutoff are not provided, they are estimated based on accuracy:

PyTorch

energies, forces = ewald_summation(
    positions=positions,
    charges=charges,
    cell=cell,
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    accuracy=1e-6,  # Target relative error
    compute_forces=True,
)

JAX

energies, forces = ewald_summation(
    positions=positions,
    charges=charges,
    cell=cell,
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    accuracy=1e-6,  # Target relative error
    compute_forces=True,
)

The estimation uses the Kolafa-Perram formula:

\[\eta = \left(\frac{V^2}{N}\right)^{1/6} / \sqrt{2\pi}\]

\[\alpha = \frac{1}{2 \cdot \eta}, \quad r_{\text{cutoff}} = \sqrt{-2 \ln \varepsilon} \cdot \eta, \quad k_{\text{cutoff}} = \sqrt{-2 \ln \varepsilon} / \eta\]

Tip

Refer to the Parameter Estimation section for API usage.

Separating real- and reciprocal-space#

When either components are required individually, the following code can be used instead of the high level wrapper to compute the contributions directly:

PyTorch

from nvalchemiops.torch.interactions.electrostatics import (
    ewald_real_space, ewald_reciprocal_space, generate_k_vectors_ewald_summation,
)

# Real-space only (short-range, damped Coulomb)
real_energies, real_forces = ewald_real_space(
    positions, charges, cell, alpha=0.3,
    neighbor_list=neighbor_list, neighbor_shifts=neighbor_shifts,
)

# Reciprocal-space only (long-range, smooth)
alpha = torch.tensor([0.3], dtype=positions.dtype, device=positions.device)
k_vectors = generate_k_vectors_ewald_summation(cell, k_cutoff=8.0)
recip_energies, recip_forces = ewald_reciprocal_space(
    positions, charges, cell, k_vectors, alpha, compute_forces=True,
)

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.electrostatics import ewald_real_space, ewald_reciprocal_space

# Real-space only (short-range, damped Coulomb)
real_energies, real_forces = ewald_real_space(
    positions, charges, cell, alpha=0.3,
    neighbor_list=neighbor_list, neighbor_shifts=neighbor_shifts,
)

# Reciprocal-space only (long-range, smooth)
recip_energies, recip_forces = ewald_reciprocal_space(
    positions, charges, cell, alpha=0.3, k_cutoff=8.0,
)

Note

The sum of real and reciprocal components gives the Ewald energy. The self-energy and background corrections are embedded within the reciprocal energy.

Particle Mesh Ewald (PME)#

Mathematical Background#

For very large atomic systems, the particle mesh Ewald (PME) algorithm provides substantial improvements in computational performance over conventional Ewald summation. PME accelerates the reciprocal-space sum using fast Fourier transforms by:

Charge Assignment: Spread charges onto a mesh using B-spline interpolation
Forward FFT: Transform charge mesh to reciprocal space
Convolution: Multiply by Green’s function in k-space
Inverse FFT: Transform back to get potentials/electric field
Force Interpolation: Gather forces at atomic positions

The B-spline interpolation introduces errors corrected by the influence function:

\[G(\mathbf{k}) = \frac{2\pi}{V} \cdot \frac{\exp(-k^2 / 4\alpha^2)}{k^2} \cdot \frac{1}{C^{2p}(\mathbf{k})}\]

where \(C(\mathbf{k})\) is the B-spline correction factor and \(p\) is the spline order.

Usage Examples#

Basic Usage#

PyTorch

from nvalchemiops.torch.interactions.electrostatics import particle_mesh_ewald

energies, forces = particle_mesh_ewald(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=0.3,
    mesh_dimensions=(32, 32, 32),  # FFT mesh size
    spline_order=4,                 # B-spline order (4 = cubic)
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.electrostatics import particle_mesh_ewald

energies, forces = particle_mesh_ewald(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=0.3,
    mesh_dimensions=(32, 32, 32),  # FFT mesh size
    spline_order=4,                 # B-spline order (4 = cubic)
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

Mesh Spacing#

Instead of explicit mesh dimensions, specify mesh spacing:

PyTorch

energies, forces = particle_mesh_ewald(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=0.3,
    mesh_spacing=0.5,  # Angstrom (or your length unit)
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

JAX

energies, forces = particle_mesh_ewald(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=0.3,
    mesh_spacing=0.5,  # Angstrom (or your length unit)
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

Automatic Parameter Estimation#

Similar to the Ewald summation interface, PME accepts an accuracy parameter that can be used to automatically determine sensible \(\alpha\) and mesh:

PyTorch

energies, forces = particle_mesh_ewald(
    positions=positions,
    charges=charges,
    cell=cell,
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    accuracy=4e-5,  # Estimates alpha and mesh dimensions
    compute_forces=True,
)

JAX

energies, forces = particle_mesh_ewald(
    positions=positions,
    charges=charges,
    cell=cell,
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    accuracy=4e-5,  # Estimates alpha and mesh dimensions
    compute_forces=True,
)

Note

We encourage users to properly benchmark performance gains afforded by accuracy on their systems of interest. The lower the value of accuracy the more precise, at the cost of higher computational requirements.

PME vs Ewald: When to Use Each#

Criterion	Ewald	PME
System size	\(<5000\) atoms	Any size
Scaling	\(O(N^2)\)	\(O(N \log N)\)
Setup overhead	Lower	Higher (FFT setup)
Accuracy control	`k_cutoff`	Mesh resolution
Memory	Low	Mesh memory \((n_x \times n_y \times n_z)\)

For small systems, direct Ewald may be faster due to lower overhead. For large systems, PME’s \(O(N \log N)\) scaling provides substantial speedup.

Damped Shifted Force (DSF)#

Note

DSF Coulomb bindings are currently available for PyTorch only. See JAX electrostatics API for available JAX functions.

Motivation#

Standard truncation of the \(1/r\) Coulomb potential at a cutoff radius introduces two fundamental problems in molecular simulations:

Charge imbalance: The truncation sphere is generally not charge-neutral, causing long-range potential oscillations and systematic errors in thermodynamic properties.
Force discontinuities: Atoms crossing the cutoff boundary experience instantaneous jumps in force, injecting energy into the system and violating energy conservation during molecular dynamics.

The Damped Shifted Force (DSF) method, introduced by Fennell and Gezelter (2006), solves both problems through a pairwise, real-space \(\mathcal{O}(N)\) electrostatic summation technique. The core idea (building on the earlier Wolf summation) is that the neglected environment beyond the cutoff can be approximated from local structure: a neutralizing “image charge” is placed on the surface of the cutoff sphere for every charge within it. A shifted-force construction then ensures both the potential energy and the force smoothly vanish at the cutoff radius \(R_c\).

Tip

DSF is particularly well-suited for non-periodic systems (clusters, droplets, interfaces) and extremely large systems where the \(\mathcal{O}(N)\) scaling provides significant speedups over Ewald-based methods.

Mathematical Background#

Shifted-Force Construction#

For a generic pair potential \(v(r)\), the shifted-force form ensures both the potential and its derivative (force) vanish at the cutoff:

\[V_{\text{SF}}(r) = v(r) - v(R_c) - v'(R_c)(r - R_c), \quad r \le R_c\]

This guarantees \(V_{\text{SF}}(R_c) = 0\) and \(F_{\text{SF}}(R_c) = -V'_{\text{SF}}(R_c) = 0\).

For DSF, the base kernel is the damped Coulomb interaction \(v(r) = \text{erfc}(\alpha r) / r\), where the complementary error function screens the interaction similarly to the real-space part of Ewald summation.

DSF Pair Potential#

The potential energy for a pair of charges \(i\) and \(j\) at distance \(r_{ij} \le R_c\):

\[V_{\text{DSF}}(r_{ij}) = q_i q_j \left[ \frac{\text{erfc}(\alpha r_{ij})}{r_{ij}} - \frac{\text{erfc}(\alpha R_c)}{R_c} + \left( \frac{\text{erfc}(\alpha R_c)}{R_c^2} + \frac{2\alpha}{\sqrt{\pi}} \frac{e^{-\alpha^2 R_c^2}}{R_c} \right)(r_{ij} - R_c) \right]\]

For \(r_{ij} > R_c\), \(V_{\text{DSF}}(r_{ij}) = 0\).

The three terms have clear physical interpretations:

Damped Coulomb (\(\text{erfc}(\alpha r)/r\)): The screened interaction between the charges.
Potential shift (\(-\text{erfc}(\alpha R_c)/R_c\)): Charge neutralization on the cutoff sphere, ensuring \(V(R_c) = 0\).
Force shift (linear in \(r - R_c\)): Ensures the derivative (force) also vanishes at \(R_c\), preventing energy drift.

DSF Force#

The force between charges at distance \(r_{ij} \le R_c\):

\[\mathbf{F}_{\text{DSF}}(r_{ij}) = q_i q_j \left[ \left( \frac{\text{erfc}(\alpha r_{ij})}{r_{ij}^2} + \frac{2\alpha}{\sqrt{\pi}} \frac{e^{-\alpha^2 r_{ij}^2}}{r_{ij}} \right) - \left( \frac{\text{erfc}(\alpha R_c)}{R_c^2} + \frac{2\alpha}{\sqrt{\pi}} \frac{e^{-\alpha^2 R_c^2}}{R_c} \right) \right] \frac{\mathbf{r}_{ij}}{r_{ij}}\]

The subtracted constant ensures the force magnitude is exactly zero at \(r_{ij} = R_c\).

Self-Energy Correction#

Each charge interacts with its own neutralizing image charge on the cutoff sphere. This self-energy must be subtracted:

\[U_i^{\text{self}} = -\left( \frac{\text{erfc}(\alpha R_c)}{2 R_c} + \frac{\alpha}{\sqrt{\pi}} \right) q_i^2\]

Total System Energy#

The total DSF electrostatic energy is:

\[U_{\text{elec}} = \frac{1}{2} \sum_{i} \sum_{j \neq i} V_{\text{DSF}}(r_{ij}) + \sum_i U_i^{\text{self}}\]

Note

The implementation assumes a full neighbor list where each pair \((i, j)\) appears in both directions. The factor of \(1/2\) accounts for this double counting.

Usage Examples#

Basic Energy and Forces#

from nvalchemiops.torch.interactions.electrostatics import dsf_coulomb
from nvalchemiops.torch.neighbors import neighbor_list

# Build full neighbor list
neighbor_list_coo, neighbor_ptr, neighbor_shifts = neighbor_list(
    positions, cutoff=10.0, cell=cell, pbc=pbc, return_neighbor_list=True
)

# Compute DSF energy and forces
energy, forces = dsf_coulomb(
    positions=positions,
    charges=charges,
    cutoff=10.0,
    alpha=0.2,
    cell=cell,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
    unit_shifts=neighbor_shifts,
    compute_forces=True,
)

With Periodic Boundary Conditions and Virial#

energy, forces, virial = dsf_coulomb(
    positions=positions,
    charges=charges,
    cutoff=10.0,
    alpha=0.2,
    cell=cell,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
    unit_shifts=neighbor_shifts,
    compute_forces=True,
    compute_virial=True,
)
# energy: (num_systems,), dtype=float64
# forces: (num_atoms, 3), dtype matches input
# virial: (num_systems, 3, 3), dtype matches input

Using Neighbor Matrix Format#

from nvalchemiops.torch.neighbors import cell_list

# Build neighbor matrix
neighbor_matrix, num_neighbors, shifts = cell_list(
    positions, cutoff=10.0, cell=cell, pbc=pbc
)

energy, forces = dsf_coulomb(
    positions=positions,
    charges=charges,
    cutoff=10.0,
    alpha=0.2,
    cell=cell,
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=shifts,
    compute_forces=True,
)

Charge Gradients for MLIP Training#

For machine learning interatomic potentials (MLIPs) with geometry-dependent charges, DSF supports charge gradient computation through PyTorch autograd:

# Charges predicted by a neural network (requires_grad flows from the model)
charges = charge_model(positions, atomic_numbers)

energy, forces = dsf_coulomb(
    positions=positions,
    charges=charges,
    cutoff=12.0,
    alpha=0.2,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
)

# Backpropagate through charges
loss = (energy - ref_energy).pow(2).sum()
loss.backward()
# charges.grad now contains dE/dq * dloss/dE

Note

Charge gradients (\(\partial E / \partial q_i\)) are computed analytically by the Warp kernel and propagated through PyTorch autograd via a “straight-through trick.” The returned energy tensor is not differentiable with respect to positions or cell through autograd – forces and virials are computed analytically by the kernel.

Batched Calculations#

import torch
from nvalchemiops.torch.interactions.electrostatics import dsf_coulomb

# Concatenate atoms from multiple systems
positions = torch.cat([pos_sys0, pos_sys1])
charges = torch.cat([charges_sys0, charges_sys1])

# System index for each atom
batch_idx = torch.cat([
    torch.zeros(len(pos_sys0), dtype=torch.int32),
    torch.ones(len(pos_sys1), dtype=torch.int32),
]).to(positions.device)

energy, forces = dsf_coulomb(
    positions=positions,
    charges=charges,
    cutoff=10.0,
    alpha=0.2,
    batch_idx=batch_idx,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
    num_systems=2,
)
# energy: (2,) -- per-system energies
# forces: (N, 3) -- per-atom forces

Undamped Shifted-Force Coulomb (alpha=0)#

Setting \(\alpha = 0\) reduces DSF to a shifted-force bare Coulomb interaction (since \(\text{erfc}(0) = 1\) and \(e^0 = 1\)):

energy, forces = dsf_coulomb(
    positions=positions,
    charges=charges,
    cutoff=12.0,
    alpha=0.0,  # Undamped: shifted-force 1/r
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
)

Parameter Guidance#

The accuracy of the DSF method is controlled by two parameters:

Parameter	Typical Range	Guidance
\(R_c\) (cutoff)	10–15	12 is a common standard; 15 recommended for higher precision
\(\alpha\) (damping)	0.0–0.25	Controls convergence vs. accuracy trade-off

Damping parameter regimes:

\(\alpha = 0.0\) (undamped): Best for structural properties (RDFs) and absolute force magnitudes. Simplest form; no erfc damping overhead.
\(\alpha \approx 0.2\text{--}0.25\): Best for long-time dynamics, collective motions, and dielectric properties. Accelerates convergence with cutoff but over-damping should be avoided.

Important

A practical convergence heuristic is to monitor \(\text{erfc}(\alpha R_c)\):

Most applications: \(\text{erfc}(\alpha R_c) < 10^{-3}\) is adequate. For example, \(\alpha = 0.2\) and \(R_c = 12\) gives \(\text{erfc}(2.4) \approx 5 \times 10^{-4}\).
High precision: \(\text{erfc}(\alpha R_c) < 10^{-5}\) is recommended. For example, \(\alpha = 0.2\) and \(R_c = 15\) gives \(\text{erfc}(3.0) \approx 2 \times 10^{-5}\).

When to Use DSF#

Criterion	DSF	Ewald	PME
Scaling	\(O(N)\)	\(O(N^2)\)	\(O(N \log N)\)
Periodicity required	No	Yes	Yes
Force continuity at cutoff	Yes	Depends on cutoff	Depends on cutoff
Self-energy correction	Built-in	Separate term	Separate term
Best for	Large systems, clusters, non-periodic	Small periodic systems	Large periodic systems
Charge gradients (dE/dq)	Analytic, via straight-through	Via autograd	Via autograd

Choose DSF when:

The system is non-periodic (clusters, droplets, interfaces) where Ewald/PME would require artificial periodic boundary conditions.
The system is extremely large and the \(O(N)\) scaling provides significant speedups and memory savings over PME.
Training MLIPs with geometry-dependent charges where analytic \(\partial E / \partial q\) is needed for backpropagation.

Choose Ewald/PME when:

High accuracy of long-range electrostatics is critical for the target property (e.g., dielectric constants, free energies of solvation).
The system is periodic and relatively small (\(< 5000\) atoms), where Ewald’s lower overhead may be advantageous.

Applicability and Limitations#

Applicability:

Large-scale MD simulations with approximate Coulomb
Non-periodic and partially periodic systems (clusters, droplets, surfaces, interfaces)

Limitations:

Dielectric properties: May slightly underestimate the dielectric constant in some liquids if the cutoff is too small or damping too high. Typical cutoffs of 12–15 provide adequate accuracy for most systems.
Molecular torques: Over-damping (\(\alpha > 0.3\)) can degrade the accuracy of torques in molecular systems. Keep \(\alpha \le 0.25\) for molecular simulations.
Low-frequency phonons: In crystal lattices, undamped DSF may deviate slightly from Ewald results for very low-frequency modes, though \(\alpha \approx 0.2\) typically resolves this.

Software Ecosystem#

The DSF method is widely implemented and validated across major simulation packages, including LAMMPS (pair_style coul/dsf), OpenMD, DL_POLY, Cassandra, JAX-MD, and CP2K. This broad adoption provides extensive cross-validation of the method and its parameters.

References#

Fennell, C. J.; Gezelter, J. D. (2006). “Is the Ewald summation still necessary? Pairwise alternatives to the accepted standard for long-range electrostatics.” J. Chem. Phys. 124, 234104. DOI: 10.1063/1.2206581
Wolf, D.; Keblinski, P.; Phillpot, S. R.; Eggebrecht, J. (1999). “Exact method for the simulation of Coulombic systems by spherically truncated, pairwise r-1 summation.” J. Chem. Phys. 110, 8254. DOI: 10.1063/1.478738

Batched Calculations#

All electrostatics functions support batched calculations for evaluating multiple independent systems simultaneously. For most use cases (except for very large systems) batching is the optimal way to amortize GPU utilization.

The API for electrostatics only needs minor modification to support batches of systems: users must provide a batch_idx tensor to both the initial neighbor list computation as well as to either the ewald_summation() and particle_mesh_ewald() methods. While \(\alpha\) can be specified independently for each system within a batch, the mesh dimensions must be the same for all systems (although each system has its own mesh grid).

Example code to perform a batched Ewald calculation:

PyTorch

import torch
from nvalchemiops.torch.interactions.electrostatics import ewald_summation
from nvalchemiops.torch.neighbors import neighbor_list

# Concatenate atoms from multiple systems
positions = torch.cat([pos_system0, pos_system1, pos_system2])
charges = torch.cat([charges_system0, charges_system1, charges_system2])

# Assign each atom to its system
batch_idx = torch.cat([
    torch.zeros(len(pos_system0), dtype=torch.int32),
    torch.ones(len(pos_system1), dtype=torch.int32),
    torch.full((len(pos_system2),), 2, dtype=torch.int32),
]).to(positions.device)

# Stack cells (B, 3, 3)
cells = torch.stack([cell0, cell1, cell2])
pbc = torch.tensor([[True, True, True]] * 3, device=positions.device)

# Build batched neighbor list
neighbor_list_coo, neighbor_ptr, neighbor_shifts = neighbor_list(
    positions, cutoff=10.0, cell=cells, pbc=pbc,
    batch_idx=batch_idx, method="batch_naive", return_neighbor_list=True
)

# Per-system alpha values (optional)
alphas = torch.tensor([0.3, 0.35, 0.3], dtype=torch.float64, device=positions.device)

# Batched calculation
energies, forces = ewald_summation(
    positions=positions,
    charges=charges,
    cell=cells,
    alpha=alphas,  # Per-system or single value
    k_cutoff=8.0,
    batch_idx=batch_idx,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

# energies: (total_atoms,) - per-atom energies
# Sum per system:
energy_per_system = torch.zeros(3, device=positions.device)
energy_per_system.scatter_add_(0, batch_idx.long(), energies)

Batch mode uses one shared set of Miller indices for the reciprocal-space calculation. If k_cutoff is supplied per system, either directly or via estimate_ewald_parameters, nvalchemiops uses the maximum cutoff across the batch to build that shared set.

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.electrostatics import ewald_summation
from nvalchemiops.jax.neighbors import neighbor_list

# Concatenate atoms from multiple systems
positions = jnp.concatenate([pos_system0, pos_system1, pos_system2])
charges = jnp.concatenate([charges_system0, charges_system1, charges_system2])

# Assign each atom to its system
batch_idx = jnp.concatenate([
    jnp.zeros(len(pos_system0), dtype=jnp.int32),
    jnp.ones(len(pos_system1), dtype=jnp.int32),
    jnp.full((len(pos_system2),), 2, dtype=jnp.int32),
])

# Stack cells (B, 3, 3)
cells = jnp.stack([cell0, cell1, cell2])
pbc = jnp.array([[True, True, True]] * 3)

# Build batched neighbor list
neighbor_list_coo, neighbor_ptr, neighbor_shifts = neighbor_list(
    positions, cutoff=10.0, cell=cells, pbc=pbc,
    batch_idx=batch_idx, method="batch_naive", return_neighbor_list=True
)

# Per-system alpha values (optional)
alphas = jnp.array([0.3, 0.35, 0.3], dtype=jnp.float64)

# Batched calculation
energies, forces = ewald_summation(
    positions=positions,
    charges=charges,
    cell=cells,
    alpha=alphas,  # Per-system or single value
    k_cutoff=8.0,
    batch_idx=batch_idx,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
    compute_forces=True,
)

# energies: (total_atoms,) - per-atom energies
# Sum per system using segment_sum:
energy_per_system = jax.ops.segment_sum(energies, batch_idx, num_segments=3)

Parameter Estimation#

ALCHEMI Toolkit-Ops provides functions to estimate sensible parameters based on desired accuracy threshold with two functions that share some functionality, but target the Ewald and PME algorithms respectively.

Ewald Parameters#

The function estimate_ewald_parameters() (PyTorch) / estimate_ewald_parameters() (JAX) is used to estimate \(\alpha\) and cutoffs for real- and reciprocal-space specifically for the Ewald algorithm:

PyTorch

from nvalchemiops.torch.interactions.electrostatics import estimate_ewald_parameters

params = estimate_ewald_parameters(
    positions=positions,
    cell=cell,
    batch_idx=None,  # or provide for batched systems
    accuracy=1e-6,
)

print(f"alpha = {params.alpha.item():.4f}")
print(f"r_cutoff = {params.real_space_cutoff.item():.4f}")
print(f"k_cutoff = {params.reciprocal_space_cutoff.item():.4f}")

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.electrostatics import estimate_ewald_parameters

params = estimate_ewald_parameters(
    positions=positions,
    cell=cell,
    batch_idx=None,  # or provide for batched systems
    accuracy=1e-6,
)

print(f"alpha = {params.alpha:.4f}")
print(f"r_cutoff = {params.real_space_cutoff:.4f}")
print(f"k_cutoff = {params.reciprocal_space_cutoff:.4f}")

This method returns an EwaldParameters dataclass, which is a light data structure that holds parameters used for the Ewald algorithm.

PME Parameters#

The function estimate_pme_parameters() (PyTorch) / estimate_pme_parameters() (JAX) is used to estimate \(\alpha\), the real-space cutoff, and mesh specifications specifically for the PME algorithm; the value of \(\alpha\) is determined the same way as for Ewald.

PyTorch

from nvalchemiops.torch.interactions.electrostatics import estimate_pme_parameters

params = estimate_pme_parameters(
    positions=positions,
    cell=cell,
    batch_idx=None,
    accuracy=1e-6,
)

print(f"alpha = {params.alpha.item():.4f}")
print(f"Mesh: {params.mesh_dimensions}")
print(f"r_cutoff = {params.real_space_cutoff.item():.4f}")

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.electrostatics import estimate_pme_parameters

params = estimate_pme_parameters(
    positions=positions,
    cell=cell,
    batch_idx=None,
    accuracy=1e-6,
)

print(f"alpha = {params.alpha:.4f}")
print(f"Mesh: {params.mesh_dimensions}")
print(f"r_cutoff = {params.real_space_cutoff:.4f}")

This method returns a PMEParameters dataclass, which is a light data structure that holds parameters used for the particle-mesh Ewald algorithm.

Units#

The electrostatics functions are unit-agnostic; they work in whatever consistent unit system you provide. Common conventions:

Unit System	Positions	Energy	Charge
Atomic units	Bohr	Hartree	e
eV-Angstrom	Angstrom	eV	e
LAMMPS “real”	Angstrom	kcal/mol	e

Important

Ensure consistency between your position units, cell units, and cutoff values. The alpha parameter has units of inverse length.

For atomic units (Bohr/Hartree), no additional constants are needed. For other unit systems, you may need to multiply energies by a Coulomb constant:

# eV-Angstrom: k_e ~ 14.3996 eV*Angstrom
# The functions assume k_e = 1 (atomic units)

Theory Background#

The Ewald Splitting#

The Coulomb potential \(1/r\) is split into short-range and long-range components using a Gaussian screening function:

\[\frac{1}{r} = \frac{\text{erfc}(\alpha r)}{r} + \frac{\text{erf}(\alpha r)}{r}\]

The \(\text{erfc}\) term decays exponentially and is computed in real space
The \(\text{erf}\) term is smooth and computed efficiently in reciprocal space

The splitting parameter \(\alpha\) controls the balance:

Large \(\alpha\): More work in reciprocal space, fewer k-vectors needed
Small \(\alpha\): More work in real space, larger neighbor cutoff needed

Charge Neutrality#

For periodic systems, overall charge neutrality is required for the electrostatic energy to be well-defined. Non-neutral systems include a background correction:

\[E_{\text{background}} = \frac{\pi}{2\alpha^2 V} Q_{\text{total}}^2\]

This term represents the interaction of the charged system with a uniform neutralizing background.

B-Spline Interpolation (PME)#

PME uses cardinal B-splines of order \(p\) for charge assignment:

Order 1: Nearest-grid-point (NGP)
Order 2: Cloud-in-cell (CIC)
Order 3: Triangular-shaped cloud (TSC)
Order 4: Cubic B-spline (recommended)
Higher orders: Smoother but wider support

Higher spline orders provide better accuracy but spread charges over more grid points. Order 4 (cubic) is the standard choice, balancing accuracy and efficiency.

Troubleshooting#

Common Issues#

Energy not converging with k_cutoff: The reciprocal-space energy should converge as k_cutoff increases. If it doesn’t, check that your cell is properly defined (lattice vectors as rows) and that the volume is computed correctly.

Force discontinuities: Ensure the real-space cutoff is compatible with your neighbor list cutoff. The neighbor list should include all pairs within the damping range of \(\text{erfc}(\alpha r)\).

NaN or Inf values:

Check for overlapping atoms (r -> 0)
Verify cell volume is positive
Ensure charges are finite

Memory issues with large meshes: PME mesh memory scales as \(n_x \times n_y \times n_z\). For very large cells, consider using coarser mesh spacing. It may also be worth comparing compute requirements between Ewald and PME algorithms.

Validation#

Note

The validation example below uses torchpme, which is a PyTorch-specific package. JAX users can validate against reference implementations in their ecosystem or compare against the PyTorch results for equivalent inputs.

You can validate PME results against reference implementations like torchpme. Here’s a simple example comparing reciprocal-space energies:

import torch
import math
from nvalchemiops.torch.interactions.electrostatics import pme_reciprocal_space

# Create a simple dipole system
device = torch.device("cuda")
dtype = torch.float64
cell_size = 10.0
separation = 2.0

# Two charges separated along x-axis
center = cell_size / 2
positions = torch.tensor(
    [
        [center - separation / 2, center, center],
        [center + separation / 2, center, center],
    ],
    dtype=dtype,
    device=device,
)
charges = torch.tensor([1.0, -1.0], dtype=dtype, device=device)
cell = torch.eye(3, dtype=dtype, device=device) * cell_size

# PME parameters
alpha = 0.3
mesh_spacing = 0.5
mesh_dims = (20, 20, 20)

# Compute reciprocal-space energy
energy = pme_reciprocal_space(
    positions=positions,
    charges=charges,
    cell=cell,
    alpha=alpha,
    mesh_dimensions=mesh_dims,
    spline_order=4,
    compute_forces=False,
)

print(f"Reciprocal-space energy: {energy.sum().item():.6f}")

# Optional: Compare with torchpme if available
try:
    from torchpme import PMECalculator
    from torchpme.potentials import CoulombPotential

    # torchpme uses sigma where Gaussian is exp(-r**2/(2 * sigma**2))
    # Standard Ewald uses exp(-alpha**2 * r**2), so sigma = 1/(2**0.5 * alpha)
    smearing = 1.0 / (math.sqrt(2.0) * alpha)
    potential = CoulombPotential(smearing=smearing).to(device=device, dtype=dtype)

    calculator = PMECalculator(
        potential=potential,
        mesh_spacing=mesh_spacing,
        interpolation_nodes=4,
        full_neighbor_list=True,
        prefactor=1.0,
    ).to(device=device, dtype=dtype)

    charges_pme = charges.unsqueeze(1)
    reciprocal_potential = calculator._compute_kspace(charges_pme, cell, positions)
    torchpme_energy = (reciprocal_potential * charges_pme).sum()

    print(f"TorchPME energy: {torchpme_energy.item():.6f}")
    print(f"Relative difference: {abs(energy.sum() - torchpme_energy) / abs(torchpme_energy):.2e}")

except ImportError:
    print("torchpme not available for comparison")

For more comprehensive validation examples, including:

Crystal structure systems (CsCl, wurtzite, zincblende)
Gradient validation against numerical finite differences
Batch processing consistency checks
Conservation law tests (momentum, translation invariance)

See the unit tests at test/interactions/electrostatics/ in the repository.

Electrostatic Interactions#

Overview of Available Methods#

Quick Start#

Data Formats#

Tensor Specifications#

Output Data Types#

Neighbor Representations#

Ewald Summation#

Mathematical Background#

Usage Examples#

Explicit Parameters#

Automatic Parameter Estimation#

Separating real- and reciprocal-space#

Particle Mesh Ewald (PME)#

Mathematical Background#

Usage Examples#

Basic Usage#

Mesh Spacing#

Automatic Parameter Estimation#

PME vs Ewald: When to Use Each#

Damped Shifted Force (DSF)#

Motivation#

Mathematical Background#

Shifted-Force Construction#

DSF Pair Potential#

DSF Force#

Self-Energy Correction#

Total System Energy#

Usage Examples#

Basic Energy and Forces#

With Periodic Boundary Conditions and Virial#

Using Neighbor Matrix Format#

Charge Gradients for MLIP Training#

Batched Calculations#

Undamped Shifted-Force Coulomb (alpha=0)#

Parameter Guidance#

When to Use DSF#

Applicability and Limitations#

Software Ecosystem#

References#

Batched Calculations#

Autograd Support#

Position Gradients (Forces)#

Charge Gradients#

Virial / Stress#

Parameter Estimation#

Ewald Parameters#

PME Parameters#

Units#

Theory Background#

The Ewald Splitting#

Charge Neutrality#

B-Spline Interpolation (PME)#

Troubleshooting#

Common Issues#

Validation#

Further Reading#