`nvalchemiops.jax.interactions.electrostatics`: Electrostatics#

The electrostatics module provides GPU-accelerated implementations of long-range electrostatic interactions for molecular simulations with JAX bindings. These functions accept standard jax.Array inputs.

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. They are compatible with jax.jit when setup-only PME parameters such as mesh_dimensions and alpha are supplied explicitly whenever those values would otherwise be estimated from traced inputs. miller_bounds is also a static shape control: under jax.jit, pass it as a concrete tuple or build k_vectors outside the compiled function. Energy derivatives are defined for positions, charges, and cell. Setup values such as alpha and mesh controls are constants; precomputed reciprocal metadata such as k_vectors, k_squared, volume, and cell_inv_t is accepted for cell-differentiated calls as static metadata that is assumed to correspond to the current cell; cache-generation derivatives are not recovered. Energy-returning Ewald, PME, and slab paths support atom-weighted losses such as (weights * energies).sum() for positions, charges, and supported cell derivatives. Monopole entry points provide the keyword-only energy_reduction option. The default, "atom", returns one energy per atom with shape (N,). Set it to "system" to sum energies within each system and return shape (B,). Other requested outputs keep their existing shapes. JAX PME supports first-order cell/strain gradients, but PME cell/strain HVPs, including full PME with slab_correction=True, are explicitly unsupported until a native transposable PME cell-HVP path is implemented and tested. Point-charge Ewald/PME inputs support float32 and float64. Keep all floating inputs and precomputed metadata in a call on a consistent dtype.

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

Compute complete Ewald summation.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic partial charges.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices.
alpha (float, jax.Array, or None, default=None) – Ewald splitting parameter. If None, estimated automatically.
k_vectors (jax.Array or None, default=None) – Reciprocal lattice vectors. Generated from cell when omitted.
k_cutoff (float, jax.Array, or None, default=None) – K-space cutoff used when generating k_vectors.
miller_bounds (tuple[int, int, int] or None, default=None, keyword-only) – Static Miller-index bounds for JIT-compatible k-vector generation.
batch_idx (jax.Array or None, default=None) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.
max_atoms_per_system (int or None, default=None) – Static batch shape control for reciprocal kernels under jax.jit.
neighbor_list (jax.Array or None) – CSR neighbor-list inputs for the real-space component.
neighbor_ptr (jax.Array or None) – CSR neighbor-list inputs for the real-space component.
neighbor_shifts (jax.Array or None) – CSR neighbor-list inputs for the real-space component.
neighbor_matrix (jax.Array or None) – Dense neighbor-matrix inputs for the real-space component.
neighbor_matrix_shifts (jax.Array or None) – Dense neighbor-matrix inputs for the real-space component.
mask_value (int or None, default=None) – Sentinel value for invalid neighbor-matrix entries.
compute_forces (bool, default=False) –

Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use JAX autodiff for differentiable forces.
compute_charge_gradients (bool, default=False) –

Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use JAX autodiff for \(\partial E / \partial q_i\).
compute_virial (bool, default=False) –

Deprecated since version 0.4.0: Deprecated direct-output flag for the virial tensor.
accuracy (float, default=1e-6) – Target accuracy for automatic parameter estimation.
hybrid_forces (bool, default=False) – Deprecated direct-output flag retained for transition compatibility.
pbc (jax.Array, optional) – Per-system periodic boundary conditions for slab correction.
slab_correction (bool, default=False) – If True, add the Yeh-Berkowitz/Ballenegger slab correction.
energy_reduction ({"atom", "system"}, keyword-only, default="atom") – Public energy layout. "atom" returns per-atom energies of shape (N,) (default, unchanged behavior). "system" returns per-system energies of shape (B,) (or (1,) for a single system), computed as a differentiable, uniform scatter-sum of the per-atom energies; it does not support nonuniform per-atom weighting. Only the energy field of the return value changes; forces, charge gradients, and the virial keep their existing atom/system layouts.

Returns:

jax.Array, shape (N,) or (B,) – Total Ewald energy when no deprecated direct-output flags are set: per-atom when energy_reduction="atom", per-system when energy_reduction="system". Gradients flow through positions, charges, and cell via the registered custom-JVP rules.
tuple[jax.Array, …] – When any deprecated flag is True: (energies,) extended by the requested outputs in order — forces of shape (N, 3), charge gradients of shape (N,), virial of shape (1, 3, 3) or (B, 3, 3) — matching the ordering of nvalchemiops.jax.interactions.electrostatics.ewald.ewald_real_space().

Return type:

Array | tuple[Array, …]

See also

