Ctrl+K

nvalchemiops.torch.interactions.electrostatics: Electrostatics#

The electrostatics module provides GPU-accelerated implementations of long-range electrostatic interactions for molecular simulations with PyTorch bindings. These functions accept standard torch.Tensor inputs and support automatic differentiation. Ewald and PME support full autograd for positions, charges, and cell parameters. DSF supports charge gradients via autograd; forces and virials are computed analytically. Setup parameters such as alpha, cutoffs, mesh controls, batch metadata, and neighbor topology are treated as constants. Cell-derived caches such as k_vectors, k_squared, volume, and cell_inv_t are accepted when cell.requires_grad is true, but they are static metadata and are assumed to correspond to the current cell; their cache-generation derivatives are not recovered. Energy-returning Ewald, PME, and slab paths support atom-weighted losses such as (weights * energies).sum() for positions, charges, and supported cell derivatives. Point-charge Ewald/PME inputs support float32 and float64. Keep all floating inputs and precomputed metadata in a call on a consistent dtype.

Tip

For the underlying framework-agnostic Warp kernels, see nvalchemiops.interactions.electrostatics: Electrostatic Interactions (Warp).

High-Level Interface#

These are the primary entry points for most users.

nvalchemiops.torch.interactions.electrostatics.ewald_summation(positions, charges, cell, alpha=None, k_vectors=None, k_cutoff=None, batch_idx=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, mask_value=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, accuracy=1e-6, hybrid_forces=False, pbc=None, slab_correction=False, *, miller_bounds=None, max_atoms_per_system=None)[source]#

Complete Ewald summation for long-range electrostatics.

