Ctrl+K

nvalchemiops.interactions.electrostatics: Electrostatic Interactions#

Electrostatics Interactions Module#

This module provides GPU-accelerated implementations of various methods for computing long-range electrostatic interactions in molecular simulations.

Available Methods#

  1. Ewald Summation (ewald) - Classical method splitting interactions into real-space and reciprocal-space - \(O(N^2)\) scaling for explicit k-vectors, good for small systems - Full autograd support

  2. Particle Mesh Ewald (PME) (pme) - FFT-based method for \(O(N \log N)\) scaling - Uses B-spline interpolation for charge assignment - Full autograd support

High-Level Interface#

These functions provide the end-user facing functions for computing electrostatic interactions using Ewald summation, Particle Mesh Ewald (PME), and direct Coulomb methods. All functions support automatic differentiation through PyTorch’s autograd system.

Tip

Check out the Electrostatic Interactions page for usage examples and a conceptual overview of the available electrostatic methods.

Ewald Summation#

Complete Ewald summation combining real-space and reciprocal-space contributions. Supports both single-system and batched calculations via the batch_idx parameter.

nvalchemiops.interactions.electrostatics.ewald.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, accuracy=1e-6)[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. Supports automatic parameter estimation, batched calculations, and automatic differentiation.

Formula#

The total Ewald energy is:

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}} = \frac{1}{2V} \sum_{k \in halfspace} G(k) \vert S(k) \vert^2 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\]
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) for single-system mode.

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 - None: Automatically estimated from accuracy using Kolafa-Perram formula Larger α shifts more computation to reciprocal space.

type alpha:

float, torch.Tensor, or None, default=None

param k_vectors:

Pre-computed reciprocal lattice vectors. Shape (K, 3) for single system, (B, K, 3) for batch. If None, generated from k_cutoff using generate_k_vectors_ewald_summation.

type k_vectors:

torch.Tensor, optional

param k_cutoff:

K-space cutoff (maximum |k| magnitude) for generating k_vectors. If None with alpha=None, estimated from accuracy. Typical values: 8-12 Å⁻¹.

type k_cutoff:

float, optional

param batch_idx:

System index for each atom (0 to B-1). Determines execution mode: - None: Single-system optimized kernels - Provided: Batched kernels for multiple independent systems

type batch_idx:

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

param neighbor_list:

Neighbor pairs in COO format. Row 0 = source indices, row 1 = target. Mutually exclusive with neighbor_matrix.

type neighbor_list:

torch.Tensor, shape (2, M), dtype=int32, optional

param neighbor_ptr:

CSR row pointers for neighbor list. neighbor_ptr[i] gives the starting index in idx_j for atom i’s neighbors. Required with neighbor_list. Provided by neighborlist module.

type neighbor_ptr:

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

param neighbor_shifts:

Periodic image shifts for each neighbor pair. Required with neighbor_list.

type neighbor_shifts:

torch.Tensor, shape (M, 3), dtype=int32, optional

param neighbor_matrix:

Dense neighbor matrix. 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.

type neighbor_matrix:

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

param neighbor_matrix_shifts:

Periodic image shifts for neighbor_matrix. Required with neighbor_matrix.

type neighbor_matrix_shifts:

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

param mask_value:

Value indicating invalid entries in neighbor_matrix. Defaults to N.

type mask_value:

int, optional

param compute_forces:

Whether to compute explicit analytical forces.

type compute_forces:

bool, default=False

param accuracy:

Target relative accuracy for automatic parameter estimation. Only used when alpha or k_cutoff is None. Smaller values increase accuracy but also computational cost.

type accuracy:

float, default=1e-6

returns:

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

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

raises ValueError:

If neither neighbor_list nor neighbor_matrix is provided.

raises TypeError:

If alpha has an unsupported type.

Examples

Automatic parameter estimation (recommended for most cases):

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

Explicit parameters for reproducibility:

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

Using neighbor matrix format:

>>> energies, forces = ewald_summation(
...     positions, charges, cell,
...     alpha=0.3, k_cutoff=8.0,
...     neighbor_matrix=nm, neighbor_matrix_shifts=nm_shifts,
...     mask_value=-1,
... )

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 = ewald_summation(
...     positions, charges, cells,  # cells shape (3, 3, 3)
...     alpha=torch.tensor([0.3, 0.35, 0.3]),
...     batch_idx=batch_idx,
...     k_cutoff=8.0,
...     neighbor_list=nl, neighbor_shifts=shifts,
... )