nvalchemiops.jax.interactions.electrostatics.ewald.ewald_real_space(): Real-space component.
nvalchemiops.jax.interactions.electrostatics.ewald.ewald_reciprocal_space(): Reciprocal-space component.
nvalchemiops.jax.interactions.electrostatics.parameters.estimate_ewald_parameters(): Automatic alpha and k-cutoff estimation.
nvalchemiops.jax.interactions.electrostatics.k_vectors.generate_k_vectors_ewald_summation(): Generates k-vectors from cell and cutoff.

nvalchemiops.jax.interactions.electrostatics.particle_mesh_ewald(positions, charges, cell, alpha=None, mesh_spacing=None, mesh_dimensions=None, spline_order=4, batch_idx=None, k_vectors=None, k_squared=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, mask_value=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, accuracy=1e-6, hybrid_forces=False, pbc=None, slab_correction=False, *, energy_reduction='atom', cell_inv_t=None, volume=None, moduli_x=None, moduli_y=None, moduli_z=None)[source]#

Complete Particle Mesh Ewald calculation for long-range electrostatics.

Computes the total Coulomb energy via the PME method, which achieves \(O(N \log N)\) scaling through FFT-based reciprocal-space calculations:

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

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic partial charges.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices with lattice vectors as rows.
alpha (float, jax.Array, or None, default=None) – Ewald splitting parameter. If None, estimated automatically.
mesh_spacing (float or None, default=None) – Target mesh spacing used when mesh_dimensions is omitted.
mesh_dimensions (tuple[int, int, int] or None, default=None) – Explicit FFT mesh dimensions.
spline_order (int, default=4) – B-spline interpolation order.
batch_idx (jax.Array or None, default=None) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.
k_vectors (jax.Array or None) – Precomputed PME reciprocal grid values.
k_squared (jax.Array or None) – Precomputed PME reciprocal grid values.
neighbor_list (jax.Array or None) – CSR neighbor-list inputs for the real-space component.
neighbor_ptr (jax.Array or None) – CSR neighbor-list inputs for the real-space component.
neighbor_shifts (jax.Array or None) – CSR neighbor-list inputs for the real-space component.
neighbor_matrix (jax.Array or None) – Dense neighbor-matrix inputs for the real-space component.
neighbor_matrix_shifts (jax.Array or None) – Dense neighbor-matrix inputs for the real-space component.
mask_value (int or None, default=None) – Sentinel value for invalid neighbor-matrix entries.
compute_forces (bool) – Deprecated direct-output flags. Compute energy and use JAX autodiff for differentiable forces, charge gradients, and strain virials.
compute_charge_gradients (bool) – Deprecated direct-output flags. Compute energy and use JAX autodiff for differentiable forces, charge gradients, and strain virials.
compute_virial (bool) – Deprecated direct-output flags. Compute energy and use JAX autodiff for differentiable forces, charge gradients, and strain virials.
accuracy (float, default=1e-6) – Target accuracy for automatic parameter estimation.
hybrid_forces (bool, default=False) – Deprecated Torch-compatibility escape hatch for charge-gradient routing.
pbc (jax.Array, optional) – Per-system periodic boundary conditions for slab correction.
slab_correction (bool, default=False) – If True, add the Yeh-Berkowitz/Ballenegger slab correction.
energy_reduction ({"atom", "system"}, keyword-only, default="atom") – Public energy layout. "atom" returns per-atom energies of shape (N,) (default, unchanged behavior). "system" returns per-system energies of shape (B,) (or (1,) for a single system), computed as a differentiable, uniform scatter-sum of the per-atom energies; it does not support nonuniform per-atom weighting. Only the energy field of the return value changes; forces, charge gradients, and the virial keep their existing atom/system layouts.
volume (jax.Array or None) – Optional precomputed PME intermediates. Cell-derived values supplied while differentiating with respect to cell are treated as static metadata that corresponds to the current cell.
cell_inv_t (jax.Array or None) – Optional precomputed PME intermediates. Cell-derived values supplied while differentiating with respect to cell are treated as static metadata that corresponds to the current cell.
moduli_x (jax.Array or None) – Optional precomputed PME intermediates. Cell-derived values supplied while differentiating with respect to cell are treated as static metadata that corresponds to the current cell.
moduli_y (jax.Array or None) – Optional precomputed PME intermediates. Cell-derived values supplied while differentiating with respect to cell are treated as static metadata that corresponds to the current cell.
moduli_z (jax.Array or None) – Optional precomputed PME intermediates. Cell-derived values supplied while differentiating with respect to cell are treated as static metadata that corresponds to the current cell.

Returns:

