Ctrl+K

nvalchemiops.jax.interactions.electrostatics: Electrostatics#

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

Tip

For the underlying framework-agnostic Warp kernels, see nvalchemiops.interactions.electrostatics: Electrostatic Interactions (Warp).

High-Level Interface#

These are the primary entry points for most users. They are compatible with jax.jit when setup-only PME parameters such as mesh_dimensions and alpha are supplied explicitly whenever those values would otherwise be estimated from traced inputs. miller_bounds is also a static shape control: under jax.jit, pass it as a concrete tuple or build k_vectors outside the compiled function. Energy derivatives are defined for positions, charges, and cell. Setup values such as alpha and mesh controls are constants; precomputed reciprocal metadata such as k_vectors, k_squared, volume, and cell_inv_t is accepted for cell-differentiated calls as static metadata that is assumed to correspond to the current cell; cache-generation derivatives are not recovered. Energy-returning Ewald, PME, and slab paths support atom-weighted losses such as (weights * energies).sum() for positions, charges, and supported cell derivatives. JAX PME supports first-order cell/strain gradients, but PME cell/strain HVPs, including full PME with slab_correction=True, are explicitly unsupported until a native transposable PME cell-HVP path is implemented and tested. Point-charge Ewald/PME inputs support float32 and float64. Keep all floating inputs and precomputed metadata in a call on a consistent dtype.

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

Compute complete Ewald summation.

Parameters:

  • positions (jax.Array, shape (N, 3)) – Atomic coordinates.

  • charges (jax.Array, shape (N,)) – Atomic partial charges.

  • cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices.

  • alpha (float, jax.Array, or None, default=None) – Ewald splitting parameter. If None, estimated automatically.

  • k_vectors (jax.Array or None, default=None) – Reciprocal lattice vectors. Generated from cell when omitted.

  • k_cutoff (float, jax.Array, or None, default=None) – K-space cutoff used when generating k_vectors.

  • miller_bounds (tuple[int, int, int] or None, default=None, keyword-only) – Static Miller-index bounds for JIT-compatible k-vector generation.

  • batch_idx (jax.Array or None, default=None) – System index for each atom. When provided, atoms must be grouped by system: batch_idx must be contiguous, nondecreasing, and use system IDs 0..B-1.

  • max_atoms_per_system (int or None, default=None) – Static batch shape control for reciprocal kernels under jax.jit.

  • neighbor_list (jax.Array or None) – CSR neighbor-list inputs for the real-space component.

  • neighbor_ptr (jax.Array or None) – CSR neighbor-list inputs for the real-space component.

  • neighbor_shifts (jax.Array or None) – CSR neighbor-list inputs for the real-space component.

  • neighbor_matrix (jax.Array or None) – Dense neighbor-matrix inputs for the real-space component.

  • neighbor_matrix_shifts (jax.Array or None) – Dense neighbor-matrix inputs for the real-space component.

  • mask_value (int or None, default=None) – Sentinel value for invalid neighbor-matrix entries.

  • compute_forces (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use JAX autodiff for differentiable forces.

  • compute_charge_gradients (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag. Compute energy and use JAX autodiff for \(\partial E / \partial q_i\).

  • compute_virial (bool, default=False) –

    Deprecated since version 0.4.0: Deprecated direct-output flag for the virial tensor.

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

  • hybrid_forces (bool, default=False) – Deprecated direct-output flag retained for transition compatibility.

  • pbc (jax.Array, optional) – Per-system periodic boundary conditions for slab correction.

  • slab_correction (bool, default=False) – If True, add the Yeh-Berkowitz/Ballenegger slab correction.

Returns:

  • jax.Array, shape (N,) – Per-atom total Ewald energy when no deprecated direct-output flags are set. 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, *, 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.

  • 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,)) – Per-atom total electrostatic energies (real + reciprocal + slab).

  • 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 (1, 3, 3) or (B, 3, 3)) – Unit cell matrix. Shape (B, 3, 3) for batched calculations.

  • cutoff (float) – Cutoff distance for interactions.

  • alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.

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

  • batch_idx (jax.Array | None, shape (N,)) – Batch indices for each atom.

Returns:

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

Return type:

jax.Array, shape (N,)

Examples

>>> # Direct Coulomb (undamped)
>>> energies = coulomb_energy(
...     positions, charges, cell, cutoff=10.0, alpha=0.0,
...     neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr,
...     neighbor_shifts=neighbor_shifts
... )
>>> total_energy = energies.sum()

>>> # Ewald/PME real-space (damped)
>>> 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:

  • descriptions. (See coulomb_energy for parameter)

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

Returns:

forces – Forces on each atom.

Return type:

jax.Array, shape (N, 3)

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 (1, 3, 3) or (B, 3, 3)) – Unit cell matrix. Shape (B, 3, 3) for batched calculations.

  • cutoff (float) – Cutoff distance for interactions.

  • alpha (float, default=0.0) – Ewald splitting parameter. Use 0.0 for undamped Coulomb.

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

  • batch_idx (jax.Array | None, shape (N,)) – Batch indices for each atom.

Returns:

  • energies (jax.Array, shape (N,)) – Per-atom energies.

  • forces (jax.Array, shape (N, 3)) – Forces on each atom.

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

Returns:

  • jax.Array, shape (N,) – Per-atom real-space Ewald energy when no derivative flags are set.

  • 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)[source]#

Compute reciprocal-space Ewald energy and optional direct outputs.

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

