`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. The lower-level ewald_reciprocal_space component follows the tangent carried by its k_vectors argument. Vectors constructed from cell in the traced computation therefore contribute their reciprocal-cell derivative. A precomputed array, or one passed through jax.lax.stop_gradient, has zero tangent and is fixed for that differentiation. Full ewald_summation(k_vectors=...) and PME precomputed metadata treat explicit vectors as static metadata. Use ewald_reciprocal_space_from_miller_indices or ewald_summation(miller_indices=...) for retained topology materialized from the live cell. 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. For changing-cell full Ewald, retain signed Miller indices with generate_ewald_miller_indices and pass them through ewald_summation; the full-Ewald custom-JVP then includes the reciprocal-cell tangent. A retained topology must cover all intended cells. Enlarging it changes K and can trigger JIT recompilation. The legacy generate_miller_indices remains the batched bounds helper for static-shape JIT generation. 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. Import nvalchemiops.jax.interactions.electrostatics before creating JAX arrays intended to be float64 or compiling JAX functions that operate on float64 arrays. On import, the module sets JAX 64-bit support process-wide, including when it was previously disabled. This cannot restore precision in arrays already created as float32 or update functions that JAX has already compiled. float32 calculations remain supported.

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, miller_indices=None, energy_reduction='atom')[source]#

Compute complete Ewald summation.

Supply explicit k_vectors as fixed metadata, retained miller_indices to materialize vectors from the current cell, or neither to generate vectors from k_cutoff and optional miller_bounds.

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) – Explicit reciprocal vectors, treated as fixed metadata. When omitted, vectors are generated from k_cutoff or materialized from miller_indices.
k_cutoff (float, jax.Array, or None, default=None) – Reciprocal cutoff used only when generating vectors internally.
miller_bounds (tuple[int, int, int] or None, default=None, keyword-only) – Static Miller-index bounds used with internally generated vectors. Ignored when explicit k_vectors are supplied for compatibility.
miller_indices (jax.Array or None, default=None, keyword-only) – Caller-retained signed integer topology of shape (K, 3). Full Ewald materializes vectors from the current cell. Do not combine it with k_vectors, k_cutoff, or miller_bounds.
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=5, 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, alpha and the mesh dimensions are estimated in the call from accuracy outside JAX transformations. Pass params.alpha from estimate_pme_parameters to use prepared parameters.
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. When both mesh controls are omitted, the mesh is estimated in the call from accuracy outside JAX transformations. Pass params.mesh_dimensions to use prepared parameters.
spline_order (int, default=5) – B-spline interpolation order. Implicit meshes provide assignment-order support.
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 may be supplied without k_vectors. If compute_virial=True and k_vectors is omitted, Cartesian vectors are regenerated.
k_squared (jax.Array or None) – Precomputed PME reciprocal grid values. k_squared may be supplied without k_vectors. If compute_virial=True and k_vectors is omitted, Cartesian vectors are regenerated.
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 used by the in-call parameter estimators when alpha or the mesh controls are omitted.
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

Call estimate_pme_parameters before building the neighbor list, build that list with its real_space_cutoff, then pass its alpha and mesh_dimensions here. Use the same spline_order during setup and execution. Reuse these values while the cell and atom counts remain fixed, inside or outside jax.jit.

When cell, alpha, or batch metadata are traced by a JAX transformation, pass explicit mesh dimensions. When cell or batch metadata are traced and alpha would be estimated, precompute alpha as well. Position and charge derivatives can use implicit setup with concrete captured cell and batch values.

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. Energy autodiff uses the tangent carried by this argument. Vectors constructed from the differentiated cell in the traced computation contribute their reciprocal-cell derivative. A precomputed array, or one passed through jax.lax.stop_gradient, has zero tangent and is fixed with respect to that differentiation. Matching values alone do not create a dependency.
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.

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

Compute the reciprocal component from retained Miller indices.

Materializes Cartesian vectors from the supplied cell, then delegates to ewald_reciprocal_space(). Direct outputs, return order, warnings, and energy reduction match that component.

Parameters:

positions (Array) – Match ewald_reciprocal_space().
charges (Array) – Match ewald_reciprocal_space().
cell (Array) – Match ewald_reciprocal_space().
alpha (float | Array) – Match ewald_reciprocal_space().
batch_idx (Array | None) – Match ewald_reciprocal_space().
max_atoms_per_system (int | None) – Match ewald_reciprocal_space().
compute_forces (bool) – Match ewald_reciprocal_space().
compute_charge_gradients (bool) – Match ewald_reciprocal_space().
compute_virial (bool) – Match ewald_reciprocal_space().
energy_reduction (Literal['atom', 'system']) – Match ewald_reciprocal_space().
miller_indices (jax.Array, shape (K, 3)) – Caller-retained signed integer topology. See k_vectors_from_miller_indices() for its preconditions.

Returns:

Same result contract as ewald_reciprocal_space().

Return type:

jax.Array or tuple[jax.Array, …]

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=5, 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=5) – 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. k_squared may be supplied without k_vectors. For batches with identical cells, k_squared may contain one shared three-dimensional grid. If compute_virial=True and k_vectors is omitted, Cartesian vectors are regenerated. 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. k_squared may be supplied without k_vectors. For batches with identical cells, k_squared may contain one shared three-dimensional grid. If compute_virial=True and k_vectors is omitted, Cartesian vectors are regenerated. 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. mesh_spacing and component-only accuracy-based mesh sizing 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).

Parameters:

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

Return type:

Array

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

Generate the positive-half-space Miller topology for Ewald summation.

Derives conservative rectangular Miller bounds from k_cutoff for the reciprocal-space Ewald sum. Supplying miller_bounds instead enumerates that rectangle directly; neither mode applies a per-vector magnitude filter. 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:: Reciprocal cutoff used to derive conservative per-axis Miller bounds. The resulting rectangular topology can contain vectors whose magnitude exceeds this value.
type k_cutoff:: float or jax.Array
param miller_bounds:: Explicit Miller half-bounds (M_h, M_k, M_l). When supplied, their rectangle is enumerated directly and k_cutoff does not select individual rows. Use the legacy generate_miller_indices() bounds helper before jax.jit; it returns bounds, not full index rows.
type miller_bounds:: tuple[int, int, int] | None, optional
returns:: Signed integer Miller indices of shape (K, 3). The rows are nonzero, unique, and in the positive half-space.
rtype:: jax.Array

Examples

Single system with explicit k_cutoff:

>>> cell = jnp.eye(3, dtype=jnp.float64) * 10.0
>>> indices = generate_ewald_miller_indices(cell, k_cutoff=8.0)
>>> k_vectors = k_vectors_from_miller_indices(cell, indices)

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.

A retained topology must conservatively cover every cell state in which it will be used. Enlarging it changes K and can recompile jax.jit.

Parameters:

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

Return type:

Array

nvalchemiops.jax.interactions.electrostatics.k_vectors_from_miller_indices(cell, miller_indices)[source]#

Materialize reciprocal vectors from caller-retained Miller indices.

The supplied cell defines the live reciprocal transform. miller_indices must be a signed integer array of shape (K, 3); (0, 3) is valid. Duplicate rows, zero rows, and half-space membership are caller preconditions and are not checked. generate_ewald_miller_indices() produces valid topology. The result has shape (K, 3) for one cell or (B, K, 3) for a batch.

Parameters:

cell (Array)
miller_indices (Array)

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) for one system or (B, 3, 3) with the same batch size as cell.

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.

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

Generate squared PME reciprocal-grid magnitudes without k-vectors.

This is the memory-efficient metadata path for PME calls that do not request a reciprocal-space virial. It sums squared Cartesian components separately and avoids materializing the Cartesian vector grid, whose final axis is three times larger than the returned scalar grid.

Parameters:

cell (jax.Array) – Unit cell matrix with lattice vectors as rows. Shape (3, 3) for a single system or (B, 3, 3) for a batch.
mesh_dimensions (tuple[int, int, int]) – PME mesh dimensions (nx, ny, nz).
reciprocal_cell (jax.Array, optional) – Precomputed reciprocal cell matrix 2*pi*cell**-1. Shape (3, 3) for one system or (B, 3, 3) with the same batch size as cell.

Returns:

Safe squared magnitudes with shape (nx, ny, nz//2+1) for one system or (B, nx, ny, nz//2+1) for a batch. The zero mode is set to 1e-12.

Return type:

jax.Array

See also

generate_k_vectors_pme: Generate both Cartesian vectors and magnitudes.

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, *, spline_order=5, fft_padding_fraction=0.25)[source]#

Estimate PME splitting parameters and FFT-friendly mesh dimensions.

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) – Caller-requested multiplier on the continuous mesh-size estimate.
spline_order (int, default=5) – B-spline assignment order used for mesh sizing.
fft_padding_fraction (float, default=0.25) – Maximum fraction of extra total mesh points for snapping to powers of two, then smooth multiples of four. Zero selects the smallest 2/3/5/7-smooth mesh. Explicit other orders keep dyadic accuracy sizing.

Returns:

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

Return type:

PMEParameters

Notes

accuracy retains the public estimator’s relative sizing convention. Independent force checks determine the measured error for a charge distribution.

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

Return accuracy-based FFT mesh dimensions with spline support.

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\]

and rounded upward to a mesh with prime factors 2, 3, 5, and 7. For the default order five, snap to powers of two first, then smooth multiples of four, within the configured budget of extra total mesh points. Other spline orders retain upstream power-of-two accuracy sizing. Each dimension covers the continuous target and assignment support. accuracy controls the historical fifth-root sizing rule. The estimate retains its relative sizing convention; independent force checks determine the measured error for a particular charge distribution.

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) – Dimensionless input to the legacy mesh-size heuristic.
mesh_safety_factor (float, default=1.0) – Caller-requested multiplier on the continuous size estimate. The default leaves the estimate unchanged.
spline_order (int, default=5) – B-spline assignment order. The mesh covers its assignment support.
fft_padding_fraction (float, default=0.25) – Maximum fraction of extra total mesh points for snapping to powers of two, then smooth multiples of four. Zero selects the smallest 2/3/5/7-smooth mesh. Explicit other orders keep dyadic accuracy sizing.

Returns:

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

Return type:

tuple[int, int, int]

Notes

Measured example on an NVIDIA H100 NVL with Torch: for 131,072 NH3 atoms, snapping from 486 cubed to 512 cubed increased mesh points by 16.92% and reduced the energy, forces, and charge-gradient call from 13.96 ms to 11.17 ms. Both runs used order five, float64, a 9 Angstrom cutoff, and accuracy 1e-6. Times are medians of five groups of ten calls after three warmups; parameter setup and neighbor construction were excluded.

nvalchemiops.jax.interactions.electrostatics.mesh_spacing_to_dimensions(cell, mesh_spacing, *, spline_order=5, fft_padding_fraction=0.25)[source]#

Convert mesh spacing to mesh dimensions.

Parameters:

cell (jax.Array) – Unit cell matrix.
mesh_spacing (float | jax.Array) – Target mesh spacing.
spline_order (int, default=5) – B-spline assignment order. Each axis covers its assignment support.
fft_padding_fraction (float, default=0.25) – Maximum fraction of extra total mesh points for snapping to powers of two, then smooth multiples of four. Zero selects the smallest 2/3/5/7-smooth mesh. These spacing preferences apply to every spline order.

Returns:

Maximum dimensions across the batch, rounded upward to 2/3/5/7-smooth integers that cover the requested spacing and assignment support.

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#