Energy-only (skips force computation for speed):

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

Using autograd for gradients:

>>> positions.requires_grad_(True)
>>> energies, forces = ewald_summation(
...     positions, charges, cell,
...     neighbor_list=nl, neighbor_shifts=shifts,
... )
>>> total_energy = energies.sum()
>>> total_energy.backward()
>>> autograd_forces = -positions.grad

Notes

Automatic Parameter Estimation (when alpha or k_cutoff is None):

Uses the Kolafa-Perram formula:

\[\begin{split}\begin{aligned} \eta &= \frac{(V^2 / N)^{1/6}}{\sqrt{2\pi}} \\ \alpha &= \frac{1}{\sqrt{2} \eta} \\ k_{\text{cutoff}} &= \frac{\sqrt{-2 \ln(\varepsilon)}}{\eta} \\ r_{\text{cutoff}} &= \sqrt{-2 \ln(\varepsilon)} \cdot \eta \end{aligned}\end{split}\]

This balances computational cost between real and reciprocal space.

Autograd Support:

All inputs (positions, charges, cell) support gradient computation. For positions, \(-\nabla E\) gives forces, which should match the explicit forces.

See also

ewald_real_space

Real-space component only

ewald_reciprocal_space

Reciprocal-space component only

estimate_ewald_parameters

Automatic parameter estimation

EwaldParameters

Container for Ewald parameters

Parameters:

  • positions (Tensor)

  • charges (Tensor)

  • cell (Tensor)

  • alpha (float | Tensor | None)

  • k_vectors (Tensor | None)

  • k_cutoff (float | None)

  • batch_idx (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)

  • accuracy (float)

Return type:

tuple[Tensor, Tensor] | Tensor

nvalchemiops.interactions.electrostatics.ewald.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=-1, batch_idx=None, compute_forces=False, compute_charge_gradients=False)[source]#

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

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.

Formula:

\[E_{\text{real}} = \frac{1}{2} \sum_{i \neq j} q_i q_j \frac{\text{erfc}(\alpha r_{ij})}{r_{ij}}\]
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).

  • alpha (torch.Tensor, shape (1,) or (B,)) – Ewald splitting parameter(s). Controls the real/reciprocal space split. Larger α shifts more computation to reciprocal space.

  • neighbor_list (torch.Tensor, shape (2, M), dtype=int32, optional) – Neighbor list in COO format. Row 0 contains source atom indices (i), row 1 contains target atom indices (j). Each pair should appear once (not symmetrized). 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 idx_j (neighbor_list[1]) for atom i’s neighbors. Required when using neighbor_list format. Provided by neighborlist module.

  • neighbor_shifts (torch.Tensor, shape (M, 3), dtype=int32, optional) – Periodic image shifts for each neighbor pair. Entry [k, :] gives the integer cell translation for pair k. Required with neighbor_list.

  • neighbor_matrix (torch.Tensor, shape (N, max_neighbors), dtype=int32, optional) – Dense neighbor matrix format. Entry [i, k] = j means atom j is the k-th neighbor of atom i. Invalid entries should be set to mask_value. More cache-friendly for small, fixed neighbor counts. 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, default=-1) – Value indicating invalid/padded entries in neighbor_matrix.

  • batch_idx (torch.Tensor, shape (N,), dtype=int32, optional) – System index for each atom (0 to B-1). Determines kernel dispatch: - None: Single-system optimized kernels - Provided: Batched kernels for multiple independent systems

  • compute_forces (bool, default=False) – Whether to compute explicit forces. Forces are computed analytically within the kernel (not via autograd).

  • compute_charge_gradients (bool, default=False) – Whether to compute analytical charge gradients (∂E/∂q_i). Useful for second-derivative training in ML potentials, as Warp requires analytical first derivatives to compute second derivatives via autograd.