Computes total Coulomb energy by combining real-space and reciprocal-space contributions with self-energy and background corrections.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • charges (torch.Tensor, shape (N,)) – Atomic partial charges.

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrices.

  • alpha (float, torch.Tensor, or None, default=None) – Ewald splitting parameter. Auto-estimated if None.

  • k_vectors (torch.Tensor, optional) – Pre-computed reciprocal lattice vectors.

  • k_cutoff (float, optional) – K-space cutoff for generating k_vectors.

  • miller_bounds (tuple[int, int, int] or torch.Tensor, optional, keyword-only) – Precomputed Miller-index half-bounds used when k_vectors is not supplied. Passing Python integer bounds avoids deriving range sizes from device tensors inside regenerated-k-vector loops.

  • max_atoms_per_system (int, optional, keyword-only) – Maximum number of atoms in any single system when batch_idx is provided. See ewald_reciprocal_space() for the sync-free launch contract.

  • batch_idx (torch.Tensor, shape (N,), optional) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

  • neighbor_list (torch.Tensor, shape (2, M), optional) – Neighbor pairs in COO format.

  • neighbor_ptr (torch.Tensor, shape (N+1,), optional) – CSR row pointers.

  • neighbor_shifts (torch.Tensor, shape (M, 3), optional) – Periodic image shifts for neighbor list.

  • neighbor_matrix (torch.Tensor, shape (N, max_neighbors), optional) – Dense neighbor matrix.

  • neighbor_matrix_shifts (torch.Tensor, shape (N, max_neighbors, 3), optional) – Periodic image shifts for neighbor_matrix.

  • mask_value (int, optional) – Value indicating invalid entries. Defaults to N.

  • compute_forces (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use torch.autograd.grad for differentiable forces.

  • compute_charge_gradients (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use torch.autograd.grad for \(\partial E / \partial q_i\).

  • compute_virial (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag for the virial tensor \(W = -\partial E / \partial \varepsilon\). Stress = -virial / volume.

  • accuracy (float, default=1e-6) – Target accuracy for parameter estimation.

  • hybrid_forces (bool, default=False) – When True, positions and cell are detached from the autograd graph and charge gradients are attached to the energy via a straight-through trick. Forces and virial are forward-only (not differentiable). See ewald_real_space() for details.

  • pbc (torch.Tensor, shape (3,) or (B, 3), dtype=bool, optional) – Per-system periodic boundary conditions. Required when slab_correction=True. Each row has True for periodic directions and False for the non-periodic (slab) direction. A (3,) tensor is accepted only for single-system calls; batched calls require explicit (B, 3) per-system pbc. This argument controls the slab correction geometry; real-space periodic images are determined by the neighbor list supplied to the Ewald real-space term.

  • slab_correction (bool, default=False) – When True, apply the Yeh-Berkowitz slab correction (with the Ballenegger et al. 2009 Eq. 29 non-neutral extension) to the total energy and to forces/charge_grads/virial when those are requested. Orthorhombic and triclinic slab cells are supported.

Returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom total Ewald energy.

  • forces (torch.Tensor, shape (N, 3), optional) –

    Deprecated since version 0.4.0: Deprecated direct forces (if compute_forces=True).

  • charge_gradients (torch.Tensor, shape (N,), optional) –

    Deprecated since version 0.4.0: Deprecated direct charge gradients (if compute_charge_gradients=True).

  • virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor (if compute_virial=True). Always last in the tuple.

Return type:

tuple[Tensor, …] | Tensor

Note

Energies are accumulated in float64 for numerical stability. Deprecated direct forces, charge gradients, and virials match the input dtype where the underlying component path returns typed outputs.

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115).

Enabled output flags are appended in order: energies, [forces], [charge_gradients], [virial]. A single output is returned unwrapped; multiple outputs are returned as a tuple.

Examples

Automatic parameter estimation (recommended for most cases):

>>> energies = ewald_summation(
...     positions, charges, cell,
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
...     accuracy=1e-6,
... )
>>> total_energy = energies.sum()

Explicit parameters with forces:

>>> energies, forces = ewald_summation(
...     positions, charges, cell,
...     alpha=0.3, k_cutoff=8.0,
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
...     compute_forces=True,
... )

Slab correction for two-dimensional periodic systems:

>>> pbc_slab = torch.tensor([[True, True, False]], device=positions.device)
>>> energies, forces = ewald_summation(
...     positions, charges, cell,
...     alpha=0.3, k_cutoff=8.0,
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
...     pbc=pbc_slab, slab_correction=True,
...     compute_forces=True,
... )
nvalchemiops.torch.interactions.electrostatics.particle_mesh_ewald(positions, charges, cell, alpha=None, mesh_spacing=None, mesh_dimensions=None, spline_order=4, batch_idx=None, k_vectors=None, k_squared=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, mask_value=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, accuracy=1e-6, hybrid_forces=False, pbc=None, slab_correction=False, *, cell_inv_t=None, volume=None, moduli_x=None, moduli_y=None, moduli_z=None)[source]#

Complete Particle Mesh Ewald (PME) calculation for long-range electrostatics.

Computes total Coulomb energy using the PME method, which achieves \(O(N \log N)\) scaling through FFT-based reciprocal space calculations. Combines: 1. Real-space contribution (short-range, erfc-damped) 2. Reciprocal-space contribution (long-range, FFT + B-spline interpolation) 3. Self-energy and background corrections

Total Energy Formula:

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

where:

\[\begin{split}\begin{aligned} E_{\text{real}} &= \frac{1}{2} \sum_{i \neq j} q_i q_j \frac{\operatorname{erfc}(\alpha r_{ij})}{r_{ij}} \\ E_{\text{reciprocal}} &= \text{FFT-based smooth long-range contribution} \\ E_{\text{self}} &= \sum_i \frac{\alpha}{\sqrt{\pi}} q_i^2 \\ E_{\text{background}} &= \frac{\pi}{2\alpha^2 V} Q_{\text{total}}^2 \end{aligned}\end{split}\]
Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates. Supports float32 or float64 dtype.

  • charges (torch.Tensor, shape (N,)) – Atomic partial charges in elementary charge units.

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrices with lattice vectors as rows. Shape (3, 3) is automatically promoted to (1, 3, 3) for single-system mode.

  • alpha (float, torch.Tensor, or None, default=None) – Ewald splitting parameter controlling real/reciprocal space balance. - float: Same \(\alpha\) for all systems - Tensor shape (B,): Per-system \(\alpha\) values - None: Automatically estimated using Kolafa-Perram formula Larger \(\alpha\) shifts more computation to reciprocal space.

  • mesh_spacing (float, optional) – Target mesh spacing in same units as cell (typically Å). Mesh dimensions computed as ceil(cell_length / mesh_spacing). Typical value: 0.8-1.2 Å. This setup path reads cell lengths into Python integers; pass explicit mesh_dimensions when cell-dependent mesh sizing is not desired.

  • mesh_dimensions (tuple[int, int, int], optional) – Explicit FFT mesh dimensions (nx, ny, nz). Power-of-2 values recommended for optimal FFT performance. If None and mesh_spacing is None, computed from accuracy parameter.

  • spline_order (int, default=4) – B-spline interpolation order. Higher orders are more accurate but slower. - 4: Cubic B-splines (standard, good accuracy/speed balance) - 5-6: Higher accuracy for demanding applications

  • batch_idx (torch.Tensor, shape (N,), dtype=int32, optional) – System index for each atom (0 to B-1). Determines execution mode: - None: Single-system optimized kernels - Provided: Batched kernels for multiple independent systems When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

  • k_vectors (torch.Tensor, shape (nx, ny, nz//2+1, 3), optional) – Precomputed k-vectors from generate_k_vectors_pme. Providing this along with k_squared skips k-vector generation (~15% speedup). Useful for fixed-cell MD simulations (NVT/NVE). When supplied while cell.requires_grad is true, the cache is assumed to correspond to the current cell.

  • k_squared (torch.Tensor, shape (nx, ny, nz//2+1), optional) – Precomputed \(|k|^2\) values. Must be provided together with k_vectors.

  • cell_inv_t (torch.Tensor, shape (3, 3) or (B, 3, 3), optional) – Precomputed transposed cell inverse \((M^{-1})^T\). When supplied, the reciprocal-space path skips the per-call torch.linalg.inv of the cell (which dispatches getrf/trsm/laswp on the 3x3 cell every iteration). This is a setup constant for fixed-cell calls and is assumed to correspond to the current cell when supplied while cell.requires_grad is true.

  • volume (torch.Tensor, shape (1,) or (B,), optional) – Precomputed cell volume \(|\det(M)|\). When supplied, both the Green’s-function normalization and the self/background correction skip torch.linalg.det (which also dispatches getrf under the hood). Same fixed-cell use-case as cell_inv_t.

  • moduli_x (torch.Tensor, optional) – Precomputed 1D B-spline modulus LUTs (sinc(m/N)^spline_order per axis) from compute_bspline_moduli_1d. When supplied, the reciprocal-space path skips the per-call fftfreq + sinc^p rebuild. The moduli only depend on mesh dimension + spline order, so callers can precompute them once for repeated calls with the same mesh and spline order.

  • moduli_y (torch.Tensor, optional) – Precomputed 1D B-spline modulus LUTs (sinc(m/N)^spline_order per axis) from compute_bspline_moduli_1d. When supplied, the reciprocal-space path skips the per-call fftfreq + sinc^p rebuild. The moduli only depend on mesh dimension + spline order, so callers can precompute them once for repeated calls with the same mesh and spline order.

  • moduli_z (torch.Tensor, optional) – Precomputed 1D B-spline modulus LUTs (sinc(m/N)^spline_order per axis) from compute_bspline_moduli_1d. When supplied, the reciprocal-space path skips the per-call fftfreq + sinc^p rebuild. The moduli only depend on mesh dimension + spline order, so callers can precompute them once for repeated calls with the same mesh and spline order.

  • neighbor_list (torch.Tensor, shape (2, M), dtype=int32, optional) – Neighbor pairs for real-space in COO format. Row 0 = source indices, row 1 = target indices. Mutually exclusive with neighbor_matrix.

  • neighbor_ptr (torch.Tensor, shape (N+1,), dtype=int32, optional) – CSR row pointers for neighbor_list. neighbor_ptr[i] gives the starting index in neighbor_list for atom i’s neighbors. Required with neighbor_list.

  • neighbor_shifts (torch.Tensor, shape (M, 3), dtype=int32, optional) – Periodic image shifts for neighbor_list. Required with neighbor_list.

  • neighbor_matrix (torch.Tensor, shape (N, max_neighbors), dtype=int32, optional) – Dense neighbor matrix format. Entry [i, k] = j means j is k-th neighbor of i. Invalid entries should be set to mask_value. Mutually exclusive with neighbor_list.

  • neighbor_matrix_shifts (torch.Tensor, shape (N, max_neighbors, 3), dtype=int32, optional) – Periodic image shifts for neighbor_matrix. Required with neighbor_matrix.

  • mask_value (int, optional) – Value indicating invalid entries in neighbor_matrix. Defaults to N.

  • compute_forces (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use torch.autograd.grad for differentiable forces.

  • compute_charge_gradients (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use torch.autograd.grad for \(\partial E/\partial q_i\).

  • compute_virial (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag for the virial tensor W = -dE/d(displacement). Stress = -virial / volume.

  • accuracy (float, default=1e-6) – Target relative accuracy for automatic parameter estimation (\(\alpha\), mesh dims). Only used when alpha or mesh_dimensions is None. Smaller values increase accuracy but also computational cost.

  • hybrid_forces (bool, default=False) – When True, positions and cell are detached from the autograd graph and charge gradients are attached to the energy via a straight-through trick. Forces and virial are forward-only (not differentiable). See ewald_real_space() for details.

  • pbc (torch.Tensor, shape (3,) or (B, 3), optional) – Per-system periodic boundary conditions for slab correction. Required when slab_correction=True. Each row has True for periodic directions and False for the non-periodic slab direction. Batched slab correction requires explicit shape (B, 3).

  • slab_correction (bool, default=False) – Whether to add the two-dimensional Yeh-Berkowitz / Ballenegger slab correction to the 3D-periodic PME result. This is only available for the full PME interface; use compute_slab_correction() explicitly when manually composing ewald_real_space and pme_reciprocal_space.

Returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom contribution to total PME energy. Sum gives total energy.

  • forces (torch.Tensor, shape (N, 3), optional) –

    Deprecated since version 0.4.0: Deprecated direct forces. Only returned if compute_forces=True.

  • charge_gradients (torch.Tensor, shape (N,), optional) –

    Deprecated since version 0.4.0: Deprecated direct charge gradients \(\partial E/\partial q_i\). Only returned if compute_charge_gradients=True.

  • virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor. Only returned if compute_virial=True. Always last in tuple.

Return type:

Tensor | tuple[Tensor, …]

Note

Internal reductions use float64 where needed for numerical stability. Returned energies, forces, and virials match the input dtype. Energy gradients are part of the public contract only for positions, charges, and cell. Caller-supplied reciprocal metadata such as k_vectors, k_squared, volume, and cell_inv_t is treated as static setup state that corresponds to the current cell.

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115).

Enabled output flags are appended in order: energies, [forces], [charge_gradients], [virial]. A single output is returned unwrapped; multiple outputs are returned as a tuple.

Raises:

  • ValueError – If neither neighbor_list nor neighbor_matrix is provided for real-space.

  • TypeError – If alpha has an unsupported type.

Parameters:
Return type:

Tensor | tuple[Tensor, …]

Examples

Automatic parameter estimation (recommended for most cases):

>>> energies = particle_mesh_ewald(
...     positions, charges, cell,
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
...     accuracy=1e-6,
... )
>>> total_energy = energies.sum()

Explicit parameters for reproducibility:

>>> energies = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     spline_order=4,
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... )
>>> forces = -torch.autograd.grad(energies.sum(), positions, create_graph=True)[0]

Using mesh spacing for automatic mesh sizing:

>>> energies = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_spacing=1.0,  # ~1 Å spacing
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... )

Batched systems (multiple independent structures):

>>> # positions: concatenated atoms from all systems
>>> # batch_idx: [0,0,0,0, 1,1,1,1, 2,2,2,2] for 4 atoms x 3 systems
>>> energies = particle_mesh_ewald(
...     positions, charges, cells,  # cells shape (3, 3, 3)
...     alpha=torch.tensor([0.3, 0.35, 0.3]),
...     batch_idx=batch_idx,
...     mesh_dimensions=(32, 32, 32),
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... )

Precomputed k-vectors for MD loop (fixed cell):

>>> from nvalchemiops.torch.interactions.electrostatics import generate_k_vectors_pme
>>> mesh_dims = (32, 32, 32)
>>> k_vectors, k_squared = generate_k_vectors_pme(cell, mesh_dims)
>>> for step in range(num_steps):
...     energies = particle_mesh_ewald(
...         positions, charges, cell,
...         alpha=0.3, mesh_dimensions=mesh_dims,
...         k_vectors=k_vectors, k_squared=k_squared,
...         neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
...     )

With charge gradients for ML training:

>>> charges.requires_grad_(True)
>>> energies = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... )
>>> charge_grads = torch.autograd.grad(energies.sum(), charges, create_graph=True)[0]

PME with slab correction:

>>> pbc_slab = torch.tensor([[True, True, False]], device=positions.device)
>>> energies, forces = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
...     compute_forces=True,
...     pbc=pbc_slab,
...     slab_correction=True,
... )

Using PyTorch autograd:

>>> positions.requires_grad_(True)
>>> energies = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... )
>>> total_energy = energies.sum()
>>> total_energy.backward()
>>> autograd_forces = -positions.grad  # Should match explicit forces

Notes

Automatic Parameter Estimation (when alpha is None):

Uses Kolafa-Perram formula:

\[\begin{split}\begin{aligned} \eta &= \frac{(V^2 / N)^{1/6}}{\sqrt{2\pi}} \\ \alpha &= \frac{1}{2\eta} \end{aligned}\end{split}\]

Mesh dimensions (when mesh_dimensions is None):

\[n_x = \left\lceil \frac{2 \alpha L_x}{3 \varepsilon^{1/5}} \right\rceil\]
Autograd Support:

All inputs (positions, charges, cell) support gradient computation.

See also

pme_reciprocal_space

Reciprocal-space component only

ewald_real_space

Real-space component (used internally)

estimate_pme_parameters

Automatic parameter estimation

PMEParameters

Container for PME parameters

Slab Correction#

Two-dimensional slab correction for systems with two periodic axes and one non-periodic axis. The high-level Ewald and PME interfaces can add this correction directly. Component-level workflows should add compute_slab_correction explicitly to ewald_real_space plus either ewald_reciprocal_space for Ewald or pme_reciprocal_space for PME.

nvalchemiops.torch.interactions.electrostatics.compute_slab_correction(positions, charges, cell, pbc, batch_idx=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False)[source]#

Yeh-Berkowitz slab correction for 2D periodic electrostatics, with the Ballenegger et al. (2009) Eq. 29 extension for non-neutral systems.

Returns the slab-correction contribution (per-atom energy and optionally per-atom force, charge gradient, and per-system virial). The caller adds these to the corresponding 3D Ewald/PME quantities; in normal usage the correction is invoked through ewald_summation(..., slab_correction=True) or particle_mesh_ewald(..., slab_correction=True).

Background-charge convention#

For systems with net charge \(Q \ne 0\), the formula corresponds to a uniform-volume neutralizing background (the same convention used by standard 3D Ewald). This matches LAMMPS and Ballenegger et al. (2009) Eq. 29.

Cell geometry#

Orthorhombic and triclinic cells are supported. For triclinic slab systems, the slab normal follows the plane spanned by the two periodic cell vectors.

param positions:

Atomic coordinates.

type positions:

torch.Tensor, shape (N, 3)

param charges:

Atomic charges.

type charges:

torch.Tensor, shape (N,)

param cell:

Unit cell matrices.

type cell:

torch.Tensor, shape (3, 3) or (B, 3, 3)

param pbc:

Per-system periodic boundary conditions. True for periodic directions, False for the non-periodic (slab) direction. Systems whose pbc is not slab-like (i.e., has anything other than exactly one False entry) contribute zero. A (3,) tensor is accepted only for single-system calls; batched calls require explicit (B, 3) per-system pbc.

type pbc:

torch.Tensor, shape (3,) or (B, 3), dtype=bool

param batch_idx:

System index for each atom. Defaults to all zeros (single system). When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

type batch_idx:

torch.Tensor, shape (N,), dtype=int32, optional

param compute_forces:

If True, return per-atom forces.

type compute_forces:

bool, default=False

param compute_charge_gradients:

If True, return per-atom charge gradients dE_slab/dq_i.

type compute_charge_gradients:

bool, default=False

param compute_virial:

If True, return per-system virial tensor using the normal-following affine strain convention W = E_slab * (I - 2 n n^T).

type compute_virial:

bool, default=False

returns:

  • energies (torch.Tensor, shape (N,), dtype=float64) – Per-atom slab correction energy.

  • forces (torch.Tensor, shape (N, 3), dtype matches positions, optional) – Per-atom slab force (only returned if compute_forces=True).

  • charge_grads (torch.Tensor, shape (N,), dtype=float64, optional) – Per-atom slab charge gradient (only if compute_charge_gradients=True).

  • virial (torch.Tensor, shape (B, 3, 3), dtype matches positions, optional) – Per-system slab virial (only if compute_virial=True).

Examples

Standalone correction for an orthorhombic slab with vacuum along z:

>>> pbc_slab = torch.tensor([[True, True, False]], device=positions.device)
>>> slab_energy, slab_forces = compute_slab_correction(
...     positions, charges, cell, pbc_slab, compute_forces=True
... )
>>> corrected_energy = ewald_energy + slab_energy

Triclinic cells use the normal to the periodic plane:

>>> triclinic_energy, triclinic_forces = compute_slab_correction(
...     positions, charges, triclinic_cell, pbc_slab, compute_forces=True
... )
Parameters:
Return type:

Tensor | tuple[Tensor, …]

Coulomb Interactions#

Direct pairwise Coulomb interactions.

nvalchemiops.torch.interactions.electrostatics.coulomb_energy(positions, charges, cell, cutoff, alpha=0.0, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, fill_value=None, batch_idx=None)[source]#

Compute Coulomb electrostatic energies.

Computes pairwise electrostatic energies using the Coulomb law, with optional erfc damping for Ewald/PME real-space calculations. Supports automatic differentiation with respect to positions, charges, and cell.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • charges (torch.Tensor, shape (N,)) – Atomic charges.

  • cell (torch.Tensor, shape (1, 3, 3) or (B, 3, 3)) – Unit cell matrix. Shape (B, 3, 3) for batched calculations.

  • cutoff (float) – Cutoff distance for interactions.

  • alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.

  • neighbor_list (torch.Tensor | None, shape (2, num_pairs)) – Neighbor pairs in COO format. Row 0 = source, Row 1 = target.

  • neighbor_ptr (torch.Tensor | None, shape (N+1,)) – CSR row pointers for neighbor list. Required with neighbor_list. Provided by neighborlist module.

  • neighbor_shifts (torch.Tensor | None, shape (num_pairs, 3)) – Integer unit cell shifts for neighbor list format.

  • neighbor_matrix (torch.Tensor | None, shape (N, max_neighbors)) – Neighbor indices in matrix format.

  • neighbor_matrix_shifts (torch.Tensor | None, shape (N, max_neighbors, 3)) – Integer unit cell shifts for matrix format.

  • fill_value (int | None) – Fill value for neighbor matrix padding.

  • batch_idx (torch.Tensor | None, shape (N,)) – Batch indices for each atom.

Returns:

energies – Per-atom energies. Sum to get total energy.

Return type:

torch.Tensor, shape (N,)

Examples

>>> # Direct Coulomb (undamped)
>>> energies = coulomb_energy(
...     positions, charges, cell, cutoff=10.0, alpha=0.0,
...     neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr,
...     neighbor_shifts=neighbor_shifts
... )
>>> total_energy = energies.sum()

>>> # Ewald/PME real-space (damped) with autograd
>>> positions.requires_grad_(True)
>>> energies = coulomb_energy(
...     positions, charges, cell, cutoff=10.0, alpha=0.3,
...     neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr,
...     neighbor_shifts=neighbor_shifts
... )
>>> energies.sum().backward()
>>> forces = -positions.grad
nvalchemiops.torch.interactions.electrostatics.coulomb_forces(positions, charges, cell, cutoff, alpha=0.0, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, fill_value=None, batch_idx=None)[source]#

Compute Coulomb electrostatic forces.

Convenience wrapper that returns only forces (no energies).

Parameters:

  • descriptions. (See coulomb_energy for parameter)

  • positions (Tensor)

  • charges (Tensor)

  • cell (Tensor)

  • cutoff (float)

  • alpha (float)

  • neighbor_list (Tensor | None)

  • neighbor_ptr (Tensor | None)

  • neighbor_shifts (Tensor | None)

  • neighbor_matrix (Tensor | None)

  • neighbor_matrix_shifts (Tensor | None)

  • fill_value (int | None)

  • batch_idx (Tensor | None)

Returns:

forces – Forces on each atom.

Return type:

torch.Tensor, shape (N, 3)

See also

coulomb_energy_forces

Compute both energies and forces

nvalchemiops.torch.interactions.electrostatics.coulomb_energy_forces(positions, charges, cell, cutoff, alpha=0.0, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, fill_value=None, batch_idx=None)[source]#

Compute Coulomb electrostatic energies and forces.

Computes pairwise electrostatic energies and forces using the Coulomb law, with optional erfc damping for Ewald/PME real-space calculations. Supports automatic differentiation with respect to positions, charges, and cell.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • charges (torch.Tensor, shape (N,)) – Atomic charges.

  • cell (torch.Tensor, shape (1, 3, 3) or (B, 3, 3)) – Unit cell matrix. Shape (B, 3, 3) for batched calculations.

  • cutoff (float) – Cutoff distance for interactions.

  • alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.

  • neighbor_list (torch.Tensor | None, shape (2, num_pairs)) – Neighbor pairs in COO format.

  • neighbor_ptr (torch.Tensor | None, shape (N+1,)) – CSR row pointers for neighbor list. Required with neighbor_list. Provided by neighborlist module.

  • neighbor_shifts (torch.Tensor | None, shape (num_pairs, 3)) – Integer unit cell shifts for neighbor list format.

  • neighbor_matrix (torch.Tensor | None, shape (N, max_neighbors)) – Neighbor indices in matrix format.

  • neighbor_matrix_shifts (torch.Tensor | None, shape (N, max_neighbors, 3)) – Integer unit cell shifts for matrix format.

  • fill_value (int | None) – Fill value for neighbor matrix padding.

  • batch_idx (torch.Tensor | None, shape (N,)) – Batch indices for each atom.

Returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom energies.

  • forces (torch.Tensor, shape (N, 3)) – Forces on each atom.

Return type:

tuple[Tensor, Tensor]

Note

Energies are always float64 for numerical stability during accumulation. Forces match the input dtype (float32 or float64).

Examples

>>> # Direct Coulomb
>>> energies, forces = coulomb_energy_forces(
...     positions, charges, cell, cutoff=10.0, alpha=0.0,
...     neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr,
...     neighbor_shifts=neighbor_shifts
... )

>>> # Ewald/PME real-space
>>> energies, forces = coulomb_energy_forces(
...     positions, charges, cell, cutoff=10.0, alpha=0.3,
...     neighbor_matrix=neighbor_matrix, neighbor_matrix_shifts=neighbor_matrix_shifts,
...     fill_value=num_atoms
... )

DSF Coulomb#

Damped Shifted Force (DSF) pairwise electrostatics with \(\mathcal{O}(N)\) scaling.

nvalchemiops.torch.interactions.electrostatics.dsf_coulomb(positions, charges, cutoff, alpha=0.2, cell=None, batch_idx=None, neighbor_list=None, neighbor_ptr=None, unit_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, fill_value=None, compute_forces=True, compute_virial=False, num_systems=None, device=None)[source]#

Compute DSF electrostatic energy, forces, and virial.

The Damped Shifted Force (DSF) method is a pairwise O(N) electrostatic summation technique that ensures both potential energy and forces smoothly vanish at a defined cutoff radius.

Supports float32 and float64 input precision. Energy is always returned in float64. Forces, virial, and charge gradients match the input precision.

Parameters:

  • positions (torch.Tensor, shape (num_atoms, 3)) – Atomic coordinates (float32 or float64).

  • charges (torch.Tensor, shape (num_atoms,)) – Atomic charges (must match positions dtype). If requires_grad=True, charge gradients (dE/dq) will be propagated through autograd.

  • cutoff (float) – Cutoff radius beyond which interactions are zero.

  • alpha (float, default 0.2) – Damping parameter. Set to 0.0 for shifted-force bare Coulomb.

  • cell (torch.Tensor, shape (num_systems, 3, 3), optional) – Unit cell matrices for periodic boundary conditions.

  • batch_idx (torch.Tensor, shape (num_atoms,), dtype=int32, optional) – System index for each atom. If None, all atoms in one system.

  • neighbor_list (torch.Tensor, shape (2, num_pairs), dtype=int32, optional) – Neighbor list in COO format. Row 1 contains destination atoms.

  • neighbor_ptr (torch.Tensor, shape (num_atoms+1,), dtype=int32, optional) – CSR row pointers (required with neighbor_list).

  • unit_shifts (torch.Tensor, shape (num_pairs, 3), dtype=int32, optional) – Integer unit cell shifts for PBC (required with neighbor_list + cell).

  • neighbor_matrix (torch.Tensor, shape (num_atoms, max_neighbors), dtype=int32, optional) – Dense neighbor matrix format.

  • neighbor_matrix_shifts (torch.Tensor, shape (num_atoms, max_neighbors, 3), dtype=int32, optional) – Integer unit cell shifts for matrix format PBC.

  • fill_value (int, optional) – Padding indicator for neighbor_matrix. Defaults to num_atoms.

  • compute_forces (bool, default True) – Whether to compute forces.

  • compute_virial (bool, default False) – Whether to compute virial tensor (requires PBC and compute_forces).

  • num_systems (int, optional) – Number of systems. Inferred from batch_idx or cell if not given.

  • device (str, optional) – Warp device string. Inferred from positions if not given.

Returns:

  • energy (torch.Tensor, shape (num_systems,), dtype=float64) – Per-system electrostatic energy (always float64). If charges.requires_grad, this tensor is connected to the autograd graph for charge gradients.

  • forces (torch.Tensor, shape (num_atoms, 3), dtype matches input) – Per-atom forces. Only returned if compute_forces=True.

  • virial (torch.Tensor, shape (num_systems, 3, 3), dtype matches input) – Per-system virial tensor. Only returned if compute_virial=True.

Return type:

tuple[Tensor] | tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]

Notes

  • Assumes a full neighbor list (each pair appears in both directions).

  • For MLIP training with geometry-dependent charges, set charges.requires_grad_(True) before calling. After energy.sum().backward(), charges.grad will contain dE/dq.

  • Charge gradients (dE/dq) are computed when charges.requires_grad=True, regardless of compute_forces.

  • The returned energy tensor is not differentiable w.r.t. positions or cell through PyTorch autograd. Forces are computed analytically by the Warp kernel, not via autograd.

Examples

>>> # Basic energy + forces
>>> energy, forces = dsf_coulomb(positions, charges, cutoff=10.0, alpha=0.2,
...     neighbor_list=nl, neighbor_ptr=ptr)

>>> # MLIP workflow with charge gradients
>>> charges = model(positions)  # Predict charges from geometry
>>> charges.requires_grad_(True)
>>> energy, forces = dsf_coulomb(positions, charges, cutoff=10.0, alpha=0.2,
...     neighbor_list=nl, neighbor_ptr=ptr)
>>> loss = (energy - ref_energy).pow(2).sum()
>>> loss.backward()  # charges.grad now contains dE/dq * dloss/dE

Ewald Components#

Individual components of the Ewald summation method.

nvalchemiops.torch.interactions.electrostatics.ewald_real_space(positions, charges, cell, alpha, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, mask_value=None, batch_idx=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, hybrid_forces=False)[source]#

Compute real-space Ewald energy and optionally forces, charge gradients, and virial.

Computes the damped Coulomb interactions for atom pairs within the real-space cutoff. The complementary error function (erfc) damping ensures rapid convergence in real space.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • charges (torch.Tensor, shape (N,)) – Atomic partial charges.

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrices.

  • alpha (torch.Tensor, shape (1,) or (B,)) – Ewald splitting parameter(s).

  • neighbor_list (torch.Tensor, shape (2, M), optional) – Neighbor list in COO format.

  • neighbor_ptr (torch.Tensor, shape (N+1,), optional) – CSR row pointers for neighbor list.

  • neighbor_shifts (torch.Tensor, shape (M, 3), optional) – Periodic image shifts for neighbor list.

  • neighbor_matrix (torch.Tensor, shape (N, max_neighbors), optional) – Dense neighbor matrix format.

  • neighbor_matrix_shifts (torch.Tensor, shape (N, max_neighbors, 3), optional) – Periodic image shifts for neighbor_matrix.

  • mask_value (int, optional) – Value indicating invalid entries in neighbor_matrix. Defaults to N.

  • batch_idx (torch.Tensor, shape (N,), optional) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

  • compute_forces (bool, default=False) – Whether to compute explicit component forces. This direct output is kept for no-autograd MD/inference use; use energy autograd for differentiable training.

  • compute_charge_gradients (bool, default=False) – Whether to compute explicit component charge gradients. This direct output follows the same no-autograd contract as compute_forces.

  • compute_virial (bool, default=False) – Whether to compute the component virial tensor \(W = -\partial E / \partial \varepsilon\). Stress = -virial / volume.

  • hybrid_forces (bool, default=False) – When True, positions and cell are detached from the autograd graph and charge gradients are attached to the energy via a straight-through trick. Forces and virial are forward-only (not differentiable). This is intended for efficient inference with geometry-dependent charges \(q = q(R)\), where explicit forces provide \(\partial E/\partial R|_q\) and autograd through the energy provides the charge chain-rule term \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\).

Returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom real-space energy.

  • forces (torch.Tensor, shape (N, 3), optional) – Direct component forces (if compute_forces=True).

  • charge_gradients (torch.Tensor, shape (N,), optional) – Direct component charge gradients (if compute_charge_gradients=True).

  • virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor (if compute_virial=True). Always last in the tuple.

Return type:

Tensor | tuple[Tensor, …]

Note

Energies are always float64 for numerical stability during accumulation. Forces, virial, and charge gradients match the input dtype (float32 or float64).

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115). Hybrid direct-output mode uses the same cached fallback connector so weighted \(q = q(R)\) losses can recover a valid energy gradient when the forward energy was detached.