Parameters:

  • positions (jax.Array, shape (N, 3)) – Atomic coordinates.

  • charges (jax.Array, shape (N,)) – Atomic partial charges.

  • cell (jax.Array, shape (3, 3) or (B, 3, 3)) – Unit cell matrices. A 2-D input is promoted to (1, 3, 3) internally.

  • k_vectors (jax.Array, shape (K, 3) or (B, K, 3)) – Reciprocal-space lattice vectors. Gradients are stopped internally so this argument must not depend on cell outside this function. Use nvalchemiops.jax.interactions.electrostatics.ewald.ewald_summation() for cell-differentiable k-vectors.

  • alpha (float or jax.Array) – Ewald splitting parameter. A scalar float or array of shape (1,) or (B,).

  • batch_idx (jax.Array or None, shape (N,), optional) – System index per atom for batched mode. Atoms must be grouped contiguously with IDs 0..B-1.

  • max_atoms_per_system (int or None, optional) – Maximum number of atoms in any single system. Required under jax.jit with batched inputs; inferred from data otherwise.

  • compute_forces (bool, default=False) – Return explicit forces \(-\partial E / \partial \mathbf{r}_i\). For differentiable force computation prefer JAX autodiff.

  • compute_charge_gradients (bool, default=False) – Deprecated. Return explicit \(\partial E / \partial q_i\). Raises DeprecationWarning when True.

  • compute_virial (bool, default=False) – Deprecated. Return explicit virial tensor. Raises DeprecationWarning when True.

Returns:

  • jax.Array, shape (N,) – Per-atom reciprocal-space Ewald energy (with self and background corrections) when no derivative flags are set.

  • tuple[jax.Array, …] – (energies, forces) when compute_forces=True; (energies, forces, charge_gradients) when compute_charge_gradients=True; additionally appends the virial tensor of shape (1, 3, 3) or (B, 3, 3) when compute_virial=True.

Return type:

Array | tuple[Array, …]

See also

nvalchemiops.jax.interactions.electrostatics.ewald.ewald_real_space()

Real-space Ewald contribution.

nvalchemiops.jax.interactions.electrostatics.ewald.ewald_summation()

Complete Ewald summation combining both components.

nvalchemiops.jax.interactions.electrostatics.k_vectors.generate_k_vectors_ewald_summation()

Generates k_vectors from a cell and cutoff.

PME Components#

Individual components of the Particle Mesh Ewald method.

nvalchemiops.jax.interactions.electrostatics.pme_reciprocal_space(positions, charges, cell, alpha, mesh_dimensions=None, mesh_spacing=None, spline_order=4, batch_idx=None, k_vectors=None, k_squared=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, hybrid_forces=False, *, 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.

  • 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,)) – Per-atom reciprocal-space energies.

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

Returns:

  • energies (jax.Array, shape (N,)) – Per-atom slab correction energy.

  • forces (jax.Array, shape (N, 3), optional) – Per-atom slab force.

  • charge_gradients (jax.Array, shape (N,), optional) – Per-atom slab charge gradient.

  • virial (jax.Array, shape (B, 3, 3), optional) – Per-system slab virial tensor.

Return type:

Array | tuple[Array, …]

K-Vector Generation#

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

Generate Miller index bounds for Ewald summation.

Parameters:

  • cell (jax.Array, shape (N, 3, 3)) – Unit cell matrices with lattice vectors as rows.

  • k_cutoff (float | jax.Array) – Maximum magnitude of k-vectors to include in reciprocal summation.

Returns:

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

Return type:

jax.Array

Notes

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

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

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

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

Half-Space Optimization#

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

The half-space condition selects k-vectors where:

  • h > 0, OR

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

  • (h == 0 AND k == 0 AND l > 0)

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

Mathematical Background#

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

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

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

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

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

param cell:

Unit cell matrix with lattice vectors as rows. Shape (3, 3) for single system or (B, 3, 3) for batch.

type cell:

jax.Array

param k_cutoff:

Maximum magnitude of k-vectors to include (\(|\mathbf{k}| \leq k_{\text{cutoff}}\)). Typical values: 8-12 \(\text{\AA}^{-1}\) for molecular systems. Higher values increase accuracy but also computational cost.

type k_cutoff:

float or jax.Array

param miller_bounds:

Precomputed maximum Miller indices (M_h, M_k, M_l) for each lattice direction. When provided, the function skips the internal computation of bounds from cell and k_cutoff, making it compatible with jax.jit (which requires static array shapes). Use generate_miller_indices() to compute these bounds eagerly before entering a JIT context. When None (default), bounds are computed automatically from cell and k_cutoff.

type miller_bounds:

tuple[int, int, int] | None, optional

returns:

Reciprocal lattice vectors within the cutoff. Shape (K, 3) for single system or (B, K, 3) for batch. Excludes k=0 and includes only half-space vectors.

rtype:

jax.Array

Examples

Single system with explicit k_cutoff:

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

With automatic parameter estimation:

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

JIT-compatible usage with precomputed bounds:

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

Notes

  • The k=0 vector is always excluded (causes division by zero in Green’s function).

  • For batch mode, the same set of Miller indices is used for all systems but transformed using each system’s reciprocal cell. If k_cutoff is given per system, the maximum cutoff across the batch determines the shared Miller bounds.

  • The number of k-vectors K scales as \(O(k_{\text{cutoff}}^3 \cdot V)\) where V is the cell volume.

  • When using inside jax.jit, you must provide miller_bounds as a concrete tuple[int, int, int]. The bounds determine array shapes (via jnp.arange), which must be statically known at trace time.

See also

ewald_reciprocal_space

Uses these k-vectors for reciprocal space energy.

estimate_ewald_parameters

Automatic parameter estimation including k_cutoff.

generate_miller_indices

Compute Miller bounds for JIT-compatible usage.

Parameters:
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:
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#

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#

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

On this page
Show Source