Returns:

  • energies (torch.Tensor, shape (N,)) – Per-atom real-space energy contribution (sum gives total E_real).

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

  • charge_gradients (torch.Tensor, shape (N,), optional) – Gradient ∂E_real/∂q_i. Only returned if compute_charge_gradients=True.

  • Return Patterns

  • —————

  • - ``compute_forces=False, compute_charge_gradients=False`` (energies)

  • - ``compute_forces=True, compute_charge_gradients=False`` ((energies, forces))

  • - ``compute_forces=False, compute_charge_gradients=True`` ((energies, charge_gradients))

  • - ``compute_forces=True, compute_charge_gradients=True`` ((energies, forces, charge_gradients))

Raises:

ValueError – If neither neighbor_list nor neighbor_matrix is provided.

Return type:

Tensor | tuple[Tensor, …]

Examples

Energy only with neighbor list:

>>> energies = ewald_real_space(
...     positions, charges, cell, alpha,
...     neighbor_list=nl, neighbor_shifts=shifts,
... )
>>> total_energy = energies.sum()

With explicit forces:

>>> energies, forces = ewald_real_space(
...     positions, charges, cell, alpha,
...     neighbor_list=nl, neighbor_shifts=shifts,
...     compute_forces=True,
... )

With charge gradients for ML training:

>>> energies, forces, charge_grads = ewald_real_space(
...     positions, charges, cell, alpha,
...     neighbor_list=nl, neighbor_shifts=shifts,
...     compute_forces=True, compute_charge_gradients=True,
... )
>>> # charge_grads can be used to compute charge Hessian via autograd:

Using neighbor matrix format:

>>> energies = ewald_real_space(
...     positions, charges, cell, alpha,
...     neighbor_matrix=nm, neighbor_matrix_shifts=nm_shifts,
...     mask_value=-1,
... )

Batched systems:

>>> # positions: concatenated atoms, batch_idx: system assignment
>>> energies = ewald_real_space(
...     positions, charges, cells,  # cells shape (B, 3, 3)
...     alpha,  # shape (B,)
...     batch_idx=batch_idx,
...     neighbor_list=nl, neighbor_shifts=shifts,
... )

See also

ewald_reciprocal_space

Reciprocal-space component of Ewald summation.

ewald_summation

Complete Ewald summation (real + reciprocal).

estimate_ewald_parameters

Automatic parameter estimation.

nvalchemiops.interactions.electrostatics.ewald.ewald_reciprocal_space(positions, charges, cell, k_vectors, alpha, batch_idx=None, compute_forces=False, compute_charge_gradients=False)[source]#

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

Computes the smooth long-range electrostatic contribution using structure factors in reciprocal space. Automatically applies self-energy and background corrections.

The total energy is given by

\[E_{\text{reciprocal}} = \frac{1}{2V} \sum_{k \in halfspace} G(k) \vert S(k) \vert^2 - E_{\text{self}} - E_{\text{background}}\]

where the components are:

  • Green’s function: \(G(k) = \frac{8\pi}{k^2} \exp\left(-\frac{k^2}{4\alpha^2}\right)\)

  • Structure factor: \(S(k) = \sum_j q_j \exp(ik \cdot r_j)\)

  • Self-energy correction: \(E_{\text{self}} = \sum_i \frac{\alpha}{\sqrt{\pi}} q_i^2\)

  • Background correction: \(E_{\text{background}} = \frac{\pi}{2\alpha^2 V} Q_{\text{total}}^2\) (for non-neutral systems)

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

  • k_vectors (torch.Tensor) – Reciprocal lattice vectors from generate_k_vectors_ewald_summation. Shape (K, 3) for single system, (B, K, 3) for batch. Must be half-space vectors (excludes k=0 and -k for each +k).

  • alpha (torch.Tensor, shape (1,) or (B,)) – Ewald splitting parameter(s). Must match values used for real-space.

  • batch_idx (torch.Tensor, shape (N,), dtype=int32, optional) – System index for each atom (0 to B-1). Determines kernel dispatch: - None: Single-system optimized kernels - Provided: Batched kernels for multiple independent systems

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

  • compute_charge_gradients (bool, default=False) – Whether to compute analytical charge gradients (∂E/∂q_i). Useful for computing charge Hessians in ML potential training.

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 including corrections. Only returned if compute_charge_gradients=True.

  • Return Patterns

  • —————

  • - ``compute_forces=False, compute_charge_gradients=False`` (energies)

  • - ``compute_forces=True, compute_charge_gradients=False`` ((energies, forces))

  • - ``compute_forces=False, compute_charge_gradients=True`` ((energies, charge_gradients))

  • - ``compute_forces=True, compute_charge_gradients=True`` ((energies, forces, charge_gradients))

