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.

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)

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.

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

• 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) – Whether to compute explicit forces.

• compute_charge_gradients (bool, default=False) – Whether to compute charge gradients dE/dq_i.

• compute_virial (bool, default=False) – Whether to compute the virial tensor W = -dE/d(epsilon). Stress = virial / volume.

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

Returns:

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

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

• charge_gradients (torch.Tensor, shape (N,), optional) – 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 always float64 for numerical stability during accumulation. Forces, charge gradients, and virial match the input dtype (float32 or float64).

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)

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:

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

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 α for all systems - Tensor shape (B,): Per-system α values - None: Automatically estimated using Kolafa-Perram formula Larger α 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 Å.

• 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

• 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).

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

• 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) – Whether to compute explicit analytical forces.

• compute_charge_gradients (bool, default=False) – Whether to compute analytical charge gradients ∂E/∂q_i. Useful for training ML potentials that require second derivatives (charge Hessians).

• compute_virial (bool, default=False) – Whether to compute the virial tensor W = -dE/d(epsilon). Stress = virial / volume.

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

Returns:

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

• forces (torch.Tensor, shape (N, 3), optional) – Forces on each atom. Only returned if compute_forces=True.

• charge_gradients (torch.Tensor, shape (N,), optional) – Charge gradients ∂E/∂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

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

Enabled flags are appended in order: energies, [forces], [charge_gradients], [virial]. A single output is returned unwrapped; multiple outputs 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:

• positions (Tensor)

• charges (Tensor)

• cell (Tensor)

• alpha (float | Tensor | None)

• mesh_spacing (float | None)

• mesh_dimensions (tuple[int, int, int] | None)

• spline_order (int)

• batch_idx (Tensor | None)

• k_vectors (Tensor | None)

• k_squared (Tensor | None)

• neighbor_list (Tensor | None)

• neighbor_ptr (Tensor | None)

• neighbor_shifts (Tensor | None)

• neighbor_matrix (Tensor | None)

• neighbor_matrix_shifts (Tensor | None)

• mask_value (int | None)

• compute_forces (bool)

• compute_charge_gradients (bool)

• compute_virial (bool)

• accuracy (float)

Return type:

Tensor | tuple[Tensor, …]

Examples

Automatic parameter estimation (recommended for most cases):

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

Explicit parameters for reproducibility:

>>> energies, forces = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     spline_order=4, neighbor_list=nl,
... neighbor_shifts=shifts, neighbor_ptr=nptr,
...     compute_forces=True,
... )

Using mesh spacing for automatic mesh sizing:

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

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 × 3 systems
>>> energies, forces = 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_shifts=shifts, neighbor_ptr=nptr,
...     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, forces = 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_shifts=shifts,
...         neighbor_ptr=nptr,
...         compute_forces=True,
...     )

With charge gradients for ML training:

>>> energies, forces, charge_grads = particle_mesh_ewald(
...     positions, charges, cell,
...     alpha=0.3, mesh_dimensions=(32, 32, 32),
...     neighbor_list=nl, neighbor_shifts=shifts,
...     neighbor_ptr=nptr,
...     compute_forces=True, compute_charge_gradients=True,
... )
>>> # Use charge_grads for training on ∂E/∂q

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_shifts=shifts,
...     neighbor_ptr=nptr,
... )
>>> 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

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)

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)

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)

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)

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)

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.

• compute_forces (bool, default=False) – Whether to compute explicit forces.

• compute_charge_gradients (bool, default=False) – Whether to compute charge gradients.

• compute_virial (bool, default=False) – Whether to compute the virial tensor W = -dE/d(epsilon). Stress = virial / volume.

Returns:

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

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

• charge_gradients (torch.Tensor, shape (N,), optional) – 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 and virial match the input dtype (float32 or float64).

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)

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.

• compute_forces (bool, default=False) – Whether to compute explicit forces.

• compute_charge_gradients (bool, default=False) – Whether to compute charge gradients.

• compute_virial (bool, default=False) – Whether to compute the virial tensor W = -dE/d(epsilon). Stress = virial / volume.

Returns:

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

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