energies (jax.Array, shape (N,) or (B,)) – Total electrostatic energies (real + reciprocal + slab): per-atom when energy_reduction="atom", per-system when energy_reduction="system".
forces (jax.Array, shape (N, 3), optional) – Per-atom forces. Only present when compute_forces=True (deprecated).
charge_gradients (jax.Array, shape (N,), optional) – Per-atom charge gradients \(\partial E/\partial q\). Only present when compute_charge_gradients=True (deprecated).
virial (jax.Array, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor. Only present when compute_virial=True (deprecated). Always last in the return tuple.

Return type:

Array | tuple[Array, Array] | tuple[Array, Array, Array] | tuple[Array, Array, Array, Array]

Notes

When cell, alpha, or batch metadata are traced by jax.jit or other JAX transformations, pass explicit mesh_dimensions. mesh_spacing and accuracy-based mesh sizing depend on concrete mesh setup values. If alpha would otherwise be estimated from traced inputs, precompute it outside the transformation and pass it explicitly.

Coulomb Interactions#

Direct pairwise Coulomb interactions.

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

Compute Coulomb electrostatic energies.

Computes pairwise electrostatic energies using the Coulomb law, with optional erfc damping for Ewald/PME real-space calculations.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic charges.
cell (jax.Array, shape (3, 3), (1, 3, 3), or (B, 3, 3)) – Unit cell matrix. A single-system (3, 3) matrix is promoted to (1, 3, 3) internally.
cutoff (float) – Cutoff distance for interactions.
alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.
neighbor_list (jax.Array | None, shape (2, num_pairs)) – Neighbor pairs in COO format. Row 0 = source, Row 1 = target.
neighbor_ptr (jax.Array | None, shape (N+1,)) – CSR row pointers for neighbor list. Required with neighbor_list. Provided by neighborlist module.
neighbor_shifts (jax.Array | None, shape (num_pairs, 3)) – Integer unit cell shifts for neighbor list format.
neighbor_matrix (jax.Array | None, shape (N, max_neighbors)) – Neighbor indices in matrix format.
neighbor_matrix_shifts (jax.Array | None, shape (N, max_neighbors, 3)) – Integer unit cell shifts for matrix format.
fill_value (int | None) – Fill value for neighbor matrix padding. Applies only to matrix format.
batch_idx (jax.Array | None, shape (N,)) – Batch indices for each atom.

Return type:

Array

Notes

Callers must supply one complete supported neighbor topology route: CSR/COO neighbor_list, neighbor_ptr, and neighbor_shifts; or matrix neighbor_matrix and neighbor_matrix_shifts. Current validation raises ValueError when no complete route is provided or when both complete routes are supplied simultaneously; it does not reject stray fields from the other representation. Shift arrays are required in both formats; for non-periodic systems pass integer all-zero shifts.

Returns:

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

Return type:

jax.Array, shape (N,)

Parameters:

positions (Array)
charges (Array)
cell (Array)
cutoff (float)
alpha (float)
neighbor_list (Array | None)
neighbor_ptr (Array | None)
neighbor_shifts (Array | None)
neighbor_matrix (Array | None)
neighbor_matrix_shifts (Array | None)
fill_value (int | None)
batch_idx (Array | None)

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)
>>> energies = coulomb_energy(
...     positions, charges, cell, cutoff=10.0, alpha=0.3,
...     neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr,
...     neighbor_shifts=neighbor_shifts
... )

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

Compute Coulomb electrostatic forces.

Convenience wrapper that returns only forces (no energies).

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic charges.
cell (jax.Array, shape (3, 3), (1, 3, 3), or (B, 3, 3)) – Unit cell matrix. A single-system (3, 3) matrix is promoted to (1, 3, 3) internally.
cutoff (float) – Cutoff distance for interactions.
alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.
neighbor_list (jax.Array | None, shape (2, num_pairs)) – Neighbor pairs in COO format. Row 0 = source, Row 1 = target.
neighbor_ptr (jax.Array | None, shape (N+1,)) – CSR row pointers for neighbor list. Required with neighbor_list. Provided by neighborlist module.
neighbor_shifts (jax.Array | None, shape (num_pairs, 3)) – Integer unit cell shifts for neighbor list format.
neighbor_matrix (jax.Array | None, shape (N, max_neighbors)) – Neighbor indices in matrix format.
neighbor_matrix_shifts (jax.Array | None, shape (N, max_neighbors, 3)) – Integer unit cell shifts for matrix format.
fill_value (int | None) – Fill value for neighbor matrix padding. Applies only to matrix format.
batch_idx (jax.Array | None, shape (N,)) – Batch indices for each atom.

Return type:

Array

Notes

Callers must supply one complete supported neighbor topology route: CSR/COO neighbor_list, neighbor_ptr, and neighbor_shifts; or matrix neighbor_matrix and neighbor_matrix_shifts. Current validation raises ValueError when no complete route is provided or when both complete routes are supplied simultaneously; it does not reject stray fields from the other representation. Shift arrays are required in both formats; for non-periodic systems pass integer all-zero shifts.