Return type:

Tensor | tuple[Tensor, …]

Examples

Generate k-vectors and compute energy:

>>> from nvalchemiops.interactions.electrostatics import (
...     generate_k_vectors_ewald_summation
... )
>>> k_vectors = generate_k_vectors_ewald_summation(cell, k_cutoff=8.0)
>>> energies = ewald_reciprocal_space(
...     positions, charges, cell, k_vectors, alpha,
... )
>>> total_recip_energy = energies.sum()

With forces:

>>> energies, forces = ewald_reciprocal_space(
...     positions, charges, cell, k_vectors, alpha,
...     compute_forces=True,
... )

With charge gradients for ML training:

>>> energies, charge_grads = ewald_reciprocal_space(
...     positions, charges, cell, k_vectors, alpha,
...     compute_charge_gradients=True,
... )

Batched systems:

>>> # k_vectors shape: (B, K, 3) with same K for all systems
>>> energies = ewald_reciprocal_space(
...     positions, charges, cells, k_vectors, alpha,
...     batch_idx=batch_idx,
... )

Notes

  • k_vectors MUST be generated using generate_k_vectors_ewald_summation, which provides half-space k-vectors. Using full k-space vectors will double-count and give incorrect energies.

  • For batch mode with varying cell sizes, use the same k_cutoff for all systems to ensure consistent K dimension.

  • The charge gradient formula includes corrections for self-energy and background, making it suitable for training on charge derivatives.

See also

ewald_real_space

Real-space component of Ewald summation.

ewald_summation

Complete Ewald summation (real + reciprocal).

generate_k_vectors_ewald_summation

Generate k-vectors for this function.

Particle Mesh Ewald (PME)#

FFT-based Ewald method achieving \(O(N \log N)\) scaling. Uses B-spline interpolation for efficient charge assignment and force interpolation.