• charge_gradients (torch.Tensor, shape (N,), optional) – 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 and virial match the input dtype (float32 or float64).

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)

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 Green’s function (multiply by G(k))

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) = G(k) \times B^2(k) \times \rho_{\text{mesh}}(k)\]

where:

• \(G(k) = (4\pi/k^2) \times \exp(-k^2/(4\alpha^2))\) is the Green’s function

• \(B(k)\) is the B-spline structure factor (interpolation correction)

• \(\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 α for all systems - Tensor shape (B,): Per-system α 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 Å.

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

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.

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.

type k_squared:

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

param compute_forces:

Whether to compute explicit reciprocal-space forces.

type compute_forces:

bool, default=False

param compute_charge_gradients:

Whether to compute analytical charge gradients ∂E/∂q_i. Useful for computing charge Hessians in ML potential training.

type compute_charge_gradients:

bool, default=False

param compute_virial:

Whether to compute the virial tensor W = -dE/d(epsilon). Stress = virial / volume.

type compute_virial:

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) – Reciprocal-space forces. Only returned if compute_forces=True.

• charge_gradients (torch.Tensor, shape (N,), optional) – Charge gradients ∂E_recip/∂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

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

Enabled flags are appended in order: energies, [forces], [charge_gradients], [virial]. A single output is returned unwrapped; multiple outputs 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.

pme_green_structure_factor

Compute Green’s function on mesh.

Parameters:

• positions (Tensor)

• charges (Tensor)

• cell (Tensor)

• alpha (float | Tensor)

• mesh_dimensions (tuple[int, int, int] | None)

• mesh_spacing (float | None)

• spline_order (int)

• batch_idx (Tensor | None)

• k_vectors (Tensor | None)

• k_squared (Tensor | None)

• compute_forces (bool)

• compute_charge_gradients (bool)

• compute_virial (bool)

Return type:

Tensor | tuple[Tensor, …]

nvalchemiops.torch.interactions.electrostatics.pme_green_structure_factor(k_squared, mesh_dimensions, alpha, cell, spline_order=4, batch_idx=None)

Compute Green’s function and B-spline structure factor correction.

Computes the Coulomb Green’s function with volume normalization and the B-spline aliasing correction factor for PME.

Green’s function (volume-normalized):

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

Structure factor correction (for B-spline deconvolution):

\[C^2(k) = \left[\text{sinc}(m_x/N_x) \cdot \text{sinc}(m_y/N_y) \cdot \text{sinc}(m_z/N_z)\right]^{2p}\]

where p is the spline order.

Supports both float32 and float64 dtypes.

Parameters:

• k_squared (torch.Tensor) – \(|k|^2\) values at each FFT grid point. - Single-system: shape (Nx, Ny, Nz_rfft) - Batch: shape (B, Nx, Ny, Nz_rfft)

• mesh_dimensions (tuple[int, int, int]) – Full mesh dimensions (Nx, Ny, Nz) before rfft.

• alpha (torch.Tensor) – Ewald splitting parameter. - Single-system: shape (1,) - Batch: shape (B,)

• cell (torch.Tensor) – Unit cell matrices. - Single-system: shape (3, 3) or (1, 3, 3) - Batch: shape (B, 3, 3)

• spline_order (int, default=4) – B-spline interpolation order (typically 4 for cubic B-splines).

• batch_idx (torch.Tensor | None, default=None) – If provided, dispatches to batch kernels.

Returns:

• green_function (torch.Tensor) – Volume-normalized Green’s function \(G(k)\). - Single-system: shape (Nx, Ny, Nz_rfft) - Batch: shape (B, Nx, Ny, Nz_rfft)

• structure_factor_sq (torch.Tensor) – Squared structure factor \(C^2(k)\) for B-spline deconvolution. Shape (Nx, Ny, Nz_rfft), shared across batch.

Return type:

tuple[Tensor, Tensor]

Notes

• \(G(k=0)\) is set to zero to avoid singularity

• The volume normalization in \(G(k)\) eliminates later divisions

• Structure factor is mesh-dependent only, so shared across batch

nvalchemiops.torch.interactions.electrostatics.pme_energy_corrections(raw_energies, charges, cell, alpha, batch_idx=None)