Returns:

forces – Forces on each atom.

Return type:

jax.Array, shape (N, 3)

Parameters:

positions (Array)
charges (Array)
cell (Array)
cutoff (float)
alpha (float)
neighbor_list (Array | None)
neighbor_ptr (Array | None)
neighbor_shifts (Array | None)
neighbor_matrix (Array | None)
neighbor_matrix_shifts (Array | None)
fill_value (int | None)
batch_idx (Array | None)

See also

coulomb_energy_forces: Compute both energies and forces

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

Compute Coulomb electrostatic energies and forces.

Computes pairwise electrostatic energies and forces using the Coulomb law, with optional erfc damping for Ewald/PME real-space calculations.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic charges.
cell (jax.Array, shape (3, 3), (1, 3, 3), or (B, 3, 3)) – Unit cell matrix. A single-system (3, 3) matrix is promoted to (1, 3, 3) internally.
cutoff (float) – Cutoff distance for interactions.
alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.
neighbor_list (jax.Array | None, shape (2, num_pairs)) – Neighbor pairs in COO format.
neighbor_ptr (jax.Array | None, shape (N+1,)) – CSR row pointers for neighbor list. Required with neighbor_list. Provided by neighborlist module.
neighbor_shifts (jax.Array | None, shape (num_pairs, 3)) – Integer unit cell shifts for neighbor list format.
neighbor_matrix (jax.Array | None, shape (N, max_neighbors)) – Neighbor indices in matrix format.
neighbor_matrix_shifts (jax.Array | None, shape (N, max_neighbors, 3)) – Integer unit cell shifts for matrix format.
fill_value (int | None) – Fill value for neighbor matrix padding. Applies only to matrix format.
batch_idx (jax.Array | None, shape (N,)) – Batch indices for each atom.

Return type:

tuple[Array, Array]

Notes

Callers must supply one complete supported neighbor topology route: CSR/COO neighbor_list, neighbor_ptr, and neighbor_shifts; or matrix neighbor_matrix and neighbor_matrix_shifts. Current validation raises ValueError when no complete route is provided or when both complete routes are supplied simultaneously; it does not reject stray fields from the other representation. Shift arrays are required in both formats; for non-periodic systems pass integer all-zero shifts.

Returns:

energies (jax.Array, shape (N,)) – Per-atom energies.
forces (jax.Array, shape (N, 3)) – Forces on each atom.

Parameters:

positions (Array)
charges (Array)
cell (Array)
cutoff (float)
alpha (float)
neighbor_list (Array | None)
neighbor_ptr (Array | None)
neighbor_shifts (Array | None)
neighbor_matrix (Array | None)
neighbor_matrix_shifts (Array | None)
fill_value (int | None)
batch_idx (Array | None)

Return type:

tuple[Array, Array]

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

Ewald Components#

Individual components of the Ewald summation method.

nvalchemiops.jax.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, *, energy_reduction='atom')[source]#

Compute real-space Ewald energy and optional direct derivative outputs.

Energy-only calls participate in JAX autodiff through a private custom-JVP wrapper. compute_forces=True remains a forward/direct escape hatch for no-autograd MD/inference loops; charge-gradient and virial direct outputs are deprecated training-style outputs and warn.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic partial charges.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices. A 2-D input is promoted to (1, 3, 3) internally.
alpha (float or jax.Array) – Ewald splitting parameter. A scalar float or array of shape (1,) or (B,).
neighbor_list (jax.Array or None, shape (2, M), optional) – Neighbor pairs in COO format; row 0 is idx_i, row 1 is idx_j. Provide either neighbor_list + neighbor_ptr + neighbor_shifts or neighbor_matrix + neighbor_matrix_shifts.
neighbor_ptr (jax.Array or None, shape (N+1,), optional) – CSR row pointers for neighbor_list.
neighbor_shifts (jax.Array or None, shape (M, 3), optional) – Integer periodic image shifts for each neighbor pair.
neighbor_matrix (jax.Array or None, shape (N, max_neighbors), optional) – Dense neighbor matrix; each row lists neighbor indices for one atom.
neighbor_matrix_shifts (jax.Array or None, shape (N, max_neighbors, 3), optional) – Integer periodic image shifts for each entry in neighbor_matrix.
mask_value (int or None, optional) – Sentinel indicating unused slots in neighbor_matrix. Defaults to N (number of atoms) when None.
batch_idx (jax.Array or None, shape (N,), optional) – System index per atom for batched mode. Atoms must be grouped contiguously with IDs 0..B-1.
compute_forces (bool, default=False) – Return explicit forces \(-\partial E / \partial \mathbf{r}_i\). For differentiable force computation prefer JAX autodiff.
compute_charge_gradients (bool, default=False) –