nvalchemiops.interactions.electrostatics.pme.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, accuracy=1e-6)[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:

\[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|² 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).

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

  • Return Patterns

  • —————

  • - ``compute_forces=False, compute_charge_gradients=False`` (energies)

  • - ``compute_forces=True, compute_charge_gradients=False`` ((energies, forces))

  • - ``compute_forces=False, compute_charge_gradients=True`` ((energies, charge_gradients))

  • - ``compute_forces=True, compute_charge_gradients=True`` ((energies, forces, charge_gradients))

Raises:

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

  • TypeError – If alpha has an unsupported type.

Return type:

Tensor | tuple[Tensor, Tensor] | tuple[Tensor, Tensor, 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.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

nvalchemiops.interactions.electrostatics.pme.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)[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 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: .. math:

\varphi_{\text{mesh}}(k) = G(k) \times B^2(k) \times \rho_{\text{mesh}}(k)

where: .. math:

G(k) = (4\pi/k^2) \times exp(-k^2/(4\alpha^2))   Green's function
B(k) = B-spline structure factor   Interpolation correction
\[\rho_{\text{mesh}}(k) = 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|² 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

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.

  • Return Patterns

  • —————

  • - ``compute_forces=False, compute_charge_gradients=False`` (energies)

  • - ``compute_forces=True, compute_charge_gradients=False`` ((energies, forces))

  • - ``compute_forces=False, compute_charge_gradients=True`` ((energies, charge_gradients))

  • - ``compute_forces=True, compute_charge_gradients=True`` ((energies, forces, charge_gradients))

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

Return type:

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

Direct Coulomb#

Direct pairwise Coulomb interactions, supporting both undamped (1/r) and damped \((\text{erfc}(\alpha r)/r)\) variants. Useful for isolated systems or as the real-space component of Ewald/PME.

nvalchemiops.interactions.electrostatics.coulomb.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.interactions.electrostatics.coulomb.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.interactions.electrostatics.coulomb.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]

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

Parameter Estimation#

Functions for automatic parameter estimation based on desired accuracy tolerance.

Ewald Parameters#

nvalchemiops.interactions.electrostatics.parameters.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). Common values: 1e-4 (low), 1e-6 (medium), 1e-8 (high).

Returns:

Dataclass containing alpha, real_space_cutoff, reciprocal_space_cutoff, and eta. For batch mode, these are tensors of shape (B,). For single-system, they are floats.

Return type:

EwaldParameters

Examples

>>> positions = torch.randn(100, 3)
>>> cell = torch.eye(3) * 20.0
>>> params = estimate_ewald_parameters(positions, cell, accuracy=1e-6)
>>> print(f"alpha={params.alpha:.4f}, r_cut={params.real_space_cutoff:.2f}")
alpha=0.2835, r_cut=7.42

Notes

The formulas are:

\[\begin{split}\begin{aligned} \eta &= \frac{(V^2 / N)^{1/6}}{\sqrt{2\pi}} \\ \alpha &= \frac{1}{\sqrt{2}\eta} \\ r_{\text{cutoff}} &= \sqrt{-2 \ln(\varepsilon)} \cdot \eta \\ k_{\text{cutoff}} &= \frac{\sqrt{-2 \ln(\varepsilon)}}{\eta} \end{aligned}\end{split}\]

See also

EwaldParameters

Container for the returned parameters

estimate_pme_parameters

Estimate PME parameters (includes mesh sizing)

class nvalchemiops.interactions.electrostatics.parameters.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 (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,)

See also

estimate_ewald_parameters

Estimate optimal Ewald parameters automatically

PMEParameters

Container for PME parameters

alpha: Tensor#
real_space_cutoff: Tensor#
reciprocal_space_cutoff: Tensor#

PME Parameters#

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

Estimate optimal PME parameters for a given accuracy.

This combines Ewald parameter estimation with PME mesh sizing.

Parameters:

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

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrix (row vectors are lattice vectors).

  • 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: torch.Tensor, shape (B,) - Ewald splitting parameter per system - mesh_dimensions: tuple[int, int, int] - Maximum mesh dimensions across batch - mesh_spacing: torch.Tensor, shape (B, 3) - Actual mesh spacing per system - real_space_cutoff: torch.Tensor, shape (B,) - Real-space cutoff per system

Return type:

PMEParameters

Examples

>>> positions = torch.randn(100, 3)
>>> cell = torch.eye(3) * 20.0
>>> params = estimate_pme_parameters(positions, cell, accuracy=1e-6)
>>> print(f"alpha={params.alpha.item():.4f}, mesh={params.mesh_dimensions}")
alpha=0.2835, mesh=(32, 32, 32)

See also

PMEParameters

Container for the returned parameters

estimate_ewald_parameters

Estimate Ewald parameters only

estimate_pme_mesh_dimensions

Estimate mesh dimensions only

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

Estimate optimal PME mesh dimensions for a given accuracy.

The formula is based on the B-spline interpolation error analysis:

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

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrix (row vectors are lattice vectors).

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

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

Returns:

Maximum mesh dimensions (nx, ny, nz) across all systems in batch. Dimensions are rounded up to powers of 2 for FFT efficiency.

Return type:

tuple[int, int, int]

Examples

>>> cell = 20.0 * torch.eye(3).unsqueeze(0)
>>> alpha = torch.tensor([0.3])
>>> mesh_dims = estimate_pme_mesh_dimensions(cell, alpha, accuracy=1e-6)
>>> print(mesh_dims)
(64, 64, 64)

See also

estimate_pme_parameters

Full PME parameter estimation

mesh_spacing_to_dimensions

Convert mesh spacing to dimensions

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

Convert mesh spacing to mesh dimensions.

Parameters:

  • cell (torch.Tensor, shape (3, 3) or (B, 3, 3)) – Unit cell matrix (row vectors are lattice vectors).

  • mesh_spacing (float | torch.Tensor) – Target mesh spacing. Can be: - float: uniform spacing for all directions and systems - torch.Tensor, shape (B,): per-system spacing (uniform in all directions) - torch.Tensor, shape (B, 3): per-system, per-direction spacing

Returns:

Mesh dimensions, rounded up to powers of 2.

Return type:

tuple[int, int, int]

See also

estimate_pme_mesh_dimensions

Estimate mesh dimensions from accuracy tolerance

PMEParameters

Container that includes mesh dimensions

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

Container for PME parameters.

For single-system calculations, alpha and real_space_cutoff are tensors of shape (1,). For batch calculations, alpha and real_space_cutoff are tensors of shape (B,). Mesh spacing is a tensor of shape (B, 3). mesh_dimensions is a list of 3 integers, which is calculated as the maximum

mesh dimensions along each axis for the entire set of structures in the batch.

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

See also

estimate_pme_parameters

Estimate optimal PME parameters automatically

EwaldParameters

Container for Ewald parameters

alpha: Tensor#
mesh_dimensions: tuple[int, int, int]#
mesh_spacing: Tensor#
real_space_cutoff: Tensor#

Utility Functions#

K-Vector Generation#

Functions for generating reciprocal-space vectors for Ewald summation and PME.

nvalchemiops.interactions.electrostatics.k_vectors.generate_k_vectors_ewald_summation(cell, k_cutoff)[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

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

  • 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.interactions.electrostatics.k_vectors.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.

B-Spline Functions#

Functions for B-spline charge spreading and gathering, used by PME.

nvalchemiops.spline.spline_spread(positions, values, cell, mesh_dims, spline_order=4, batch_idx=None, cell_inv_t=None)[source]#

Spread values from atoms to mesh grid using B-spline interpolation.

Parameters:

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

  • values (torch.Tensor, shape (N,)) – Values to spread (e.g., charges).

  • cell (torch.Tensor, shape (3, 3), (1, 3, 3), or (B, 3, 3)) – Unit cell matrix. For batched, shape should be (B, 3, 3).

  • mesh_dims (tuple[int, int, int]) – Mesh dimensions (nx, ny, nz).

  • spline_order (int, default=4) – B-spline order (1-4, where 4=cubic).

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

  • cell_inv_t (torch.Tensor | None, default=None) – Precomputed transpose of cell inverse. If provided, skips inverse computation. Shape (1, 3, 3) for single-system or (B, 3, 3) for batch.

Returns:

mesh – For single-system: shape (nx, ny, nz) For batch: shape (B, nx, ny, nz)

Return type:

torch.Tensor

nvalchemiops.spline.spline_gather(positions, mesh, cell, spline_order=4, batch_idx=None, cell_inv_t=None)[source]#

Gather values from mesh to atoms using B-spline interpolation.

Parameters:

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

  • mesh (torch.Tensor) – For single-system: shape (nx, ny, nz) For batch: shape (B, nx, ny, nz)

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

  • spline_order (int, default=4) – B-spline order.

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

  • cell_inv_t (torch.Tensor | None, default=None) – Precomputed transpose of cell inverse. If provided, skips inverse computation. Shape (1, 3, 3) for single-system or (B, 3, 3) for batch.

Returns:

values – Interpolated values at atomic positions.

Return type:

torch.Tensor, shape (N,)

nvalchemiops.spline.spline_gather_vec3(positions, charges, mesh, cell, spline_order=4, batch_idx=None, cell_inv_t=None)[source]#

Gather 3D vector values from mesh to atoms using B-spline interpolation.

This is useful for interpolating vector fields like electric fields.

Parameters:

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

  • mesh (torch.Tensor) – For single-system: shape (nx, ny, nz, 3) For batch: shape (B, nx, ny, nz, 3)

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

  • spline_order (int, default=4) – B-spline order.

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

  • cell_inv_t (torch.Tensor | None, default=None) – Precomputed transpose of cell inverse. If provided, skips inverse computation. Shape (1, 3, 3) for single-system or (B, 3, 3) for batch.

  • charges (Tensor)

Returns:

vectors – Interpolated 3D vectors at atomic positions.

Return type:

torch.Tensor, shape (N, 3)

nvalchemiops.spline.spline_gather_gradient(positions, charges, mesh, cell, spline_order=4, batch_idx=None, cell_inv_t=None)[source]#

Gather gradient from mesh to atoms using B-spline derivatives.

Computes forces:

\[F_i = -q_i \sum_g \phi(g) \nabla w(r_i, g)\]
Parameters:

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

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

  • mesh (torch.Tensor) – For single-system: shape (nx, ny, nz) For batch: shape (B, nx, ny, nz)

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

  • spline_order (int, default=4) – B-spline order.

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

  • cell_inv_t (torch.Tensor | None, default=None) – Precomputed transpose of cell inverse. If provided, skips inverse computation. Shape (1, 3, 3) for single-system or (B, 3, 3) for batch.

Returns:

forces – Forces on atoms.

Return type:

torch.Tensor, shape (N, 3)

On this page
Show Source