Apply self-energy and background corrections to PME energies.

Converts raw interpolated potential to energy and subtracts corrections:

\[E_i = q_i \phi_i - E_{\text{self},i} - E_{\text{background},i}\]

Self-energy correction (removes Gaussian self-interaction):

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

Background correction (for non-neutral systems):

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

Parameters:

• raw_energies (torch.Tensor, shape (N,) or (N_total,)) – Raw potential values \(\phi_i\) from mesh interpolation.

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

• cell (torch.Tensor) – Unit cell matrices. - Single-system: shape (3, 3) or (1, 3, 3) - Batch: shape (B, 3, 3)

• alpha (torch.Tensor) – Ewald splitting parameter. - Single-system: shape (1,) - Batch: shape (B,)

• batch_idx (torch.Tensor | None, default=None) – System index for each atom. If provided, uses batch kernels.

Returns:

corrected_energies – Final per-atom reciprocal-space energy with corrections applied.

Return type:

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

Notes

• For neutral systems, background correction is zero

• Matches torchpme’s self_contribution and background_correction formulas

• Supports both float32 and float64 dtypes

nvalchemiops.torch.interactions.electrostatics.pme_energy_corrections_with_charge_grad(raw_energies, charges, cell, alpha, batch_idx=None)

Apply corrections and compute charge gradients for PME energies.

Computes both corrected energies and analytical charge gradients:

E_i = q_i * φ_i - E_self_i - E_background_i ∂E/∂q_i = 2*φ_i - 2*(α/√π)*q_i - (π/(α²V))*Q_total

The factor of 2 on φ_i arises because changing q_i affects both the direct energy term (q_i * φ_i) and all other potentials through the structure factor (∑_j q_j * ∂φ_j/∂q_i = φ_i).

Parameters:

• raw_energies (torch.Tensor, shape (N,) or (N_total,)) – Raw potential values φ_i from mesh interpolation.

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

• cell (torch.Tensor) – Unit cell matrices. - Single-system: shape (3, 3) or (1, 3, 3) - Batch: shape (B, 3, 3)

• alpha (torch.Tensor) – Ewald splitting parameter. - Single-system: shape (1,) - Batch: shape (B,)

• batch_idx (torch.Tensor | None, default=None) – System index for each atom. If provided, uses batch kernels.

Returns:

• corrected_energies (torch.Tensor, shape (N,) or (N_total,)) – Final per-atom reciprocal-space energy with corrections applied.

• charge_gradients (torch.Tensor, shape (N,) or (N_total,)) – Analytical charge gradients ∂E/∂q_i.

Return type:

tuple[Tensor, Tensor]

K-Vector Generation

nvalchemiops.torch.interactions.electrostatics.generate_k_vectors_ewald_summation(cell, k_cutoff)

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

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³ · 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:

• cell (Tensor)

• k_cutoff (float | Tensor)

Return type:

Tensor

nvalchemiops.torch.interactions.electrostatics.generate_k_vectors_pme(cell, mesh_dimensions, reciprocal_cell=None)

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)

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)

Estimate optimal PME parameters for a given 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.

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

Returns:

Dataclass containing alpha, mesh dimensions, spacing, and cutoffs. Tensor fields are torch.Tensor objects.

Return type:

PMEParameters

nvalchemiops.torch.interactions.electrostatics.estimate_pme_mesh_dimensions(cell, alpha, accuracy=1e-6)

Estimate optimal PME mesh dimensions for a given accuracy.

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 accuracy.

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)

Convert mesh spacing to mesh dimensions.

Parameters:

• cell (torch.Tensor) – Unit cell matrix.

• mesh_spacing (float | torch.Tensor) – Target mesh spacing.

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)

Container for Ewald summation parameters.

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

Parameters:

• alpha (Tensor)

• real_space_cutoff (Tensor)

• reciprocal_space_cutoff (Tensor)

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)

Container for PME parameters.

Parameters:

• alpha (Tensor)

• mesh_dimensions (tuple[int, int, int])

• mesh_spacing (Tensor)

• real_space_cutoff (Tensor)

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,)