Deprecated since version 0.4.0: Deprecated. Return explicit \(\partial E / \partial q_i\). Raises DeprecationWarning when True.
compute_virial (bool, default=False) –

Deprecated since version 0.4.0: Deprecated. Return explicit virial tensor. Raises DeprecationWarning when True.
energy_reduction ({"atom", "system"}, keyword-only, default="atom") – Public energy layout. "atom" returns per-atom energies of shape (N,) (default, unchanged behavior). "system" returns per-system energies of shape (B,) (or (1,) for a single system), computed as a differentiable, uniform scatter-sum of the per-atom energies; it does not support nonuniform per-atom weighting. Only the energy field of the return value changes; forces, charge gradients, and the virial keep their existing atom/system layouts.

Returns:

jax.Array, shape (N,) or (B,) – Real-space Ewald energy when no derivative flags are set: per-atom when energy_reduction="atom", per-system when energy_reduction="system".
tuple[jax.Array, …] – (energies, forces) when compute_forces=True; (energies, forces, charge_gradients) when compute_charge_gradients=True; additionally appends the virial tensor of shape (1, 3, 3) or (B, 3, 3) when compute_virial=True.

Return type:

Array | tuple[Array, …]

See also

nvalchemiops.jax.interactions.electrostatics.ewald.ewald_reciprocal_space(): Reciprocal-space Ewald contribution.
nvalchemiops.jax.interactions.electrostatics.ewald.ewald_summation(): Complete Ewald summation combining both components.

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

Compute reciprocal-space Ewald energy and optional direct outputs.

Includes self-energy and background (net-charge) corrections so the returned energies are the full reciprocal contribution to the Ewald sum. Energy-only calls participate in JAX autodiff through a private custom-JVP wrapper. compute_forces=True remains a forward/direct escape hatch for no-autograd MD/inference loops; charge-gradient and virial direct outputs are deprecated training-style outputs and warn.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic partial charges.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices. A 2-D input is promoted to (1, 3, 3) internally.
k_vectors (jax.Array, shape (K, 3) or (B, K, 3)) – Reciprocal-space lattice vectors. Gradients are stopped internally so this argument must not depend on cell outside this function. Use nvalchemiops.jax.interactions.electrostatics.ewald.ewald_summation() for cell-differentiable k-vectors.
alpha (float or jax.Array) – Ewald splitting parameter. A scalar float or array of shape (1,) or (B,).
batch_idx (jax.Array or None, shape (N,), optional) – System index per atom for batched mode. Atoms must be grouped contiguously with IDs 0..B-1.
max_atoms_per_system (int or None, optional) – Maximum number of atoms in any single system. Required under jax.jit with batched inputs; inferred from data otherwise.
compute_forces (bool, default=False) – Return explicit forces \(-\partial E / \partial \mathbf{r}_i\). For differentiable force computation prefer JAX autodiff.
compute_charge_gradients (bool, default=False) – Deprecated. Return explicit \(\partial E / \partial q_i\). Raises DeprecationWarning when True.
compute_virial (bool, default=False) – Deprecated. Return explicit virial tensor. Raises DeprecationWarning when True.
energy_reduction ({"atom", "system"}, keyword-only, default="atom") – Public energy layout. "atom" returns per-atom energies of shape (N,) (default, unchanged behavior). "system" returns per-system energies of shape (B,) (or (1,) for a single system), computed as a differentiable, uniform scatter-sum of the per-atom energies; it does not support nonuniform per-atom weighting. Only the energy field of the return value changes; forces, charge gradients, and the virial keep their existing atom/system layouts.

Returns:

jax.Array, shape (N,) or (B,) – Reciprocal-space Ewald energy (with self and background corrections) when no derivative flags are set: per-atom when energy_reduction="atom", per-system when energy_reduction="system".
tuple[jax.Array, …] – (energies, forces) when compute_forces=True; (energies, forces, charge_gradients) when compute_charge_gradients=True; additionally appends the virial tensor of shape (1, 3, 3) or (B, 3, 3) when compute_virial=True.

Return type:

Array | tuple[Array, …]

See also

nvalchemiops.jax.interactions.electrostatics.ewald.ewald_real_space(): Real-space Ewald contribution.
nvalchemiops.jax.interactions.electrostatics.ewald.ewald_summation(): Complete Ewald summation combining both components.
nvalchemiops.jax.interactions.electrostatics.k_vectors.generate_k_vectors_ewald_summation(): Generates k_vectors from a cell and cutoff.

PME Components#

Individual components of the Particle Mesh Ewald method.

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

Compute PME reciprocal-space contribution.

