`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=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. 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)
charges (Array)
cell (Array)
alpha (float | Array)
batch_idx (Array | None)
max_atoms_per_system (int | None)
miller_indices (Array)
compute_forces (bool)
compute_charge_gradients (bool)
compute_virial (bool)
energy_reduction (Literal['atom', 'system'])

Return type:

Array | tuple[Array, …]

:param : :param compute_forces: Match ewald_reciprocal_space(). :param compute_charge_gradients: Match ewald_reciprocal_space(). :param compute_virial: Match ewald_reciprocal_space(). :param energy_reduction: Match ewald_reciprocal_space(). :param miller_indices: 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, …]

Parameters:

positions (Array)
charges (Array)
cell (Array)
miller_indices (jax.Array, shape (K, 3))
alpha (float | Array)
batch_idx (Array | None)
max_atoms_per_system (int | None)
compute_forces (bool)
compute_charge_gradients (bool)
compute_virial (bool)
energy_reduction (Literal['atom', 'system'])

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

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