nvalchemiops.torch.interactions.electrostatics.ewald_reciprocal_space(positions, charges, cell, k_vectors, alpha, batch_idx=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, hybrid_forces=False, *, max_atoms_per_system=None)[source]#

Compute reciprocal-space Ewald energy and optionally forces, charge gradients, virial.

Computes the smooth long-range electrostatic contribution using structure factors in reciprocal space.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • charges (torch.Tensor, shape (N,)) – Atomic partial charges.

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrices.

  • k_vectors (torch.Tensor) – Reciprocal lattice vectors. Shape (K, 3) for single system, (B, K, 3) for batch.

  • alpha (torch.Tensor, shape (1,) or (B,)) – Ewald splitting parameter(s).

  • batch_idx (torch.Tensor, shape (N,), optional) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

  • compute_forces (bool, default=False) – Whether to compute explicit component forces. This direct output is kept for no-autograd MD/inference use; use energy autograd for differentiable training.

  • compute_charge_gradients (bool, default=False) – Whether to compute explicit component charge gradients. This direct output follows the same no-autograd contract as compute_forces.

  • compute_virial (bool, default=False) – Whether to compute the component virial tensor \(W = -\partial E / \partial \varepsilon\). Stress = -virial / volume.

  • hybrid_forces (bool, default=False) – When True, positions and cell are detached from the autograd graph and charge gradients are attached to the energy via a straight-through trick. Forces and virial are forward-only (not differentiable). See ewald_real_space() for details.

  • max_atoms_per_system (int, optional, keyword-only) – Maximum number of atoms in any single system when batch_idx is provided. Passing this host-known upper bound avoids CUDA host synchronization from launch-size inference in the reciprocal kernel. Overestimates are safe but may launch extra blocks. When omitted, the bound is inferred from atom_start / atom_end and may synchronize on CUDA.

Returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom reciprocal-space energy.

  • forces (torch.Tensor, shape (N, 3), optional) – Direct component forces (if compute_forces=True).

  • charge_gradients (torch.Tensor, shape (N,), optional) – Direct component charge gradients (if compute_charge_gradients=True).

  • virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor (if compute_virial=True). Always last in the tuple.

Return type:

Tensor | tuple[Tensor, …]

Note

Energies are always float64 for numerical stability during accumulation. Forces, virial, and charge gradients match the input dtype (float32 or float64). k_vectors are setup metadata. Caller-supplied vectors are treated as static values that correspond to the current cell.

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115).

PME Components#

Individual components of the Particle Mesh Ewald method.

nvalchemiops.torch.interactions.electrostatics.pme_reciprocal_space(positions, charges, cell, alpha, mesh_dimensions=None, mesh_spacing=None, spline_order=4, batch_idx=None, k_vectors=None, k_squared=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, hybrid_forces=False, *, cell_inv_t=None, volume=None, moduli_x=None, moduli_y=None, moduli_z=None)[source]#

Compute PME reciprocal-space energy and optionally forces and/or charge gradients.

Performs the FFT-based reciprocal-space calculation using the Particle Mesh Ewald algorithm. This achieves \(O(N \log N)\) scaling through:

  1. B-spline charge interpolation to mesh (spreading)

  2. FFT of charge mesh to reciprocal space

  3. Convolution with raw Green’s function and B-spline deconvolution

  4. Inverse FFT back to real space (potential mesh)

  5. B-spline interpolation of potential to atoms (gathering)

  6. Self-energy and background corrections

Formula#

The reciprocal-space energy is computed via the mesh potential:

\[\varphi_{\text{mesh}}(k) = \frac{G(k)}{C^2(k)} \rho_{\text{mesh}}(k)\]

where:

  • \(G(k) = (2\pi/(V k^2)) \times \exp(-k^2/(4\alpha^2))\) is the volume-normalized PME Green’s function used by this implementation

  • \(C^2(k)\) is the squared B-spline structure factor

  • \(\rho_{\text{mesh}}(k)\) is the FFT of interpolated charges

param positions:

Atomic coordinates. Supports float32 or float64 dtype.

type positions:

torch.Tensor, shape (N, 3)

param charges:

Atomic partial charges in elementary charge units.

type charges:

torch.Tensor, shape (N,)

param cell:

Unit cell matrices with lattice vectors as rows. Shape (3, 3) is automatically promoted to (1, 3, 3).

type cell:

torch.Tensor, shape (3, 3) or (B, 3, 3)

param alpha:

Ewald splitting parameter controlling real/reciprocal space balance. - float: Same \(\alpha\) for all systems - Tensor shape (B,): Per-system \(\alpha\) values

type alpha:

float or torch.Tensor

param mesh_dimensions:

Explicit FFT mesh dimensions (nx, ny, nz). Power-of-2 values are optimal for FFT performance. Either mesh_dimensions or mesh_spacing must be provided.

type mesh_dimensions:

tuple[int, int, int], optional

param mesh_spacing:

Target mesh spacing in same units as cell. Mesh dimensions computed as ceil(cell_length / mesh_spacing). Typical value: ~1 Å. This setup path reads cell lengths into Python integers; pass explicit mesh_dimensions when cell-dependent mesh sizing is not desired.