Energy-only calls use a custom JVP so JAX does not attempt to differentiate the Warp spline/FFT FFI path. compute_forces=True remains a forward/direct escape hatch for no-autograd MD/inference loops; charge gradients, virial, and hybrid direct outputs are deprecated training-style outputs and warn.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic partial charges.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices with lattice vectors as rows.
alpha (jax.Array) – Ewald splitting parameter.
mesh_dimensions (tuple[int, int, int] or None, default=None) – Explicit FFT mesh dimensions. Required when cell, alpha, or batch metadata are traced by jax.jit or other JAX transformations.
mesh_spacing (float or None, default=None) – Target mesh spacing for eager-only mesh-size inference.
spline_order (int, default=4) – B-spline interpolation order.
batch_idx (jax.Array or None, default=None) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.
k_vectors (jax.Array or None) – Optional precomputed reciprocal grid values. These are setup constants for the JAX custom-JVP path; tangents through them are ignored. When supplied while differentiating with respect to cell, they are assumed to correspond to the current cell.
k_squared (jax.Array or None) – Optional precomputed reciprocal grid values. These are setup constants for the JAX custom-JVP path; tangents through them are ignored. When supplied while differentiating with respect to cell, they are assumed to correspond to the current cell.
compute_forces (bool) – Direct-output flags. compute_forces=True remains supported for no-autograd MD/inference use; charge-gradient and virial direct outputs are deprecated for differentiable training.
compute_charge_gradients (bool) – Direct-output flags. compute_forces=True remains supported for no-autograd MD/inference use; charge-gradient and virial direct outputs are deprecated for differentiable training.
compute_virial (bool) – Direct-output flags. compute_forces=True remains supported for no-autograd MD/inference use; charge-gradient and virial direct outputs are deprecated for differentiable training.
hybrid_forces (bool, default=False) – Deprecated charge-gradient injection mode for compatibility.
energy_reduction ({"atom", "system"}, keyword-only, default="atom") – Public energy layout. "atom" returns per-atom energies of shape (N,) (default, unchanged behavior). "system" returns per-system energies of shape (B,) (or (1,) for a single system), computed as a differentiable, uniform scatter-sum of the per-atom energies; it does not support nonuniform per-atom weighting. Only the energy field of the return value changes; forces, charge gradients, and the virial keep their existing atom/system layouts.
cell_inv_t (jax.Array or None) – Optional precomputed PME intermediates. These are setup constants for JAX and are not differentiable inputs. Cell-derived metadata such as cell_inv_t and volume is accepted while differentiating with respect to cell and is assumed to correspond to the current cell.
volume (jax.Array or None) – Optional precomputed PME intermediates. These are setup constants for JAX and are not differentiable inputs. Cell-derived metadata such as cell_inv_t and volume is accepted while differentiating with respect to cell and is assumed to correspond to the current cell.
moduli_x (jax.Array or None) – Optional precomputed PME intermediates. These are setup constants for JAX and are not differentiable inputs. Cell-derived metadata such as cell_inv_t and volume is accepted while differentiating with respect to cell and is assumed to correspond to the current cell.
moduli_y (jax.Array or None) – Optional precomputed PME intermediates. These are setup constants for JAX and are not differentiable inputs. Cell-derived metadata such as cell_inv_t and volume is accepted while differentiating with respect to cell and is assumed to correspond to the current cell.
moduli_z (jax.Array or None) – Optional precomputed PME intermediates. These are setup constants for JAX and are not differentiable inputs. Cell-derived metadata such as cell_inv_t and volume is accepted while differentiating with respect to cell and is assumed to correspond to the current cell.

Returns:

energies (jax.Array, shape (N,) or (B,)) – Reciprocal-space energies: per-atom when energy_reduction="atom", per-system when energy_reduction="system".
forces (jax.Array, shape (N, 3), optional) – Per-atom forces. Only present when compute_forces=True.
charge_gradients (jax.Array, shape (N,), optional) – Per-atom charge gradients \(\partial E/\partial q\). Only present when compute_charge_gradients=True (deprecated direct-output flag).
virial (jax.Array, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor. Only present when compute_virial=True (deprecated direct-output flag). Always last in the return tuple.

Return type:

Array | tuple[Array, Array] | tuple[Array, Array, Array] | tuple[Array, Array, Array, Array]

Notes

When cell or batch metadata are traced by jax.jit or other JAX transformations, pass explicit mesh_dimensions. If alpha would otherwise be estimated, precompute and pass it explicitly as well. mesh_spacing and accuracy-based parameter estimation depend on concrete setup values.

JAX PME higher-order support is limited to tested position and charge losses. Stress/cell/strain HVPs, alpha HVPs, and precomputed-metadata HVPs are unsupported until explicitly implemented and tested.

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

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

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

Parameters:

miller_indices (jax.Array, shape (N,)) – Integer Miller indices for one mesh axis, e.g. from jnp.fft.fftfreq(N, d=1.0/N) or jnp.fft.rfftfreq(N, d=1.0/N).
mesh_N (int) – Number of mesh points along this axis.
spline_order (int) – B-spline interpolation order (e.g. 4 for cubic B-splines).

Returns:

Per-Miller-index B-spline modulus sinc(m/N)^spline_order.

Return type:

jax.Array, shape (N,)

Slab Correction#

Explicit-output Yeh-Berkowitz/Ballenegger slab correction for systems with two periodic directions. Component-level calls can request energies, forces, charge gradients, and virials with the same flags used by the Ewald and PME wrappers. The high-level Ewald and PME wrappers can include the slab term in their energy autodiff path.

nvalchemiops.jax.interactions.electrostatics.compute_slab_correction(positions, charges, cell, pbc, batch_idx=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, *, energy_reduction='atom')[source]#

Yeh-Berkowitz/Ballenegger slab correction for 2D periodic systems.

Returns the standalone slab correction contribution for JAX electrostatics APIs. The caller can add the returned energy, force, charge-gradient, and virial terms to 3D-periodic Ewald or PME component outputs. Energy-only calls use explicit Warp-backed derivative paths; direct-output flags remain forward compatibility paths.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
charges (jax.Array, shape (N,)) – Atomic charges.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices.
pbc (jax.Array, shape (3,) or (B, 3), dtype=bool) – Per-system periodic boundary conditions. True marks periodic directions and False marks the non-periodic slab direction. Systems whose pbc row is not slab-like contribute zero. A shape (3,) array is accepted only for single-system calls.
batch_idx (jax.Array, shape (N,), dtype=int32, optional) – System index for each atom. Defaults to all zeros for a single system. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.
compute_forces (bool, default=False) – If True, return per-atom slab forces.
compute_charge_gradients (bool, default=False) – If True, return per-atom slab charge gradients dE_slab/dq_i.
compute_virial (bool, default=False) – If True, return per-system slab virial tensors.
energy_reduction ({"atom", "system"}, keyword-only, default="atom") – Public energy layout. "atom" returns per-atom energies of shape (N,) (default, unchanged behavior). "system" returns per-system energies of shape (B,) (or (1,) for a single system), computed as a differentiable, uniform scatter-sum of the per-atom energies; it does not support nonuniform per-atom weighting. Only the energy field of the return value changes; forces, charge gradients, and the virial keep their existing atom/system layouts.

Returns:

energies (jax.Array, shape (N,) or (B,)) – Slab correction energy: per-atom when energy_reduction="atom", per-system when energy_reduction="system".
forces (jax.Array, shape (N, 3), optional) – Per-atom slab force.
charge_gradients (jax.Array, shape (N,), optional) – Per-atom slab charge gradient.
virial (jax.Array, shape (B, 3, 3), optional) – Per-system slab virial tensor.

Return type:

Array | tuple[Array, …]

K-Vector Generation#

nvalchemiops.jax.interactions.electrostatics.generate_miller_indices(cell, k_cutoff)[source]#

Generate Miller index bounds for Ewald summation.

Parameters:

cell (jax.Array, shape (N, 3, 3)) – Unit cell matrices with lattice vectors as rows.
k_cutoff (float | jax.Array) – Maximum magnitude of k-vectors to include in reciprocal summation.

Returns:

Array of shape (3,) containing the maximum Miller indices (M_h, M_k, M_l) for each lattice direction.

Return type:

jax.Array

Notes

For batch mode, one shared set of Miller bounds is used for all systems. If k_cutoff is provided per system, the maximum cutoff across the batch is used to build those shared bounds.

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

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

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

Half-Space Optimization#

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

The half-space condition selects k-vectors where:

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

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

Mathematical Background#

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

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

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

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

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

param cell:: Unit cell matrix with lattice vectors as rows. Shape (3, 3) for single system or (B, 3, 3) for batch.
type cell:: jax.Array
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 jax.Array
param miller_bounds:: Precomputed maximum Miller indices (M_h, M_k, M_l) for each lattice direction. When provided, the function skips the internal computation of bounds from cell and k_cutoff, making it compatible with jax.jit (which requires static array shapes). Use generate_miller_indices() to compute these bounds eagerly before entering a JIT context. When None (default), bounds are computed automatically from cell and k_cutoff.
type miller_bounds:: tuple[int, int, int] | None, optional
returns:: Reciprocal lattice vectors within the cutoff. Shape (K, 3) for single system or (B, K, 3) for batch. Excludes k=0 and includes only half-space vectors.
rtype:: jax.Array

Examples

Single system with explicit k_cutoff:

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

With automatic parameter estimation:

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

JIT-compatible usage with precomputed bounds:

>>> from nvalchemiops.jax.interactions.electrostatics import generate_miller_indices
>>> cell = jnp.eye(3, dtype=jnp.float64)[None, ...] * 10.0
>>> bounds = generate_miller_indices(cell, k_cutoff=8.0)
>>> miller_bounds = (int(bounds[0]), int(bounds[1]), int(bounds[2]))
>>> # This can now be called inside @jax.jit
>>> k_vectors = generate_k_vectors_ewald_summation(cell, k_cutoff=8.0, miller_bounds=miller_bounds)

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_{\text{cutoff}}^3 \cdot V)\) where V is the cell volume.
When using inside jax.jit, you must provide miller_bounds as a concrete tuple[int, int, int]. The bounds determine array shapes (via jnp.arange), which must be statically known at trace time.