type mesh_spacing:

float, optional

param spline_order:

B-spline interpolation order. Higher orders are more accurate but slower. - 4: Cubic B-splines (good balance, most common) - 5-6: Higher accuracy for demanding applications - Must be >= 3 for smooth interpolation

type spline_order:

int, default=4

param batch_idx:

System index for each atom (0 to B-1). Determines kernel dispatch: - None: Single-system optimized kernels - Provided: Batched kernels for multiple independent systems When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

type batch_idx:

torch.Tensor, shape (N,), dtype=int32, optional

param k_vectors:

Precomputed k-vectors from generate_k_vectors_pme. Providing this along with k_squared skips k-vector generation (~15% speedup). Can be precomputed once and reused when cell and mesh are unchanged. When supplied while cell.requires_grad is true, the cache is assumed to correspond to the current cell.

type k_vectors:

torch.Tensor, shape (nx, ny, nz//2+1, 3), optional

param k_squared:

Precomputed \(|k|^2\) values. Must be provided together with k_vectors. PME metadata tensors are setup constants and are detached from public autograd outputs.

type k_squared:

torch.Tensor, shape (nx, ny, nz//2+1), optional

param compute_forces:

Whether to compute explicit component reciprocal-space forces. This direct output is kept for no-autograd MD/inference use; use energy autograd for differentiable training.

type compute_forces:

bool, default=False

param compute_charge_gradients:

Whether to compute explicit component charge gradients \(\partial E/\partial q_i\). This direct output follows the same no-autograd contract as compute_forces.

type compute_charge_gradients:

bool, default=False

param compute_virial:

Whether to compute the component virial tensor W = -dE/d(displacement) for the row-vector displacement recipe. Stress = -virial / volume.

type compute_virial:

bool, default=False

param hybrid_forces:

When True, positions and cell are detached from the autograd graph and charge gradients are attached to the energy via a straight-through trick. Forces and virial are forward-only (not differentiable). See ewald_real_space() for details.

type hybrid_forces:

bool, default=False

returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom reciprocal-space energy (includes self and background corrections).

  • forces (torch.Tensor, shape (N, 3), optional) – Direct reciprocal-space forces. Only returned if compute_forces=True.

  • charge_gradients (torch.Tensor, shape (N,), optional) – Direct charge gradients \(\partial E_{\text{recip}}/\partial q_i\). Only returned if compute_charge_gradients=True.

  • virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor. Only returned if compute_virial=True. Always last in tuple.

Note

Internal reductions use float64 where needed for numerical stability. Returned energies, forces, and virials match the input dtype. Energy gradients are part of the public contract only for positions, charges, and cell. Caller-supplied reciprocal metadata such as k_vectors, k_squared, volume, and cell_inv_t is treated as static setup state that corresponds to the current cell.

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115).

torch.compile is supported by the public wrapper tests, although custom Warp operators and FFTs can still limit compiler fusion for PME workloads.

Enabled output flags are appended in order: energies, [forces], [charge_gradients], [virial]. A single output is returned unwrapped; multiple outputs are returned as a tuple.

raises ValueError:

If neither mesh_dimensions nor mesh_spacing is provided.

Examples

Energy only with explicit mesh dimensions:

>>> energies = pme_reciprocal_space(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
... )
>>> total_recip_energy = energies.sum()

With forces using mesh spacing:

>>> energies, forces = pme_reciprocal_space(
...     positions, charges, cell,
...     alpha=0.3, mesh_spacing=1.0,
...     compute_forces=True,
... )

Precomputed k-vectors for MD loop (fixed cell):

>>> from nvalchemiops.torch.interactions.electrostatics import generate_k_vectors_pme
>>> mesh_dims = (32, 32, 32)
>>> k_vectors, k_squared = generate_k_vectors_pme(cell, mesh_dims)
>>> for step in range(num_steps):
...     energies = pme_reciprocal_space(
...         positions, charges, cell,
...         alpha=0.3, mesh_dimensions=mesh_dims,
...         k_vectors=k_vectors, k_squared=k_squared,
...     )

With charge gradients for ML training:

>>> energies, charge_grads = pme_reciprocal_space(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     compute_charge_gradients=True,
... )

See also

particle_mesh_ewald

Complete PME calculation (real + reciprocal).

generate_k_vectors_pme

Generate k-vectors for this function.

Parameters:
Return type:

Tensor | tuple[Tensor, …]

nvalchemiops.torch.interactions.electrostatics.compute_bspline_moduli_1d(miller_indices, mesh_N, spline_order)[source]#

Precompute the 1D B-spline modulus LUT for one PME mesh axis.

Returns b[i] = sinc(m_i / N)^spline_order for each miller index m_i (with sinc(x) = sin(pi*x)/(pi*x), sinc(0) = 1). The three-axis product b_x[i] * b_y[j] * b_z[k] is the B-spline structure factor consumed by the factory-backed convolve kernel after a 1e-10 clamp + square. Precomputing the LUT lets the convolve kernel replace three sinc transcendentals + an order-dependent power loop per (i, j, k) thread with three reads + two multiplies.

Parameters:

  • miller_indices (torch.Tensor, shape (M,)) – Integer Miller indices along one mesh axis, typically produced by torch.fft.fftfreq or torch.fft.rfftfreq scaled by mesh_N.

  • mesh_N (int) – Number of mesh points along this axis.

  • spline_order (int) – B-spline interpolation order p. The modulus is \(\operatorname{sinc}(m/N)^p\).

Returns:

1D B-spline modulus values, one per Miller index. Same dtype as miller_indices.

Return type:

torch.Tensor, shape (M,)

K-Vector Generation#

nvalchemiops.torch.interactions.electrostatics.generate_k_vectors_ewald_summation(cell, k_cutoff, miller_bounds=None)[source]#

Generate reciprocal lattice vectors for Ewald summation (half-space).

Creates k-vectors within the specified cutoff for the reciprocal space summation in the Ewald method. Uses half-space optimization to reduce computational cost by approximately 2x.

Half-Space Optimization#

This function generates k-vectors in the positive half-space only, exploiting the symmetry S(-k) = S*(k) where S(k) is the structure factor. For each pair of k-vectors (k, -k), only one is included.

The half-space condition selects k-vectors where:

  • h > 0, OR

  • (h == 0 AND k > 0), OR

  • (h == 0 AND k == 0 AND l > 0)

The kernels in ewald_kernels.py compensate by doubling the Green’s function (using \(8\pi\) instead of \(4\pi\)), so energies, forces, and charge gradients are computed correctly.

Mathematical Background#

For a direct lattice defined by basis vectors {a, b, c} (rows of cell matrix), the reciprocal lattice vectors are:

\[ \begin{align}\begin{aligned}\mathbf{a}^* &= \frac{2\pi (\mathbf{b} \times \mathbf{c})}{V}\\\mathbf{b}^* &= \frac{2\pi (\mathbf{c} \times \mathbf{a})}{V}\\\mathbf{c}^* &= \frac{2\pi (\mathbf{a} \times \mathbf{b})}{V}\end{aligned}\end{align} \]

where \(V = \mathbf{a} \cdot (\mathbf{b} \times \mathbf{c})\) is the cell volume.

In matrix form: \(\text{reciprocal_matrix} = 2\pi \cdot (\text{cell}^T)^{-1}\)

Each k-vector is: \(\mathbf{k} = h \mathbf{a}^* + k \mathbf{b}^* + l \mathbf{c}^*\) where (h, k, l) are Miller indices (integers).

param cell:

Unit cell matrix with lattice vectors as rows. Shape (3, 3) for single system or (B, 3, 3) for batch.

type cell:

torch.Tensor

param k_cutoff:

Maximum magnitude of k-vectors to include (\(|\mathbf{k}| \leq k_{\text{cutoff}}\)). Typical values: 8-12 \(\text{\AA}^{-1}\) for molecular systems. Higher values increase accuracy but also computational cost.

type k_cutoff:

float or torch.Tensor

param miller_bounds:

Precomputed Miller-index half-bounds as returned by _generate_miller_indices(). Supplying Python integer bounds skips the device-to-host synchronization needed to derive FFT range sizes from cell and k_cutoff inside tight regenerated-k-vector loops. Tensor bounds are accepted for convenience but are read back to Python integers during setup.

type miller_bounds:

tuple[int, int, int] or torch.Tensor, optional

returns:

Reciprocal lattice vectors within the cutoff. Shape (K, 3) for single system or (B, K, 3) for batch. Excludes k=0 and includes only half-space vectors.

rtype:

torch.Tensor

Examples

Single system with explicit k_cutoff:

>>> cell = torch.eye(3, dtype=torch.float64) * 10.0
>>> k_vectors = generate_k_vectors_ewald_summation(cell, k_cutoff=8.0)
>>> k_vectors.shape
torch.Size([...])  # Number depends on cell size and cutoff

With automatic parameter estimation:

>>> from nvalchemiops.torch.interactions.electrostatics import estimate_ewald_parameters
>>> params = estimate_ewald_parameters(positions, cell)
>>> k_vectors = generate_k_vectors_ewald_summation(cell, params.reciprocal_space_cutoff)

Notes

  • The k=0 vector is always excluded (causes division by zero in Green’s function).

  • For batch mode, the same set of Miller indices is used for all systems but transformed using each system’s reciprocal cell. If k_cutoff is given per system, the maximum cutoff across the batch determines the shared Miller bounds.

  • The number of k-vectors K scales as O(k_cutoff^3 * V) where V is the cell volume.

See also

ewald_reciprocal_space

Uses these k-vectors for reciprocal space energy.

estimate_ewald_parameters

Automatic parameter estimation including k_cutoff.

Parameters:
Return type:

Tensor

nvalchemiops.torch.interactions.electrostatics.generate_k_vectors_pme(cell, mesh_dimensions, reciprocal_cell=None)[source]#

Generate reciprocal lattice vectors for Particle Mesh Ewald (PME).

Creates k-vectors on a regular grid compatible with FFT-based reciprocal space calculations in PME. Uses rfft conventions (half-size in z-dimension) to exploit Hermitian symmetry of real-valued charge densities.

Notes

For a direct lattice defined by basis vectors {a, b, c} (rows of cell matrix), the reciprocal lattice vectors are:

\[\begin{split}\begin{aligned} \mathbf{a}^* &= \frac{2\pi (\mathbf{b} \times \mathbf{c})}{V} \\ \mathbf{b}^* &= \frac{2\pi (\mathbf{c} \times \mathbf{a})}{V} \\ \mathbf{c}^* &= \frac{2\pi (\mathbf{a} \times \mathbf{b})}{V} \end{aligned}\end{split}\]

where \(V = \mathbf{a} \cdot (\mathbf{b} \times \mathbf{c})\) is the cell volume.

In matrix form:

\[\text{reciprocal_matrix} = 2\pi \cdot (\text{cell}^T)^{-1}\]

Each k-vector is then:

\[\mathbf{k} = h \mathbf{a}^* + k \mathbf{b}^* + l \mathbf{c}^*\]

where (h, k, l) are Miller indices (integers).

Parameters:

  • cell (torch.Tensor) – Unit cell matrix with lattice vectors as rows. Shape (3, 3) for single system or (B, 3, 3) for batch.

  • mesh_dimensions (tuple[int, int, int]) – PME mesh grid dimensions (nx, ny, nz). Should typically be chosen such that mesh spacing is \(\sim 1 \text{\AA}\) or finer. Power-of-2 dimensions are optimal for FFT performance.

  • reciprocal_cell (torch.Tensor, optional) – Precomputed reciprocal cell matrix (\(2\pi \cdot \text{cell}^{-1}\)). If provided, skips the inverse computation. Shape (3, 3) or (B, 3, 3).

Returns:

  • k_vectors (torch.Tensor, shape (nx, ny, nz//2+1, 3)) – Cartesian k-vectors at each grid point. Uses rfft convention where z-dimension is halved due to Hermitian symmetry.

  • k_squared_safe (torch.Tensor, shape (nx, ny, nz//2+1)) – Squared magnitude \(|\mathbf{k}|^2\) for each k-vector, with k=0 set to a small positive value (1e-12) to avoid division by zero.

Return type:

tuple[Tensor, Tensor]

Examples

Basic usage:

>>> cell = torch.eye(3, dtype=torch.float64) * 10.0
>>> mesh_dims = (32, 32, 32)
>>> k_vectors, k_squared = generate_k_vectors_pme(cell, mesh_dims)
>>> k_vectors.shape
torch.Size([32, 32, 17, 3])

With precomputed reciprocal cell:

>>> reciprocal_cell = 2 * torch.pi * torch.linalg.inv(cell)
>>> k_vectors, k_squared = generate_k_vectors_pme(
...     cell, mesh_dims, reciprocal_cell=reciprocal_cell
... )

Notes

  • The z-dimension output size is nz//2+1 due to rfft symmetry.

  • Miller indices follow torch.fft.fftfreq convention (0, 1, 2, …, -2, -1).

  • k_squared_safe has k=0 replaced with 1e-12 to prevent division by zero in Green’s function calculations.

See also

pme_reciprocal_space

Uses these k-vectors for PME reciprocal space energy.

pme_green_structure_factor

Computes Green’s function using k_squared.

Parameter Estimation#

Functions for automatic parameter estimation based on desired accuracy tolerance.

nvalchemiops.torch.interactions.electrostatics.estimate_ewald_parameters(positions, cell, batch_idx=None, accuracy=1e-6)[source]#

Estimate optimal Ewald summation parameters for a given accuracy.

Uses the Kolafa-Perram formula to balance real-space and reciprocal-space contributions for optimal efficiency at the target accuracy.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrix.

  • batch_idx (torch.Tensor, shape (N,), dtype=int32, optional) – System index for each atom. If None, single-system mode.

  • accuracy (float, default=1e-6) – Target accuracy (relative error tolerance).

Returns:

Dataclass containing alpha, real_space_cutoff, reciprocal_space_cutoff as torch.Tensor objects.

Return type:

EwaldParameters

nvalchemiops.torch.interactions.electrostatics.estimate_pme_parameters(positions, cell, batch_idx=None, accuracy=1e-6, real_space_cutoff=None, mesh_safety_factor=1.0)[source]#

Estimate PME parameters for a given accuracy.

Uses the closed-form Essmann/Kolafa-Perram derivation: a single length scale \(\eta = (V^2 / N)^{1/6} / \sqrt{2\pi}\) determines both rc and alpha. Callers who want to pin a specific cutoff (e.g. tied to neighbor-list update frequency in MD) should pass real_space_cutoff.

Parameters:

  • positions (torch.Tensor, shape (N, 3)) – Atomic coordinates.

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrix.

  • batch_idx (torch.Tensor, shape (N,), dtype=int32, optional) – System index for each atom.

  • accuracy (float, default=1e-6) – Target accuracy.

  • real_space_cutoff (float, optional) – Caller-supplied cutoff. When given, alpha is derived from it via \(\alpha = \sqrt{-\log\varepsilon} / r_c\); otherwise rc and alpha come from eta.

  • mesh_safety_factor (float, default=1.0) – Multiplier on the standard mesh-size heuristic \(K = 2\alpha L / (3\varepsilon^{1/5})\). Raise for extra safety at tight \(\varepsilon\).

Returns:

Dataclass containing alpha, mesh dimensions, spacing, and cutoffs.

Return type:

PMEParameters

nvalchemiops.torch.interactions.electrostatics.estimate_pme_mesh_dimensions(cell, alpha, accuracy=1e-6, mesh_safety_factor=1.0)[source]#

Estimate PME mesh dimensions for a given accuracy.

The mesh size along each axis is chosen as

\[K_i = \left\lceil \text{mesh\_safety\_factor} \cdot \frac{2\alpha L_i}{3\varepsilon^{1/5}} \right\rceil\]

rounded up to the next power of 2. The fifth-root scaling \(\varepsilon^{1/5}\) is the standard heuristic used by production PME codes; it grows the safety margin faster than \(\sqrt{-\ln\varepsilon}\) as \(\varepsilon\) tightens, which is empirically necessary to cover both the Gaussian-decay truncation and the B-spline aliasing error at the accuracies typically requested (1e-3 to 1e-6) across a wide (alpha, L, spline_order) envelope.

The canonical Essmann lower bound \(2\alpha L\sqrt{-\ln\varepsilon}/\pi\) is the Gaussian-decay term only; it can under-allocate by 2-4x at low alpha (large rc), where the B-spline aliasing term dominates.

Parameters:

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrix.

  • alpha (torch.Tensor, shape (B,)) – Ewald splitting parameter.

  • accuracy (float, default=1e-6) – Target relative accuracy.

  • mesh_safety_factor (float, default=1.0) – Multiplier on the standard heuristic. 1.0 is the well-tested default that meets accuracy across the configurations covered by the convergence script. Raise for extra paranoia at tight accuracy. Lower at your own risk: values below 1.0 can fail the accuracy guarantee on low-\(\alpha\) / large-L systems (verify with the convergence script before using).

Returns:

Maximum mesh dimensions (nx, ny, nz) across all systems in batch.

Return type:

tuple[int, int, int]

nvalchemiops.torch.interactions.electrostatics.mesh_spacing_to_dimensions(cell, mesh_spacing)[source]#

Convert mesh spacing to mesh dimensions.

Parameters:
Returns:

Mesh dimensions, rounded up to powers of 2.

Return type:

tuple[int, int, int]

class nvalchemiops.torch.interactions.electrostatics.EwaldParameters(alpha, real_space_cutoff, reciprocal_space_cutoff)[source]#

Container for Ewald summation parameters.

All values are tensors of shape (B,), for single system calculations, the shape is (1,).

Parameters:
alpha#

Ewald splitting parameter (inverse length units).

Type:

torch.Tensor, shape (B,)

real_space_cutoff#

Real-space cutoff distance.

Type:

torch.Tensor, shape (B,)

reciprocal_space_cutoff#

Reciprocal-space cutoff (\(|k|\) in inverse length units).

Type:

torch.Tensor, shape (B,)

class nvalchemiops.torch.interactions.electrostatics.PMEParameters(alpha, mesh_dimensions, mesh_spacing, real_space_cutoff)[source]#

Container for PME parameters.

Parameters:
alpha#

Ewald splitting parameter.

Type:

torch.Tensor, shape (B,)

mesh_dimensions#

Mesh dimensions (nx, ny, nz).

Type:

tuple[int, int, int], shape (3,)

mesh_spacing#

Actual mesh spacing in each direction.

Type:

torch.Tensor, shape (B, 3)

real_space_cutoff#

Real-space cutoff distance.

Type:

torch.Tensor, shape (B,)

Multipole Electrostatics#

GTO-smeared multipole electrostatics for systems carrying per-atom charges, dipoles, and quadrupoles (l_max 0/1/2). Moments are passed as a single packed multipole_moments tensor built with pack_multipole_moments().

High-Level Interface#

nvalchemiops.torch.interactions.electrostatics.multipole_ewald_summation(positions, multipole_moments, cell, idx_j, neighbor_ptr, unit_shifts, *, sigma, alpha=None, k_cutoff=None, batch_idx=None, accuracy=1e-6, cost_ratio=1.0, half_neighbor_list=False, cache=None)[source]#

Full GTO-Ewald multipole electrostatic total energy.

Composes the three canonical Ewald pieces:

  • real-space: GTO-Ewald-damped pair sum on a CSR neighbor list, via multipole_real_space_energy() or its batched analog, using \(T^{(0)}(r) = [\operatorname{erf}(r/(2\sigma)) - \operatorname{erf}(r/(2\sigma_c))] / r\) with \(\sigma_c = \sqrt{\sigma^2 + 1/(4\alpha^2)}\).

  • reciprocal-space: GTO-smeared Fourier sum with Ewald damping \(\exp(-k^2/(4\alpha^2))/k^2\), via multipole_reciprocal_space_energy() or its batched analog (raw total, no self-subtract).

  • self-energy correction: the analytical per-atom term from _multipole_ewald_self_energy_per_atom(), subtracted to remove the \(i=j\), \(n=0\) image that the reciprocal sum includes.

The total is mathematically identical to multipole_electrostatic_energy() (direct k-space) for any \((\sigma, \alpha)\).

Single-system vs batched dispatch#

Controlled by the optional batch_idx argument:

  • batch_idx=None (default) — single system. cell is \((3, 3)\) or \((1, 3, 3)\); returns per-atom \((N,)\) float64.

  • batch_idx provided (shape \((N_\text{total},)\)) — B systems packed into flat per-atom tensors. cell must be \((B, 3, 3)\); each atom’s neighbors must live in the same system. Returns per-atom \((N_\text{total},)\) flat across systems.

sigma and alpha are scalar floats in both modes (uniform across the batch).

param positions:

type positions:

torch.Tensor, shape (N, 3) or (N_total, 3)

param multipole_moments:

type multipole_moments:

torch.Tensor, shape (N, (l_max+1)**2) matching positions

param cell:

type cell:

torch.Tensor, shape (3, 3), (1, 3, 3), or (B, 3, 3)

param idx_j:

Flat CSR neighbor list (int32 / int32 / vec3i).

type idx_j:

torch.Tensor

param neighbor_ptr:

Flat CSR neighbor list (int32 / int32 / vec3i).

type neighbor_ptr:

torch.Tensor

param unit_shifts:

Flat CSR neighbor list (int32 / int32 / vec3i).

type unit_shifts:

torch.Tensor

param sigma:

GTO density-basis width.

type sigma:

float

param alpha:

Ewald splitting parameter (positive). When None (default) it is auto-estimated from sigma and the system geometry via estimate_multipole_ewald_parameters() at the requested accuracy. The caller must build the neighbor list with the matching real-space cutoff.

type alpha:

float, optional

param k_cutoff:

Maximum \(|k|\) for the reciprocal sum. Auto-estimated from the Kolafa-Perram balance when None (single mode); batched mode always builds its own k-grid per system. Unnecessary when a cache is supplied (the cache already encodes the k-grid).

type k_cutoff:

float, optional

param batch_idx:

\((N_\text{total},)\) int32. None selects single-system mode.

type batch_idx:

torch.Tensor, optional

param accuracy:

Target relative-energy accuracy used by the auto-estimator when alpha and/or k_cutoff are None. Ignored if both are supplied.

type accuracy:

float, default 1e-6

param cost_ratio:

Hardware-empirical \(C_r / C_k\) (per-real-space-pair vs per-k cost ratio) passed to estimate_multipole_ewald_parameters(). 1.0 reproduces canonical Kolafa-Perram; higher values shift the optimum toward smaller real-space cutoff and larger reciprocal cutoff. Ignored if alpha and k_cutoff are supplied.

type cost_ratio:

float, default 1.0

param half_neighbor_list:

Forwarded to the real-space pair sum. Set when the CSR neighbor list stores each pair once (only affects the \(l_{max}=2\) path’s cell-gradient kernel; no effect otherwise).

type half_neighbor_list:

bool, default False

param cache:

Pre-built reciprocal cache (from prepare_multipole_scf_cache()). When given, the per-call reciprocal-cache rebuild (k-grid + GTO-Fourier phi_hat + per-k-factor tables) is skipped — the steady-state / MD path, analogous to passing precomputed k_squared to multipole_particle_mesh_ewald(). alpha is taken from the cache when not supplied, and k_cutoff becomes unnecessary. Stress caveat: a pre-built cache holds a fixed (detached) k-grid/volume, so grad(E, cell) does not flow through the reciprocal term — pass cache=None for stress / cell-gradient training. Forces (grad(E, positions)) are unaffected.

type cache:

MultipoleSCFCache, optional

returns:

Per-atom \((N,)\) \(\text{float64}\) (single) or \((N_\text{total},)\) (batched, flat across systems) on positions.device. Call .sum() for the total energy or torch.zeros(B).scatter_add(0, batch_idx, E) for per-system totals; forces/stress/charge-grads flow from grad(E.sum(), ...). Autograd-connected to positions and multipole_moments.

rtype:

torch.Tensor

Parameters:
Return type:

torch.Tensor

nvalchemiops.torch.interactions.electrostatics.pme_multipole.multipole_particle_mesh_ewald(positions, multipole_moments, cell, idx_j, neighbor_ptr, unit_shifts, *, sigma, alpha=None, mesh_dimensions=None, spline_order=4, cell_inv_t=None, batch_idx=None, accuracy=1e-6, cost_ratio=1.0, volume=None, moduli=None, k_squared=None)[source]#

Total multipole PME energy (real + reciprocal − self − background).

Single-system or batched, dispatched by batch_idx. Mirrors the API of multipole_ewald_summation (single function with optional batch_idx) but routes the \(O(N \cdot N_k)\) reciprocal half through PME, dropping the asymptotic cost to \(O(M \log M + N \cdot p^3)\) at the cost of a small spline-truncation residual.

The composition matches the direct-k multipole_ewald_summation:

  1. Real-space pair sum via multipole_real_space_energy() (single, or batched through its batch_idx= path). Per-atom output is multiplied by coulomb_scale = F/(4*pi).

  2. Reciprocal piece via multipole_pme_reciprocal_space() — returns E_recip - E_self - E_bg already in F units.

  3. E_total = E_real + E_recip - E_self - E_bg.

Direct-k parity holds at the spline-truncation floor (rtol ~ 1e-4 at mesh = 60^3, L = 10).

Parameters:

  • positions (torch.Tensor, shape (N, 3) or (N_total, 3)) – Cartesian atom positions; N_total in batched mode.

  • multipole_moments (torch.Tensor, shape (N, 1), (N, 4), or) – (N, 9) (or the N_total analog in batched mode). e3nn spherical-harmonic packing: [q] (l_max=0), [q, mu_y, mu_z, mu_x] (l_max=1), or the l_max=1 block plus the five traceless l=2 channels (l_max=2). The trailing dim selects the l_max path.

  • cell (torch.Tensor, shape (3, 3), (1, 3, 3), or (B, 3, 3)) – Unit-cell matrix (rows are lattice vectors); (B, 3, 3) batched.

  • idx_j (torch.Tensor) – Real-space CSR neighbor list. In batched mode the list is flat across systems; each atom’s neighbors must live in the same system (caller’s responsibility).

  • neighbor_ptr (torch.Tensor) – Real-space CSR neighbor list. In batched mode the list is flat across systems; each atom’s neighbors must live in the same system (caller’s responsibility).

  • unit_shifts (torch.Tensor) – Real-space CSR neighbor list. In batched mode the list is flat across systems; each atom’s neighbors must live in the same system (caller’s responsibility).

  • sigma (float) – GTO basis width — uniform across batch.

  • alpha (float, optional) – Ewald splitting parameter (positive). When None (default) it is auto-estimated from sigma and the system geometry via estimate_multipole_pme_parameters() at the requested accuracy. The caller is still responsible for having built the neighbor list with the matching real-space cutoff (also available via the same estimator).

  • mesh_dimensions (tuple[int, int, int], optional) – FFT mesh dimensions — shared across batch. Auto-estimated from the same Kolafa-Perram balance when None. Override this if you need to lock the mesh resolution (e.g. for kernel reuse).

  • spline_order (int, default 4) – B-spline interpolation order p (used for both spread and gather).

  • cell_inv_t (torch.Tensor, optional) – Pre-computed transpose(inv(cell)) — shape (3, 3) / (1, 3, 3) for single-system, (B, 3, 3) for batched.

  • batch_idx (torch.Tensor, optional) – (N_total,) int32. Triggers the batched path when provided.

  • accuracy (float, default 1e-6) – Target relative-energy accuracy used by the auto-estimator when alpha and/or mesh_dimensions are None. Same semantics as the monopole particle_mesh_ewald().

  • cost_ratio (float, default 1.0) – Hardware-empirical per-real-space-pair vs per-reciprocal cost ratio passed through to estimate_multipole_pme_parameters(). 1.0 reproduces canonical Kolafa-Perram. Note that PME’s true reciprocal cost (FFT + spread/gather) doesn’t have the same shape as Ewald’s per-k cost — the Ewald optimum may not transfer directly. See estimate_multipole_pme_parameters() for details. Ignored if alpha and mesh_dimensions are supplied.

  • volume (torch.Tensor or None, optional) – Cell volume(s) forwarded to the reciprocal half; () / (1,) single-system, (B,) batched. When None it is computed from cell. Passing a tensor keeps cell autograd (stress) alive.

  • moduli (tuple[torch.Tensor, torch.Tensor, torch.Tensor] or None, optional) – Pre-computed 1-D B-spline modulus LUTs (b_x, b_y, b_z) of shapes (Nx,), (Ny,), (Nz_rfft,) for MD steady-state reuse; forwarded to multipole_pme_reciprocal_space(). When None they are computed from mesh_dimensions and spline_order.

  • k_squared (torch.Tensor or None, optional) – Pre-computed \(|k|^2\) rfft grid of shape (Nx, Ny, Nz_rfft) (single-system) or (B, Nx, Ny, Nz_rfft) (batched), cacheable across MD steps when the cell is fixed. When None it is built from cell and mesh_dimensions.

Returns:

energy – Per-atom \((N,)\) (single) or \((N_\text{total},)\) (batched, flat across systems). Call .sum() for the total Coulomb energy or torch.zeros(B).scatter_add(0, batch_idx, E) for per-system totals; forces/stress/charge-grads flow from grad(E.sum(), ...).

Return type:

torch.Tensor, float64

Energy Components#

nvalchemiops.torch.interactions.electrostatics.multipole_electrostatic_energy(positions, multipole_moments, cell, *, batch_idx=None, sigma, k_cutoff=None, k_vectors=None, normalize=NormMode.MULTIPOLES, include_self_interaction=False)[source]#

Total PBC electrostatic energy via direct k-space summation.

Computes

\[E \;=\; \frac{1}{2} \cdot \frac{V}{(2\pi)^6} \sum_{\mathbf{k}} 2\,\text{Re}\!\left[\rho^{*}(\mathbf{k})\, V(\mathbf{k})\right] \;-\; \tfrac{1}{2} E_{\text{self}},\]

where \(\rho(\mathbf{k})\) and \(V(\mathbf{k}) = F \cdot \rho(\mathbf{k}) / k^2\) are assembled from per-atom multipole_moments via the Warp kernels. Matches the customer reference GTOElectrostaticEnergy bit-for-bit at l_max in {0, 1} under matched inputs.

Single-system vs batched dispatch#

Mirrors multipole_ewald_summation(): pass cell of shape (3, 3) (single) or (B, 3, 3) (batched) and use batch_idx to select the batched path (returns per-atom \((N_\text{total},)\)). Batched mode requires k_cutoff (a pre-generated k_vectors is single-system only).

param positions:

Atomic positions, shape (N, 3) or (N_total, 3) (flat across systems in the batched case), float32 or float64.

type positions:

torch.Tensor

param multipole_moments:

Packed per-atom multipole moments, shape (N, (l_max+1)**2), in e3nn spherical layout: [q] (l_max=0), [q, mu_y, mu_z, mu_x] (l_max=1), or the l_max=1 block plus the five traceless l=2 channels (l_max=2). The l=2 quadrupole is expanded to the Cartesian symmetric (N, 3, 3) form and threaded through the SCF-cache Q channel.

type multipole_moments:

torch.Tensor

param cell:

Unit-cell matrix (lattice vectors as rows), shape (3, 3), or B per-system cells (B, 3, 3) (batched).

type cell:

torch.Tensor

param batch_idx:

Per-atom system index (expected sorted). Required when cell is (B, 3, 3); must be None for a single (3, 3) cell.

type batch_idx:

torch.Tensor, optional, shape (N_total,), int32

param sigma:

Density-basis Gaussian width. Used for both the source GTO basis and the self-interaction overlap (matches GTOElectrostaticEnergy).

type sigma:

float

param k_cutoff:

Maximum |k| to include in the reciprocal-space sum. Required when k_vectors is not supplied; ignored when it is.

type k_cutoff:

float, optional

param k_vectors:

Pre-computed k-grid, shape (N_k, 3), float64. Must include ``(0, 0, 0)`` as the first row — the kernel’s V(k=0) = 0 convention expects the origin explicitly and downstream indexing assumes row 0 is it. Pass this when amortizing k-vector generation across many energy evaluations for the same geometry (MD steps at fixed cell, SCF iterations, benchmark loops). Must live on positions.device. When omitted, the function generates k-vectors internally via generate_k_vectors_ewald_summation(cell, k_cutoff) and prepends the origin.

type k_vectors:

torch.Tensor, optional

param normalize:

Normalization convention for the density basis. Defaults to NormMode.MULTIPOLES (the only physically meaningful choice for source moments; the other modes exist for debugging / cross-checks).

type normalize:

NormMode | int | str

param include_self_interaction:

If False (default), subtracts \(0.5 \cdot E_{\text{self}}\) where \(E_\text{self} = \sum_i \mathrm{oc}[0]\,q_i^2 + \mathrm{oc}[1]\,|\boldsymbol{\mu}_i|^2\) and oc comes from nvalchemiops.torch.math.compute_overlap_constants().

type include_self_interaction:

bool

returns:

Per-atom \((N,)\) \(\text{float64}\) (single) or \((N_\text{total},)\) (batched, flat across systems) on positions.device. Call .sum() for the total energy or torch.zeros(B).scatter_add(0, batch_idx, E) for per-system totals; forces/stress/charge-grads flow from grad(E.sum(), ...). Autograd-connected to positions and multipole_moments.

rtype:

torch.Tensor

Parameters:
Return type:

Tensor

nvalchemiops.torch.interactions.electrostatics.multipole_real_space_energy(positions, multipole_moments, cell, idx_j, neighbor_ptr, unit_shifts, sigma, alpha, *, batch_idx=None, half_neighbor_list=False)[source]#

GTO-Ewald real-space multipole energy.

Unified entry point covering \(l_{max} \in \{0, 1, 2\}\). The trailing dim of multipole_moments selects the path:

  • \((N, 1)\) -> \(l_{max}=0\) (charges only)

  • \((N, 4)\) -> \(l_{max}=1\) (charges + dipoles, e3nn \((y, z, x)\) order)

  • \((N, 9)\) -> \(l_{max}=2\) (adds the quadrupole channels)

Channel [:, 0] is the charge and [:, 1:4] is the dipole in e3nn \((y, z, x)\) spherical order; the wrapper permutes to Cartesian before calling the Warp kernel.

Returns per-atom \((N,)\) \(\text{float64}\) for all paths; the caller owns the atom-global reduction (.sum() for total energy, scatter_add for per-system totals). Non-uniform per-atom backward weights are supported across all \(l_{max}\).

Single-system vs batched dispatch#

Mirrors multipole_ewald_summation(): pass cell of shape (3, 3) / (1, 3, 3) (single) or (B, 3, 3) (batched) and use batch_idx to select the batched path. In batched mode sigma and alpha are per-system (B,) tensors and the return is per-atom \((N_\text{total},)\) for all \(l_{max}\).

param positions:

Atomic positions ((N_total, 3) flat across systems when batched).

type positions:

torch.Tensor, shape (N, 3)

param multipole_moments:

Per-atom multipole moments in e3nn spherical layout.

type multipole_moments:

torch.Tensor, shape (N, (l_max + 1)**2)

param cell:

Unit cell matrix (row vectors = lattice vectors); (B, 3, 3) when batched.

type cell:

torch.Tensor, shape (3, 3) / (1, 3, 3) or (B, 3, 3)

param idx_j:

CSR neighbor list (int32 / int32 / vec3i).

type idx_j:

torch.Tensor

param neighbor_ptr:

CSR neighbor list (int32 / int32 / vec3i).

type neighbor_ptr:

torch.Tensor

param unit_shifts:

CSR neighbor list (int32 / int32 / vec3i).

type unit_shifts:

torch.Tensor

param sigma:

GTO density-basis width \(\sigma\) (per-system (B,) when batched).

type sigma:

torch.Tensor, shape (1,) or (B,)

param alpha:

Ewald splitting parameter \(\alpha\) (per-system (B,) when batched).

type alpha:

torch.Tensor, shape (1,) or (B,)

param batch_idx:

Per-atom system index (expected sorted). Required when cell is (B, 3, 3); must be None for a single cell.

type batch_idx:

torch.Tensor, optional, shape (N_total,), int32

param half_neighbor_list:

Forwarded to the \(l_{max}=2\) path (no effect otherwise).

type half_neighbor_list:

bool, default False

returns:

Per-atom real-space energy (single), or per-atom \((N_\text{total},)\) batched — for all \(l_{max}\).

rtype:

torch.Tensor, shape (N,), float64

Parameters:
Return type:

Tensor

nvalchemiops.torch.interactions.electrostatics.multipole_reciprocal_space_energy(positions, multipole_moments, cell, *, batch_idx=None, sigma, alpha, k_cutoff=None, normalize=NormMode.MULTIPOLES, cache=None)[source]#

Reciprocal-space half of an Ewald-split multipole electrostatic energy.

Same pipeline as multipole_electrostatic_energy() but with a Gaussian-damped per-k kernel:

\[V(\mathbf{k}) = \frac{F \, e^{-|\mathbf{k}|^2 / (4\alpha^2)}}{|\mathbf{k}|^2} \, \rho(\mathbf{k}),\]

(k = 0 zeroed). Intended to be paired with a real-space erfc-damped contribution (see multipole_real_space_energy()) at the same \(\alpha\) to assemble the full Ewald-split Coulomb sum.

Single-system vs batched dispatch#

Mirrors multipole_ewald_summation(): pass cell of shape (3, 3) (single) or (B, 3, 3) (batched) and use batch_idx to select the batched path (returns per-atom \((N_\text{total},)\)). Both single and batched modes build their k-grid from k_cutoff (or reuse a pre-built cache).

param positions:

Same as multipole_electrostatic_energy().

param multipole_moments:

Same as multipole_electrostatic_energy().

param cell:

Same as multipole_electrostatic_energy().

param sigma:

Same as multipole_electrostatic_energy().

param k_cutoff:

Same as multipole_electrostatic_energy().

param normalize:

Same as multipole_electrostatic_energy().

param batch_idx:

Per-atom system index (expected sorted). Required when cell is (B, 3, 3); must be None for a single (3, 3) cell.

type batch_idx:

torch.Tensor, optional, shape (N_total,), int32

param alpha:

Ewald splitting parameter (must be positive). The caller’s real-space kernel should use the same alpha.

type alpha:

float

param cache:

Pre-built reciprocal cache (from prepare_multipole_scf_cache()) holding the position-independent k-grid / GTO-Fourier (phi_hat) / per-k-factor tables. When given, the per-call cache rebuild is skipped (MD / inference steady state) — the analog of passing precomputed k_squared to PME. cell/sigma/alpha/k_cutoff are then ignored for the reciprocal (the cache already encodes them); the caller owns matching the cache to the system.

Warning

A pre-built cache holds a fixed (detached) k-grid and volume, so grad(E, cell) (stress) does not flow through it — pass cache=None for stress / cell-gradient training. Forces (grad(E, positions)) are unaffected (positions enter per call).

type cache:

MultipoleSCFCache, optional

returns:

Per-atom \((N,)\) \(\text{float64}\) (single) or \((N_\text{total},)\) (batched, flat across systems) on positions.device. Does not subtract any self-interaction correction — the caller combines this with the real-space and self / background terms to get the full Ewald total. Call .sum() for the total or torch.zeros(B).scatter_add(0, batch_idx, E) for per-system totals.

rtype:

torch.Tensor

Parameters:
Return type:

torch.Tensor

Atom-Centered Features#

nvalchemiops.torch.interactions.electrostatics.multipole_electrostatic_features(positions, multipole_moments, cell, *, batch_idx=None, sigma, receiver_sigmas, k_cutoff=None, k_vectors=None, feature_max_l=1, density_normalize=NormMode.MULTIPOLES, feature_normalize=NormMode.RECEIVER, include_self_interaction=False)[source]#

Atom-centered electrostatic features via direct k-space projection.

Computes

\[f_{i, \sigma_r, l, m} \;=\; \frac{2}{(2\pi)^3} \sum_{\mathbf{k}} w(\mathbf{k}) \cdot \text{Re}\!\left[V^{*}(\mathbf{k})\, \hat\phi_{l,m}^{\sigma_r}(\mathbf{k})\, e^{i\mathbf{k}\cdot\mathbf{r}_i}\right],\]

where \(V(\mathbf{k})\) is the periodic electrostatic potential assembled from multipole_moments as in the companion energy binding, \(\hat\phi_{l,m}^{\sigma_r}\) is the receiver-basis GTO Fourier coefficient, and \(w(\mathbf{k})\) is \(0.5\) at \(k = 0\) and \(1\) elsewhere (the half-space-with-origin convention matching the reference k_factor_proj).

Bit-for-bit parity with the customer reference GTOElectrostaticFeatures at density_max_l in \(\{0, 1\}\), feature_max_l = 1, under matched inputs.

Single-system vs batched dispatch#

Mirrors multipole_ewald_summation(): pass cell of shape (3, 3) (single) or (B, 3, 3) (batched) and use batch_idx to select the batched path. Batched mode requires k_cutoff (a pre-generated k_vectors is single-system only).

param positions:

Atomic positions; flat across systems in the batched case.

type positions:

torch.Tensor, shape (N, 3) or (N_total, 3)

param multipole_moments:

Packed per-atom source moments in e3nn spherical layout. Source l_max is inferred; an \(l_{max}=2\) (N, 9) tensor enriches the projected potential with its Cartesian-quadrupole channel.

type multipole_moments:

torch.Tensor, shape (N, (l_max+1)**2)

param cell:

Single unit cell or B per-system cells (batched).

type cell:

torch.Tensor, shape (3, 3) or (B, 3, 3)

param batch_idx:

Per-atom system index (expected sorted). Required when cell is (B, 3, 3); must be None for a single (3, 3) cell.

type batch_idx:

torch.Tensor, optional, shape (N_total,), int32

param sigma:

Density-side Gaussian width.

type sigma:

float

param receiver_sigmas:

Multi-\(\sigma\) receiver basis widths. Must be non-empty.

type receiver_sigmas:

list of floats, tuple, or 1-D tensor

param k_cutoff:

Same semantics as multipole_electrostatic_energy(). Pass k_vectors to amortize setup across calls for fixed geometry (single-system only). Batched mode requires k_cutoff.

param k_vectors:

Same semantics as multipole_electrostatic_energy(). Pass k_vectors to amortize setup across calls for fixed geometry (single-system only). Batched mode requires k_cutoff.

param feature_max_l:

Receiver angular cap: how many l-blocks are projected out per \(\sigma\), independent of the source l_max. 1 -> (N_sigma * 4) features (\(l \le 1\)); 2 -> (N_sigma * 9) (adds the 5 \(l = 2\) receiver channels). The \(l = 2\) self-interaction subtract uses the source’s e3nn \(l = 2\) moment (zero when the source has no quadrupole).

type feature_max_l:

int, default 1

param density_normalize:

Source-basis normalization. Defaults to MULTIPOLES.

type density_normalize:

NormMode | int | str

param feature_normalize:

Receiver-basis normalization. Defaults to RECEIVER, matching GTOElectrostaticFeatures’s integral_normalization default.

type feature_normalize:

NormMode | int | str

param include_self_interaction:

If False (default), subtract the self-interaction term using compute_overlap_constants().

type include_self_interaction:

bool

returns:

float64 on positions.device, shape (N, N_sigma * (feature_max_l + 1)**2) in the reference permuted-flat layout (grouped by l-block). Autograd-connected to positions and multipole_moments.

rtype:

torch.Tensor

Parameters:
Return type:

Tensor

Moment Packing#

nvalchemiops.torch.interactions.electrostatics.pack_multipole_moments(charges, dipoles=None, quadrupoles=None, *, trace_atol=1e-8)[source]#

Build packed e3nn multipole_moments from Cartesian channels.

Parameters:

  • charges ((N,))

  • dipoles ((N, 3) Cartesian (x, y, z) or None.)

  • quadrupoles ((N, 3, 3) symmetric Cartesian or None. The l=2) – channel is traceless; a non-negligible input trace (> trace_atol) is dropped and raises a warning.

  • trace_atol (float, default 1e-8) – Absolute tolerance on max |Tr Q| above which the dropped quadrupole trace triggers a warning.

Returns:

Packed e3nn multipole_moments (last-dim 1, 4, or 9 depending on which channels were supplied), contiguous, on charges.device.

Return type:

torch.Tensor, shape (N, (l_max+1)**2)

Raises:

ValueError – If quadrupoles is given without dipoles (the packed l_max=2 layout requires the (N, 4) charge+dipole block).

SCF Cache (Amortized Workflow)#

Reuse the position-independent reciprocal-space state across many evaluations at fixed cell (MD steps / SCF iterations).

nvalchemiops.torch.interactions.electrostatics.prepare_multipole_scf_cache(cell, *, sigma, receiver_sigmas, k_cutoff=None, k_vectors=None, l_max=1, feature_max_l=1, density_normalize=NormMode.MULTIPOLES, feature_normalize=NormMode.RECEIVER, alpha=None, device=None)[source]#

Build a MultipoleSCFCache from the position-independent inputs.

Runs the position-independent direct-k-space geometry kernels (eval_gto_fourier_dipole for the source basis, eval_receiver_gto_fourier_dipole for the receiver basis) and precomputes the per-k and per-\(\sigma\) factor tables that the step functions consume.

The structure-factor table (cos/sin) is not part of the cache: it depends on atomic positions, which flow through MultipoleRhoFunction with autograd, so the table is recomputed from positions on every step to wire up the position gradient.

Single-system vs batched dispatch#

cell of shape (3, 3) builds a single-system cache. cell of shape (B, 3, 3) builds a batched cache (n_systems == B) whose per-k tensors carry a leading-B zero-padded layout; in that case k_cutoff is required (a pre-generated k_vectors is not supported for the batched build).

param cell:

Single unit cell, or B per-system unit cells (batched build).

type cell:

torch.Tensor, shape (3, 3) or (B, 3, 3)

param sigma:

Density-basis Gaussian width.

type sigma:

float

param receiver_sigmas:

Multi-sigma receiver widths. Must be non-empty.

type receiver_sigmas:

list / tuple / 1-D tensor of floats

param k_cutoff:

Same semantics as the step functions: either pass k_cutoff to generate the k-grid internally (with origin prepended), or pass a pre-generated k_vectors tensor ((N_k, 3) float64 with origin at row 0). The batched build (cell (B, 3, 3)) requires k_cutoff.

param k_vectors:

Same semantics as the step functions: either pass k_cutoff to generate the k-grid internally (with origin prepended), or pass a pre-generated k_vectors tensor ((N_k, 3) float64 with origin at row 0). The batched build (cell (B, 3, 3)) requires k_cutoff.

param l_max:

Source multipole order the cache is being built for. 0 or 1. Only affects the feature_overlap_constants layout (the l=1 column is zeroed when l_max = 0); all other tensors are l_max independent.

type l_max:

int

param feature_max_l:

Receiver feature angular cap (independent of the source l_max). 0, 1, or 2. Selects the receiver \(\hat\phi\) width ((..., 4, 2) for \(\le 1\), (..., 9, 2) for 2) and the feature output width (feature_max_l + 1)**2 per sigma.

type feature_max_l:

int, default 1

param density_normalize:

Normalization conventions.

type density_normalize:

NormMode | int | str

param feature_normalize:

Normalization conventions.

type feature_normalize:

NormMode | int | str

param alpha:

Ewald splitting parameter. None (default) selects the direct-k-space Coulomb factor \(\text{per\_k\_factor} = F / k^2\) with the origin zeroed. A positive value selects the Ewald-damped reciprocal-space factor: \(\text{per\_k\_factor} = F \exp(-k^2/(4\alpha^2)) / k^2\) with the origin zeroed, the Gaussian-smoothed reciprocal-space half of an Ewald split, paired with a real-space erfc contribution to assemble the full Coulomb sum.

type alpha:

float, optional

param device:

Device for all cache tensors. Defaults to cell.device.

type device:

torch.device | str, optional

returns:

Frozen dataclass. All tensors live on the resolved device.

rtype:

MultipoleSCFCache

Parameters:
Return type:

MultipoleSCFCache

nvalchemiops.torch.interactions.electrostatics.multipole_scf_step_energy(cache, positions, source_feats, *, batch_idx=None, include_self_interaction=False, quadrupoles=None)[source]#

Per-atom PBC electrostatic energies for a single or batched SCF step.

Consumes the position-independent tensors in cache and the per-step (positions, source_feats). Returns per-atom \((N,)\) \(\text{float64}\) (single) or \((N_\text{total},)\) (batched, flat across all systems). The result is autograd-connected to both inputs.

The caller owns the reduction: .sum() for the total energy, or torch.zeros(B).scatter_add(0, batch_idx, E) for per-system totals. Forces, stress, and charge-grads are obtained via grad(E.sum(), ...). The per-atom energies sum to the same total bit-for-bit as the old scalar / (B,) convention.

Parameters:

  • cache (MultipoleSCFCache) – Prebuilt cache from prepare_multipole_scf_cache(). Holds the position-independent state (k-vectors, \(\hat\phi\), per-k factors, overlap constants, LUTs). Single or batched.

  • positions (torch.Tensor, shape (N_atoms, 3) or (N_total, 3)) – Atomic positions (flat across systems in the batched case).

  • source_feats (torch.Tensor, shape (N, (l_max + 1)**2)) – Packed per-atom moments in e3nn spherical layout. \((N, 1)\) for l_max=0, \((N, 4)\) for l_max=1 ([q, mu_y, mu_z, mu_x]). Must match cache.l_max.

  • batch_idx (torch.Tensor, optional, shape (N_total,), int32) – Per-atom system index (expected sorted so atoms group by system). Required when cache is batched; must be None for a single system. Mirrors multipole_ewald_summation()’s convention.

  • include_self_interaction (bool) – If False (default), subtract \(0.5 \cdot E_\text{self}\) using cache.source_overlap_constants. True returns the raw reciprocal sum (the Ewald path subtracts the self term itself).

  • quadrupoles (torch.Tensor, optional, shape (N, 3, 3)) – Cartesian symmetric source quadrupole. When supplied, the additive \(\rho_Q(k)\) channel and its \(l=2\) self term are included (requires a cache built with source l_max>=2). None (default) is the \(l_{max} \le 1\) path.

Returns:

Per-atom \((N,)\) \(\text{float64}\) (single) or \((N_\text{total},)\) (batched, flat across systems) on cache.device. Autograd-connected to positions and source_feats. Call .sum() for the total energy or torch.zeros(B).scatter_add(0, batch_idx, E) for per-system totals; forces/stress/charge-grads flow from grad(E.sum(), ...).

Return type:

torch.Tensor

nvalchemiops.torch.interactions.electrostatics.multipole_scf_step_features(cache, positions, source_feats, *, batch_idx=None, include_self_interaction=False, quadrupoles=None)[source]#

Atom-centered multipole features for a single or batched SCF step.

Consumes the position-independent cache plus the per-step (positions, source_feats) and returns a \((N, N_\sigma \cdot (\text{feature\_max\_l}+1)^2)\) features tensor in the reference permuted flat layout (grouped by l-block). With a batched cache (cache.is_batched and batch_idx supplied) the rows span all systems flat-packed.

Parameters:

  • cache (MultipoleSCFCache) – Single or batched.

  • positions (torch.Tensor, shape (N_atoms, 3) or (N_total, 3))

  • source_feats (torch.Tensor, shape (N, (l_max + 1)**2)) – Packed per-atom moments in e3nn spherical layout.

  • batch_idx (torch.Tensor, optional, shape (N_total,), int32) – Per-atom system index (expected sorted). Required when cache is batched; must be None for a single system.

  • include_self_interaction (bool) – If False (default), subtract the self-interaction correction using cache.feature_overlap_constants.

  • quadrupoles (torch.Tensor, optional, shape (N, 3, 3)) – Cartesian-quadrupole source moments. When supplied, the additive \(\rho_Q(k)\) channel enriches the projected potential. Requires a cache built with source l_max>=2. This is the source l=2 contribution, decoupled from the receiver feature_max_l (which lives on the cache and controls how many l-blocks are projected out).

Returns:

\(\text{float64}\) on cache.device, shape \((N, N_\sigma \cdot (\text{feature\_max\_l}+1)^2)\) in the reference permuted flat layout (grouped by l-block). Autograd-connected to positions and source_feats.

Return type:

torch.Tensor

class nvalchemiops.torch.interactions.electrostatics.MultipoleSCFCache(k_vectors, k_norm2, source_phi_hat, receiver_phi_hat, per_k_factor, k_factor_proj, source_overlap_constants, feature_overlap_constants, out_col_lut_natural, out_col_lut_permuted, out_col_inv_perm, volume, cell, sigma, alpha, receiver_sigmas, l_max, density_normalize, feature_normalize, n_systems=1, valid_k_counts=None, feature_max_l=1, source_coeff2=None, feature_overlap_l2=None)[source]#

Frozen bundle of geometry-only direct-k-space tensors for the SCF step functions.

Represents a single system (n_systems == 1) or a batch of B systems (n_systems == B) under one unified dataclass. Built by prepare_multipole_scf_cache() from a (3, 3) (single) or (B, 3, 3) (batched) cell; consumed by multipole_scf_step_energy / multipole_scf_step_features (single) and their batched branches.

For the batched case every per-k tensor carries a leading-B layout and is uniform-shape (B, K_max, ...) with zero padding beyond each system’s K_b valid k-vectors. Pad rows get k_vectors = 0 (so k_alpha = 0 in any k-weighted sum), per_k_factor = 0 (zero Coulomb contribution), and k_factor_proj = 0 (zero feature-projection weight); that is sufficient to make every kernel in the direct-k-space pipeline ignore pad rows without in-kernel branching. valid_k_counts records each system’s K_b.

All tensor fields are float64 on the same device.

Parameters:
k_vectors#

Single: (N_k, 3) reciprocal-lattice k-grid with (0, 0, 0) at row 0. Batched: (B, K_max, 3) per-system grids zero-padded to K_max.

Type:

torch.Tensor, float64

k_norm2#

\(|k|^2\). Single (N_k,) / batched (B, K_max).

Type:

torch.Tensor, float64

source_phi_hat#

Source-basis GTO Fourier coefficients \(\hat\phi_{l,m}^{\sigma}(\mathbf{k})\). Single (N_k, 4, 2) / batched (B, K_max, 4, 2) (pad rows zeroed).

Type:

torch.Tensor, float64

receiver_phi_hat#

Receiver-basis GTO Fourier coefficients across all receiver sigma. Single (N_k, N_sigma, 4|9, 2) / batched (B, K_max, N_sigma, 4|9, 2).

Type:

torch.Tensor, float64

per_k_factor#

Coulomb multiplier \(\text{FIELD\_CONSTANT} / k^2\) (or the Ewald damped form) with the k = 0 entry zeroed. Single (N_k,) / batched (B, K_max) (pad + k = 0 rows zeroed).

Type:

torch.Tensor, float64

k_factor_proj#

Feature projection weight: 0.5 at real k = 0, 1 at real nonzero k, 0 at pad rows. Single (N_k,) / batched (B, K_max).

Type:

torch.Tensor, float64

source_overlap_constants#

Per-l self-overlap constants for the source basis (l=0, l=1, l=2), shared across the batch.

Type:

torch.Tensor, shape (3,), float64

feature_overlap_constants#

Per-(sigma, l) self-overlap constants for the receiver basis, shared across the batch. The l=1 column is zeroed when l_max == 0.

Type:

torch.Tensor, shape (N_sigma, 2), float64

out_col_lut_natural#

Natural row-major output LUT for the feature projection kernel.

Type:

torch.Tensor, shape (N_sigma, 4|9), int32

out_col_lut_permuted#

Permuted output LUT for the feature projection kernel.

Type:

torch.Tensor, shape (N_sigma, 4|9), int32

out_col_inv_perm#

Precomputed inverse permutation argsort(out_col_lut_permuted) that maps the natural feature layout to the reference permuted flat layout. Cached here (it is position-independent) so the feature step does not re-run torch.argsort every call.

Type:

torch.Tensor, shape (N_sigma * (feature_max_l+1)**2,), int64

volume#

|det(cell)|. Single shape () / batched (B,).

Type:

torch.Tensor, float64

cell#

The original unit-cell matrix/matrices. Single (3, 3) / batched (B, 3, 3).

Type:

torch.Tensor, float64

n_systems#

Number of systems: 1 for single, B for batched.

Type:

int

valid_k_counts#

Batched only: (B,) int32 of per-system valid k-counts K_b (K_max = valid_k_counts.max()). None for the single-system cache.

Type:

torch.Tensor or None

sigma#

Density-side Gaussian width.

Type:

float

alpha#

Ewald splitting parameter the cache was built with. None selects the direct-k-space Coulomb factor per_k_factor = F / k^2; a positive value selects the Ewald-damped reciprocal-space factor F exp(-k^2/(4 alpha^2)) / k^2.

Type:

float or None

receiver_sigmas#

Receiver (feature) \(\sigma\) widths, as an immutable tuple.

Type:

tuple of float

l_max#

Effective source multipole order this cache was built for (0 for charges-only, 1 otherwise). Used by the step functions for bookkeeping; the kernels always run the l_max=1 path with zeros for missing components.

Type:

int

density_normalize, feature_normalize

Normalization modes used when the cache was built, stored so multipole_scf_step_* can validate consistency.

Type:

NormMode

property batch_size: int#

Number of systems in the batch (n_systems; 1 if single).

property device: device#

Device all tensors live on.

feature_max_l: int = 1#

Receiver feature angular cap, independent of l_max (the source cap). receiver_phi_hat is (..., 4, 2) for feature_max_l <= 1 and (..., 9, 2) for feature_max_l == 2; feature output width is (feature_max_l + 1)**2 per sigma.

feature_overlap_l2: Tensor | None = None#

l=2 receiver self-overlap constant, shape (N_sigma,) float64, shared across the batch. None when feature_max_l < 2. Kept separate from the (N_sigma, 2) feature_overlap_constants so the l<=1 projection kernel’s shape contract is unchanged.

property is_batched: bool#

True when this cache holds a batch (cell is (B, 3, 3)).

property n_k: int#

Number of k-vectors in the (single-system) grid, including origin row 0.

property n_k_max: int#

Padded per-system k-vector count K_max (batched layout).

property n_sigma: int#

Number of receiver sigma widths.

source_coeff2: Tensor | None = None#

Cartesian-quadrupole per-k coefficient coeff2(k) = -0.5*phi0(k). Single (N_k,) / batched (B, K_max) float64. None when l_max < 2. Geometry-only (depends on k + sigma, not moments); consumed by the l=2 reciprocal channel.

Parameter Estimation#

nvalchemiops.torch.interactions.electrostatics.estimate_multipole_ewald_parameters(positions, cell, sigma, batch_idx=None, accuracy=1e-6, cost_ratio=1.0)[source]#

Estimate GTO-Ewald multipole parameters at a given target accuracy.

Mirrors estimate_ewald_parameters() semantics (“relative energy-error” accuracy via the Kolafa-Perram envelope), adjusted for the multipole case where the effective Ewald-split width is sigma_c = sqrt(sigma**2 + 1/(4 alpha**2)) rather than 1/(alpha sqrt(2)).

Derivation#

Both the real-space tail (erfc(r/(2 sigma_c))) and the reciprocal-space envelope (exp(-k**2 sigma_c**2)) decay with the same effective width sigma_c * sqrt(2). Substituting that for the monopole’s eta = 1/(alpha * sqrt(2)) in Kolafa-Perram gives rcut = error_factor * eta, kcut = error_factor / eta, and alpha = 1 / (sqrt(2) * sqrt(eta**2 - 2 * sigma**2)). The sigma -> 0 limit recovers the monopole formula.

Cost-ratio correction#

The textbook Kolafa-Perram balance assumes the per-real-space-pair cost equals the per-k-vector cost. On real hardware (and especially for the lmax=1 multipole tile kernels) those costs differ — measured C_r / C_k is in the 20-40x range for our cluster-pair tile kernels at fp64. The cost-balanced optimum scales as eta_eff = eta_KP / cost_ratio**(1/6): a 30x cost ratio shrinks rcut by ~1.76x (and grows kcut by the same factor), which can cut the real-space pair count by ~5x at the same target accuracy.

The math: with cost ratio R = C_r / C_k, the cost-balanced formula becomes eta_eff = (V^2 / (N * R))^(1/6) / sqrt(2 pi). Setting R = 1 (default) reproduces the canonical KP estimator.

param positions:

Atomic coordinates.

type positions:

torch.Tensor, shape (N, 3) or (N_total, 3)

param cell:

Unit cell matrix (matches the multipole-Ewald convention).

type cell:

torch.Tensor, shape (3, 3) or (B, 3, 3)

param sigma:

GTO basis width — same value used by the multipole-Ewald kernel. Scalar or shape (B,).

type sigma:

float or torch.Tensor

param batch_idx:

System index per atom. None selects single-system mode.

type batch_idx:

torch.Tensor, shape (N_total,), int32, optional

param accuracy:

Target relative accuracy (matches monopole convention).

type accuracy:

float, default 1e-6

param cost_ratio:

Empirical C_r / C_k ratio — per-pair real-space cost divided by per-k-vector reciprocal cost on the target hardware. 1.0 (default) reproduces canonical Kolafa-Perram. Higher values shift the optimum toward smaller rcut (fewer pairs) + larger kcut (more k-vectors), which wins when the per-pair cluster-pair tile kernel dominates. For the lmax=1 multipole kernels on a GB10-class GPU, cost_ratio ~ 30 is a reasonable starting point — measure on your own hardware via the per-pair / per-k timing probe (see docs/learnings/ if archived). Setting below 1 is allowed but rarely useful (the formula is symmetric).

type cost_ratio:

float, default 1.0

returns:

(alpha, sigma, real_space_cutoff, reciprocal_space_cutoff).

rtype:

MultipoleEwaldParameters

raises ValueError:

If any system has eta_eff <= sigma * sqrt(2) — meaning the cost-balanced Ewald split is degenerate at this size + sigma combination. Note: the validity threshold scales as cost_ratio**(-1/6) — large cost_ratio makes the split more likely to be invalid for small/dense systems.

Parameters:
Return type:

MultipoleEwaldParameters

nvalchemiops.torch.interactions.electrostatics.estimate_multipole_pme_parameters(positions, cell, sigma, batch_idx=None, accuracy=1e-6, cost_ratio=1.0)[source]#

Estimate GTO-Ewald multipole PME parameters at a given target accuracy.

Same Kolafa-Perram backbone as estimate_multipole_ewald_parameters(). Mesh dimensions follow the standard B-spline-error formula n_per_dim = 2 alpha_eff L / (3 accuracy**0.2) with alpha_eff = 1 / (sqrt(2) eta) — the monopole-equivalent alpha at the same eta.

The cost_ratio knob has the same meaning as in estimate_multipole_ewald_parameters() (per-pair vs per-k cost asymmetry), and shifts the rcut/kcut/mesh balance the same way: eta_eff = eta_KP / cost_ratio**(1/6). A larger cost_ratio grows the FFT mesh and shrinks the real-space cutoff. Note that PME’s true reciprocal cost is FFT (M log M) plus spread/gather (N p**3), which is not the same shape as the Ewald per-k-vector cost — so the optimal cost_ratio for PME may differ from the Ewald optimum even on the same hardware. Default 1.0 (canonical KP) is a safe starting point.

Returns:

(alpha, sigma, mesh_dimensions, mesh_spacing, real_space_cutoff).

Return type:

MultipolePMEParameters

Parameters:
class nvalchemiops.torch.interactions.electrostatics.MultipoleEwaldParameters(alpha, sigma, real_space_cutoff, reciprocal_space_cutoff)[source]#

Container for GTO-Ewald multipole parameters.

Like EwaldParameters but with the GTO basis width sigma propagated through. The Kolafa-Perram balance for the multipole case has the same rcut / kcut formulas as the monopole case, but alpha differs because the effective Ewald split width is sigma_c = sqrt(sigma**2 + 1/(4 alpha**2)) rather than 1/(alpha sqrt(2)).

Parameters:
alpha#

Ewald splitting parameter (inverse length units).

Type:

torch.Tensor, shape (B,)

sigma#

GTO basis width (passed through; physics).

Type:

torch.Tensor, shape (B,)

real_space_cutoff#

Real-space cutoff distance.

Type:

torch.Tensor, shape (B,)

reciprocal_space_cutoff#

Reciprocal-space cutoff (|k| in inverse length units).

Type:

torch.Tensor, shape (B,)

class nvalchemiops.torch.interactions.electrostatics.MultipolePMEParameters(alpha, sigma, mesh_dimensions, mesh_spacing, real_space_cutoff)[source]#

Container for GTO-Ewald multipole PME parameters.

Parameters:
alpha#

Ewald splitting parameter.

Type:

torch.Tensor, shape (B,)

sigma#

GTO basis width (passed through; physics).

Type:

torch.Tensor, shape (B,)

mesh_dimensions#

Mesh dimensions (nx, ny, nz) (max across batch).

Type:

tuple[int, int, int]

mesh_spacing#

Actual mesh spacing per direction.

Type:

torch.Tensor, shape (B, 3)

real_space_cutoff#

Real-space cutoff distance.

Type:

torch.Tensor, shape (B,)

On this page
Show Source