See also

ewald_reciprocal_space: Uses these k-vectors for reciprocal space energy.
estimate_ewald_parameters: Automatic parameter estimation including k_cutoff.
generate_miller_indices: Compute Miller bounds for JIT-compatible usage.

Parameters:

cell (Array)
k_cutoff (float | Array)
miller_bounds (tuple[int, int, int] | None)

Return type:

Array

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

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

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

Notes

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

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

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

In matrix form:

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

Each k-vector is then:

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

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

Parameters:

cell (jax.Array) – 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 (jax.Array, 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 (jax.Array, 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 (jax.Array, 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[Array, Array]

Examples

Basic usage:

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

With precomputed reciprocal cell:

>>> reciprocal_cell = 2 * jnp.pi * jnp.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 jnp.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.jax.interactions.electrostatics.estimate_ewald_parameters(positions, cell, batch_idx=None, accuracy=1e-6)[source]#

Estimate optimal Ewald summation parameters for a given accuracy.

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

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrix.
batch_idx (jax.Array, 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 jax.Array objects.

Return type:

EwaldParameters

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

Estimate PME parameters for a given accuracy.

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

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates.
cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrix.
batch_idx (jax.Array, shape (N,), dtype=int32, optional) – System index for each atom.
accuracy (float, default=1e-6) – Target accuracy.
real_space_cutoff (float, optional) – Caller-supplied cutoff. When given, \(\alpha\) is derived from it via \(\alpha = \sqrt{-\log \varepsilon} / r_c\); otherwise rc and \(\alpha\) come from \(\eta\).
mesh_safety_factor (float, default=1.0) – Multiplier on the standard mesh-size heuristic \(K = 2 \alpha L / (3 \varepsilon^{1/5})\). Raise for extra safety at tight \(\varepsilon\).

Returns:

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

Return type:

PMEParameters

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

Estimate PME mesh dimensions for a given accuracy.

The mesh size along each axis is chosen as

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

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

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

Parameters:

cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrix.
alpha (jax.Array, shape (B,)) – Ewald splitting parameter.
accuracy (float, default=1e-6) – Target relative accuracy.
mesh_safety_factor (float, default=1.0) – Multiplier on the standard heuristic. 1.0 is the well-tested default that meets accuracy across the configurations covered by the convergence script. Raise for extra paranoia at tight accuracy. Lower at your own risk: values below 1.0 can fail the accuracy guarantee on low-\(\alpha\) / large-L systems (verify with the convergence script before using).

Returns:

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

Return type:

tuple[int, int, int]

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

Convert mesh spacing to mesh dimensions.

Parameters:

cell (jax.Array) – Unit cell matrix.
mesh_spacing (float | jax.Array) – Target mesh spacing.

Returns:

Mesh dimensions, rounded up to powers of 2.

Return type:

tuple[int, int, int]

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

Container for Ewald summation parameters.

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

Parameters:

alpha (Array)
real_space_cutoff (Array)
reciprocal_space_cutoff (Array)

alpha#

Ewald splitting parameter (inverse length units).

Type:: jax.Array, shape (B,)

real_space_cutoff#

Real-space cutoff distance.

Type:: jax.Array, shape (B,)

reciprocal_space_cutoff#

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

Type:: jax.Array, shape (B,)

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

Container for PME parameters.

Parameters:

alpha (Array)
mesh_dimensions (tuple[int, int, int])
mesh_spacing (Array)
real_space_cutoff (Array)

alpha#

Ewald splitting parameter.

Type:: jax.Array, 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:: jax.Array, shape (B, 3)

real_space_cutoff#

Real-space cutoff distance.

Type:: jax.Array, shape (B,)

nvalchemiops.jax.interactions.electrostatics: Electrostatics#

High-Level Interface#

Coulomb Interactions#

Ewald Components#

PME Components#

Slab Correction#

K-Vector Generation#

Half-Space Optimization#

Mathematical Background#

Parameter Estimation#

`nvalchemiops.jax.interactions.electrostatics`: Electrostatics#