Models module (BaseModelMixin, ModelConfig, wrappers)

Every potential in nvalchemi —machine-learned or classical—is exposed through the same BaseModelMixin interface and described by a ModelConfig, so it drops into any BaseDynamics engine or training loop unchanged. This page is a per-model reference: what each model is, how to install it, and (for the classical models) the physics it implements. For the wrapping API, composition patterns, and output conventions, see the models user guide.

Core classes

ModelConfig	Unified model configuration combining capability declaration and runtime control.
NeighborConfig	Configuration for on-the-fly neighbor list construction.
BaseModelMixin	Abstract mixin providing a standardized interface for model wrappers.

Demo utilities

A minimal analytic model used throughout the tests and examples; useful as a template when wrapping a new potential.

DemoModel	This model is a simple demo model that computes the energies and conservative forces given an atomic point cloud (not graph!).
DemoModelWrapper	Wrapper for the demo model that implements the BaseModelMixin interface.

Machine-learned potentials

Wrappers for third-party machine-learned interatomic potentials (MLIPs). Each is an optional dependency installed through an extra; the wrapper adapts the model’s native inputs/outputs to the nvalchemi interface and exposes a from_checkpoint constructor that resolves named foundation-model checkpoints or local files. Energies come from the underlying model and forces/stresses are obtained by autograd on positions (conservative), unless the model provides them directly.

MACE

Equivariant message-passing network built on higher-order (E(3)-equivariant) features. MACEWrapper accepts any MACE variant (MACE, ScaleShiftMACE, cuEquivariance-converted, or torch.compile-d), builds a GPU one-hot node_attrs table to avoid per-step CPU round-trips, and supports PBC via neighbor-list shifts. from_checkpoint() loads local files or named MACE-MP foundation models (e.g. "medium-0b2").

Install with the mace extra (the equivariance kernels live in the CUDA group, so pair it with a cu extra):

uv sync --extra mace --extra cu13

MACEWrapper

Wrapper for any MACE model implementing the BaseModelMixin interface.

AIMNet2

Message-passing network with explicit charge equilibration. AIMNet2Wrapper computes energy as the primitive differentiable output (forces/stresses via autograd) and additionally exposes partial charges and per-atom AIM feature embeddings. It declares an external MATRIX-format neighbor list at the model’s AEV cutoff, satisfied by a NeighborListHook (or the pipeline). from_checkpoint() resolves checkpoints via AIMNet2Calculator.

Install with the aimnet extra:

uv sync --extra aimnet

AIMNet2Wrapper

Wrapper for AIMNet2 interatomic potentials.

UMA

fairchem’s Universal Models for Atoms — a single multi-task backbone with heads for OMol25 (molecules), OMat24 (crystals), and related datasets. UMAWrapper wraps a fairchem MLIPPredictUnit (the level at which energy/forces/stress are produced; the bare backbone yields only embeddings). The task is fixed at construction and active_outputs reflects what it supports. from_checkpoint() accepts registered names ("uma-s-1p1", "uma-s-1p2", "uma-m-1p1") or a local .pt.

The uma (fairchem) stack conflicts with the CUDA/MACE dependencies, so install it in a dedicated environment if you plan on using other MLIPs like MACE:

UV_PROJECT_ENVIRONMENT=.venv-uma uv sync --extra uma --extra ase

UMAWrapper

Wrapper for fairchem's UMA (Universal Models for Atoms).

Physical / classical models

Analytic potentials backed by Warp GPU kernels in nvalchemiops. Like the integrator kernels, the equations each implements are documented here so the physics is transparent. All work in the nvalchemi unit system (length \(\mathrm{\AA}\), energy \(\mathrm{eV}\), force \(\mathrm{eV}/\mathrm{\AA}\), charge in units of \(e\)) unless noted.

Lennard-Jones

A pairwise \(12\)–\(6\) potential summed over neighbor pairs within the cutoff:

\[E = \sum_{i<j} 4\varepsilon\left[ \left(\frac{\sigma}{r_{ij}}\right)^{12} - \left(\frac{\sigma}{r_{ij}}\right)^{6}\right].\]

An optional \(C^2\)-continuous switching function tapers the energy and its first two derivatives to zero over switch_width before the cutoff, avoiding force discontinuities; switch_width=0 gives a hard cutoff. Parameters: \(\varepsilon\) (well depth, eV), \(\sigma\) (zero-crossing distance, \(\mathrm{\AA}\)), and cutoff (\(\mathrm{\AA}\)).

LennardJonesModelWrapper

Warp-accelerated Lennard-Jones potential as a model wrapper.

Underlying nvalchemiops kernels

nvalchemiops.interactions.lj.lj_energy_forces(positions, cell, epsilon, sigma, cutoff, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, fill_value=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, batch_idx=None, switch_width=0.0, half_neighbor_list=True, device=None, energies_out=None, forces_out=None)

Compute Lennard-Jones energies and forces.

Parameters:

• positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic coordinates.

• cell (wp.array, shape (1,) or (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix. Use shape (B,) for batched mode.

• epsilon (float) – LJ energy parameter (well depth).

• sigma (float) – LJ length parameter (zero-crossing distance).

• cutoff (float) – Cutoff distance for interactions.

• neighbor_matrix (wp.array, shape (N, max_neighbors), dtype=wp.int32, optional) – Neighbor indices in matrix format. Provide either this or neighbor_list.

• neighbor_matrix_shifts (wp.array, shape (N, max_neighbors), dtype=wp.vec3i, optional) – Periodic shifts for each entry in neighbor_matrix.

• num_neighbors (wp.array, shape (N,), dtype=wp.int32, optional) – Valid neighbor count per atom; required when using matrix format.

• fill_value (int, optional) – Sentinel value used to pad neighbor_matrix rows.

• neighbor_list (wp.array, shape (2, M) or (M,), dtype=wp.int32, optional) – Neighbor target indices in COO/CSR adjacency form; alternative to matrix format.

• neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32, optional) – CSR row pointers; required when neighbor_list is provided.

• neighbor_shifts (wp.array, shape (M,), dtype=wp.vec3i, optional) – Periodic shifts for each edge in neighbor list format.

• batch_idx (wp.array, shape (N,), dtype=wp.int32, optional) – System index per atom (0..B-1). Pass None for single-system mode.

• switch_width (float, optional) – Width of the C2 switching region applied before cutoff. A value of 0.0 (default) disables switching.

• half_neighbor_list (bool, optional) – True (default) if the neighbor structure contains each pair once. Set to False for full neighbor lists where each pair appears twice.

• device (str, optional) – Warp device. If None, inferred from positions.

• energies_out (wp.array, shape (N,), dtype=wp.float32 or wp.float64, optional) – Pre-allocated output buffer for per-atom energies. Modified in-place (zeroed before use). If None, a new array is allocated.

• forces_out (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d, optional) – Pre-allocated output buffer for forces. Modified in-place (zeroed before use). If None, a new array is allocated.

Returns:

• wp.array, shape (N,), dtype=wp.float32 or wp.float64 – Per-atom LJ energies (matches input positions dtype).

• wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d – Forces on each atom (matches positions dtype).

Return type:

tuple[array, array]

nvalchemiops.interactions.lj.lj_energy_forces_virial(positions, cell, epsilon, sigma, cutoff, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, fill_value=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, batch_idx=None, switch_width=0.0, half_neighbor_list=True, device=None, energies_out=None, forces_out=None, virial_out=None)

Compute Lennard-Jones energies, forces, and virial tensor.

The virial tensor is needed for pressure/stress calculations in NPT/NPH.

Parameters:

• positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic coordinates.

• cell (wp.array, shape (1,) or (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix. Use shape (B,) for batched mode.

• epsilon (float) – LJ energy parameter (well depth).

• sigma (float) – LJ length parameter (zero-crossing distance).

• cutoff (float) – Cutoff distance for interactions.

• neighbor_matrix (wp.array, shape (N, max_neighbors), dtype=wp.int32, optional) – Neighbor indices in matrix format. Provide either this or neighbor_list.

• neighbor_matrix_shifts (wp.array, shape (N, max_neighbors), dtype=wp.vec3i, optional) – Periodic shifts for each entry in neighbor_matrix.

• num_neighbors (wp.array, shape (N,), dtype=wp.int32, optional) – Valid neighbor count per atom; required when using matrix format.

• fill_value (int, optional) – Sentinel value used to pad neighbor_matrix rows.

• neighbor_list (wp.array, shape (2, M) or (M,), dtype=wp.int32, optional) – Neighbor target indices in COO/CSR adjacency form; alternative to matrix format.

• neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32, optional) – CSR row pointers; required when neighbor_list is provided.

• neighbor_shifts (wp.array, shape (M,), dtype=wp.vec3i, optional) – Periodic shifts for each edge in neighbor list format.

• batch_idx (wp.array, shape (N,), dtype=wp.int32, optional) – System index per atom (0..B-1). Pass None for single-system mode.

• switch_width (float, optional) – Width of the C2 switching region applied before cutoff. A value of 0.0 (default) disables switching.

• half_neighbor_list (bool, optional) – True (default) if the neighbor structure contains each pair once. Set to False for full neighbor lists where each pair appears twice.

• device (str, optional) – Warp device. If None, inferred from positions.

• energies_out (wp.array, shape (N,), dtype=wp.float32 or wp.float64, optional) – Pre-allocated output buffer for per-atom energies. Modified in-place (zeroed before use). If None, a new array is allocated.

• forces_out (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d, optional) – Pre-allocated output buffer for forces. Modified in-place (zeroed before use). If None, a new array is allocated.

• virial_out (wp.array, shape (9,) or (B, 9), dtype=wp.float32 or wp.float64, optional) – Pre-allocated output buffer for the virial tensor. Modified in-place (zeroed before use). If None, a new array is allocated.

Returns:

• wp.array, shape (N,), dtype=wp.float32 or wp.float64 – Per-atom LJ energies (matches input positions dtype).

• wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d – Forces on each atom (matches positions dtype).

• wp.array, shape (9,) or (B, 9), dtype=wp.float32 or wp.float64 – Global virial tensor flattened as [xx, xy, xz, yx, yy, yz, zx, zy, zz] (matches input dtype). Shape is (B, 9) in batched mode.

Return type:

tuple[array, array, array]

DFT-D3 dispersion

Grimme’s DFT-D3 dispersion correction with Becke-Johnson (BJ) damping — a geometry-dependent \(C_6\)/\(C_8\) two-body dispersion term used to add long-range van der Waals attraction to a DFT or MLIP base energy:

\[E_\text{disp} = -\sum_{i<j}\ \sum_{n=6,8} s_n\,\frac{C_n^{ij}}{r_{ij}^{n} + \left(a_1 R_0^{ij} + a_2\right)^{n}}, \qquad R_0^{ij} = \sqrt{\frac{C_8^{ij}}{C_6^{ij}}} .\]

The coefficients \(C_6^{ij}\) interpolate with the atomic coordination number (controlled by k1/k3), so the correction responds to the local environment. Positions are supplied in \(\mathrm{\AA}\) and converted to Bohr internally; energies are returned in \(\mathrm{eV}\). Functional-specific parameters a1 (dimensionless), a2 (Bohr), and s8 (dimensionless) are required; reference \(C_n\)/\(R_0\) tables load from a cached dftd3_parameters.pt.

DFTD3ModelWrapper

DFT-D3(BJ) dispersion correction as a model wrapper.

Underlying nvalchemiops kernels

nvalchemiops.torch.interactions.dispersion.dftd3(positions, numbers, a1, a2, s8, k1=16.0, k3=-4.0, s6=1.0, s5_smoothing_on=1e10, s5_smoothing_off=1e10, fill_value=None, d3_params=None, covalent_radii=None, r4r2=None, c6_reference=None, coord_num_ref=None, batch_idx=None, cell=None, neighbor_matrix=None, neighbor_matrix_shifts=None, neighbor_list=None, neighbor_ptr=None, unit_shifts=None, compute_virial=False, num_systems=None, device=None)

Compute DFT-D3(BJ) dispersion energy and forces using Warp with optional periodic boundary condition support and smoothing function.

DFT-D3 parameters must be explicitly provided using one of three methods:

1. D3Parameters dataclass: Supply a D3Parameters instance (recommended). Individual parameters can override dataclass values if both are provided.

2. Explicit parameters: Supply all four parameters individually: covalent_radii, r4r2, c6_reference, and coord_num_ref.

3. Dictionary: Provide a d3_params dictionary with keys: "rcov", "r4r2", "c6ab", and "cn_ref". Individual parameters can override dictionary values if both are provided.

See examples/dispersion/utils.py for parameter generation utilities.

This wrapper can be launched by either supplying a neighbor matrix or a neighbor list, both of which can be generated by the nvalchemiops.torch.neighbors.neighbor_list() function where the latter can be returned by setting the return_neighbor_list parameter to True.

Parameters:

• positions (torch.Tensor) – Atomic coordinates [num_atoms, 3] as float32 or float64, in consistent distance units (conventionally Bohr when using standard D3 parameters)

• numbers (torch.Tensor) – Atomic numbers [num_atoms] as int32

• a1 (float) – Becke-Johnson damping parameter 1 (functional-dependent, dimensionless)

• a2 (float) – Becke-Johnson damping parameter 2 (functional-dependent), in same units as positions

• s8 (float) – \(C_8\) term scaling factor (functional-dependent, dimensionless)

• k1 (float, optional) – CN counting function steepness parameter, in inverse distance units (typically 16.0 1/Bohr for atomic units)

• k3 (float, optional) – CN interpolation Gaussian width parameter (typically -4.0, dimensionless)

• s6 (float, optional) – \(C_6\) term scaling factor (typically 1.0, dimensionless)

• s5_smoothing_on (float, optional) – Distance where S5 switching begins, in same units as positions. Set greater or equal to s5_smoothing_off to disable smoothing. Default: 1e10

• s5_smoothing_off (float, optional) – Distance where S5 switching completes, in same units as positions. Default: 1e10 (effectively no cutoff)

• fill_value (int | None, optional) – Value indicating padding in neighbor_matrix. If None, defaults to num_atoms. Entries with neighbor_matrix[i, k] >= fill_value are treated as padding. Default: None

• d3_params (D3Parameters | dict[str, torch.Tensor] | None, optional) – DFT-D3 parameters provided as either: - D3Parameters dataclass instance (recommended) - Dictionary with keys: “rcov”, “r4r2”, “c6ab”, “cn_ref” Individual parameters below can override values from d3_params.

• covalent_radii (torch.Tensor | None, optional) – Covalent radii [max_Z+1] as float32, indexed by atomic number, in same units as positions. If provided, overrides the value in d3_params.

• r4r2 (torch.Tensor | None, optional) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values [max_Z+1] as float32 for \(C_8\) computation (dimensionless). If provided, overrides the value in d3_params.

• c6_reference (torch.Tensor | None, optional) – \(C_6\) reference values [max_Z+1, max_Z+1, 5, 5] as float32 in energy \(\times\) distance\(^6\) units. If provided, overrides the value in d3_params.

• coord_num_ref (torch.Tensor | None, optional) – CN reference grid [max_Z+1, max_Z+1, 5, 5] as float32 (dimensionless). If provided, overrides the value in d3_params.

• batch_idx (torch.Tensor or None, optional) – Batch indices [num_atoms] as int32. If None, all atoms are assumed to be in a single system (batch 0). For batched calculations, atoms with the same batch index belong to the same system. Default: None

• cell (torch.Tensor or None, optional, as float32 or float64) – Unit cell lattice vectors [num_systems, 3, 3] for PBC, in same dtype and units as positions. Convention: cell[s, i, :] is i-th lattice vector for system s. If None, non-periodic calculation. Default: None

• neighbor_matrix (torch.Tensor | None, optional) – Neighbor indices [num_atoms, max_neighbors] as int32 in dense row format. Row i lists the neighbor atom indices of atom i; unused slots are padded with values >= fill_value. Requires a symmetric neighbor representation (each pair appears in both rows). Mutually exclusive with neighbor_list. Default: None

• neighbor_matrix_shifts (torch.Tensor or None, optional) – Integer unit cell shifts [num_atoms, max_neighbors, 3] as int32 for PBC with neighbor_matrix format. If None, non-periodic calculation. If provided along with cell, Cartesian shifts are computed. Mutually exclusive with unit_shifts. Default: None

• neighbor_list (torch.Tensor or None, optional) – Neighbor pairs [2, num_pairs] as int32 in COO format, where row 0 contains source atom indices and row 1 contains target atom indices. Alternative to neighbor_matrix for sparse neighbor representations. Mutually exclusive with neighbor_matrix. Must be used together with neighbor_ptr (both are returned by the neighbor list API when return_neighbor_list=True). Default: None

• neighbor_ptr (torch.Tensor or None, optional) – CSR row pointers [num_atoms+1] as int32. Required when using neighbor_list. Indicates that neighbor_list[1, :] contains destination atoms in CSR format where neighbor_ptr[i]:neighbor_ptr[i+1] gives the range of neighbors for atom i. Returned by the neighbor list API when return_neighbor_list=True. Default: None

• unit_shifts (torch.Tensor or None, optional) – Integer unit cell shifts [num_pairs, 3] as int32 for PBC with neighbor_list format. If None, non-periodic calculation. If provided along with cell, Cartesian shifts are computed. Mutually exclusive with neighbor_matrix_shifts. Default: None

• compute_virial (bool, optional) – If True, allocate and compute virial tensor. Ignored if virial parameter is provided. Default: False

• num_systems (int, optional) – Number of systems in batch. In none provided, inferred from cell or from batch_idx (introcudes CUDA synchronization overhead). Default: None

• device (str or None, optional) – Warp device string (e.g., ‘cuda:0’, ‘cpu’). If None, inferred from positions tensor. Default: None

Returns:

• energy (torch.Tensor) – Total dispersion energy [num_systems] as float32. Units are energy (Hartree when using standard D3 parameters).

• forces (torch.Tensor) – Atomic forces [num_atoms, 3] as float32. Units are energy/distance (Hartree/Bohr when using standard D3 parameters).

• coord_num (torch.Tensor) – Coordination numbers [num_atoms] as float32 (dimensionless)

• virial (torch.Tensor, optional) – Virial tensor [num_systems, 3, 3] as float32. Units are energy (Hartree when using standard D3 parameters). Only returned if compute_virial=True.

Return type:

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

Notes

• Unit consistency: All inputs must use consistent units. Standard D3 parameters from the Grimme group use atomic units (Bohr for distances, Hartree for energy), so using atomic units throughout is recommended and conventional.

• Float32 or float64 precision for positions and cell; outputs always float32

• Neighbor formats: Supports both neighbor_matrix (dense) and neighbor_list (sparse COO) formats. Choose neighbor_list for sparse systems or when memory efficiency is important.

• Padding atoms indicated by numbers[i] == 0

• Requires symmetric neighbor representation (each pair appears twice)

• Two-body only: Computes pairwise \(C_6\) and \(C_8\) dispersion terms; three-body Axilrod-Teller-Muto (ATM/\(C_9\)) terms are not included

• Virial computation requires periodic boundary conditions.

• Bulk stress tensor can be obtained by dividing virial by system volume.

Neighbor Format Selection:

• Use neighbor_matrix for dense systems or when max_neighbors is small

• Use neighbor_list for sparse systems, large cutoffs, or memory-constrained scenarios

• Both formats compute the same model and support PBC; results may differ by floating-point roundoff due to traversal order

PBC Handling:

• Matrix format: Provide cell and neighbor_matrix_shifts

• List format: Provide cell and unit_shifts

• Non-periodic: Omit both cell and shift parameters

See also

D3Parameters

Dataclass for organizing DFT-D3 reference parameters

_dftd3_matrix_op()

Internal custom operator for neighbor matrix format (non-PBC)

_dftd3_matrix_pbc_op()

Internal custom operator for neighbor matrix format (PBC)

_dftd3_op()

Internal custom operator for neighbor list format (non-PBC)

_dftd3_pbc_op()

Internal custom operator for neighbor list format (PBC)

Ewald summation

Exact long-range electrostatics for periodic systems, splitting the Coulomb sum into a short-range real-space part and a smooth reciprocal-space part with a Gaussian screening width \(\alpha\):

\[E = \underbrace{\frac{k_e}{2}\sum_{i \ne j} q_i q_j \frac{\operatorname{erfc}(\alpha r_{ij})}{r_{ij}}}_{\text{real space}} + \underbrace{\frac{k_e}{2V}\sum_{\mathbf{k}\ne 0} \frac{4\pi}{k^2}\,e^{-k^2/4\alpha^2}\,|S(\mathbf{k})|^2}_{\text{reciprocal space}} - \underbrace{\frac{k_e\,\alpha}{\sqrt{\pi}}\sum_i q_i^2}_{\text{self}},\]

with the structure factor \(S(\mathbf{k}) = \sum_i q_i e^{i\mathbf{k}\cdot\mathbf{r}_i}\) (a neutralising background term is added for charged cells). The real-space part uses a neighbor matrix within cutoff; \(\alpha\) and the reciprocal cutoff are chosen automatically from the requested accuracy. The Coulomb prefactor is \(k_e = 14.3996\ \mathrm{eV}\cdot\mathrm{\AA}/e^2\), and an optional slab correction removes spurious periodic images along one axis for 2-D systems.

EwaldModelWrapper

Ewald summation electrostatics potential as a model wrapper.

Underlying nvalchemiops kernels

nvalchemiops.torch.interactions.electrostatics.ewald.ewald_real_space(positions, charges, cell, alpha, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, neighbor_matrix=None, neighbor_matrix_shifts=None, mask_value=None, batch_idx=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, hybrid_forces=False, *, energy_reduction='atom')

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

Computes the damped Coulomb interactions for atom pairs within the real-space cutoff. The complementary error function (erfc) damping ensures rapid convergence in real space.

Parameters:

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

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

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

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

• neighbor_list (torch.Tensor, shape (2, M), optional) – Neighbor list in COO format.

• neighbor_ptr (torch.Tensor, shape (N+1,), optional) – CSR row pointers for neighbor list.

• neighbor_shifts (torch.Tensor, shape (M, 3), optional) – Periodic image shifts for neighbor list.

• neighbor_matrix (torch.Tensor, shape (N, max_neighbors), optional) – Dense neighbor matrix format.

• neighbor_matrix_shifts (torch.Tensor, shape (N, max_neighbors, 3), optional) – Periodic image shifts for neighbor_matrix.

• mask_value (int, optional) – Value indicating invalid entries in neighbor_matrix. Defaults to N.

• batch_idx (torch.Tensor, shape (N,), optional) – 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.

• compute_forces (bool, default=False) – Whether to compute explicit component forces. This direct output is kept for no-autograd MD/inference use; use energy autograd for differentiable training.

• compute_charge_gradients (bool, default=False) – Whether to compute explicit component charge gradients. This direct output follows the same no-autograd contract as compute_forces.

• compute_virial (bool, default=False) – Whether to compute the component virial tensor \(W = -\partial E / \partial \varepsilon\). Stress = -virial / volume.

• hybrid_forces (bool, default=False) – Enables the legacy direct-output path. When charges.requires_grad, ordinary first-order losses whose cotangent is uniform within each system use detached positions/cell and cached charge gradients through a straight-through connector. Non-uniform per-atom losses and create_graph=True rebuild the eager energy graph with geometry and charge-chain derivatives. Fixed-charge hybrid calls remain forward-only. Forces and virial are forward-only. Do not add direct analytical forces to a fallback-derived full geometry gradient.

• energy_reduction ({"atom", "system"}, default="atom") – Return per-atom energies (N,) or summed per-system energies (B,).

Returns:

• energies (torch.Tensor, shape (N,) or (B,)) – Real-space Ewald energy: per-atom when energy_reduction="atom", per-system when energy_reduction="system".

• forces (torch.Tensor, shape (N, 3), optional) – Direct component forces (if compute_forces=True).

• charge_gradients (torch.Tensor, shape (N,), optional) – Direct component charge gradients (if compute_charge_gradients=True).

• virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor (if compute_virial=True). Always last in the tuple.

Return type:

Tensor | tuple[Tensor, …]

Note

Energies are always float64 for numerical stability during accumulation. Forces, virial, and charge gradients match the input dtype (float32 or float64).

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115). Hybrid direct-output mode uses the same cached fallback connector so weighted \(q = q(R)\) losses can recover a valid energy gradient when the forward energy was detached.

nvalchemiops.torch.interactions.electrostatics.ewald.ewald_reciprocal_space(positions, charges, cell, k_vectors, alpha, batch_idx=None, compute_forces=False, compute_charge_gradients=False, compute_virial=False, hybrid_forces=False, *, max_atoms_per_system=None, energy_reduction='atom')

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

Computes the smooth long-range electrostatic contribution using structure factors in reciprocal space.

Parameters:

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

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

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

• k_vectors (torch.Tensor) – Reciprocal lattice vectors. Shape (K, 3) for single system, (B, K, 3) for batch. When both cell and k_vectors require gradients, the supplied k-vector graph is preserved. Physical strain derivatives require vectors generated from the same differentiable cell. Non-differentiable vectors, and grad-bearing leaf vectors without a cell edge, are fixed Cartesian metadata for cell derivatives.

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

• batch_idx (torch.Tensor, shape (N,), optional) – 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.

• compute_forces (bool, default=False) – Whether to compute explicit component forces. This direct output is kept for no-autograd MD/inference use; use energy autograd for differentiable training.

• compute_charge_gradients (bool, default=False) – Whether to compute explicit component charge gradients. This direct output follows the same no-autograd contract as compute_forces.

• compute_virial (bool, default=False) – Whether to compute the component virial tensor \(W = -\partial E / \partial \varepsilon\). Stress = -virial / volume.

• hybrid_forces (bool, default=False) – Enables the legacy direct-output path. With charges.requires_grad, uniform first-order cotangents use cached charge gradients; non-uniform per-atom losses and create_graph=True rebuild the eager energy graph with geometry and charge-chain derivatives. Fixed-charge hybrid calls remain forward-only. See ewald_real_space() for the complete contract.

• max_atoms_per_system (int, optional, keyword-only) – Maximum number of atoms in any single system when batch_idx is provided. Passing this host-known upper bound avoids CUDA host synchronization from launch-size inference in the reciprocal kernel. Overestimates are safe but may launch extra blocks. When omitted, the bound is inferred from atom_start / atom_end and may synchronize on CUDA.

• energy_reduction ({"atom", "system"}, default="atom") – Return per-atom energies (N,) or summed per-system energies (B,).

Returns:

• energies (torch.Tensor, shape (N,) or (B,)) – Reciprocal-space Ewald energy: per-atom when energy_reduction="atom", per-system when energy_reduction="system".

• forces (torch.Tensor, shape (N, 3), optional) – Direct component forces (if compute_forces=True).

• charge_gradients (torch.Tensor, shape (N,), optional) – Direct component charge gradients (if compute_charge_gradients=True).

• virial (torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional) – Virial tensor (if compute_virial=True). Always last in the tuple.

Return type:

Tensor | tuple[Tensor, …]

Note

Energies are always float64 for numerical stability during accumulation. Forces, virial, and charge gradients match the input dtype (float32 or float64). For eager execution, a differentiable cell paired with fixed Cartesian k_vectors emits a warning because its cell derivative is not the physical Ewald strain virial. This advisory warning is suppressed under torch.compile. Generate vectors from the differentiable cell with fixed Miller bounds for physical strain derivatives.

When charges is a non-leaf tensor that may depend on positions (\(q = q(R)\)), ordinary first-order losses may use cached partial derivatives and let PyTorch apply \(\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R\) once. Weighted losses and higher-order derivatives recompute safe partials or connected gradients as needed to avoid double-counting that chain term (issue #115).

Particle-mesh Ewald (PME)

The same real/reciprocal Ewald decomposition, but the reciprocal sum is evaluated by spreading charges onto a grid with B-spline interpolation and using an FFT, giving \(O(N \log N)\) scaling that is far cheaper than direct Ewald for large systems. Parameters: cutoff (\(\mathrm{\AA}\)), the mesh resolution via mesh_spacing (\(\mathrm{\AA}\)) or explicit mesh_dimensions, and the B-spline spline_order (higher is smoother and more accurate).

PMEModelWrapper

Particle Mesh Ewald electrostatics potential as a model wrapper.

Underlying nvalchemiops kernels

The real-space term reuses ewald_real_space() (documented above under Ewald); the reciprocal term is evaluated on the mesh via B-spline spread/gather:

nvalchemiops.torch.spline.spline_spread(positions, values, cell, mesh_dims, spline_order=4, batch_idx=None, cell_inv_t=None)

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

Parameters:

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

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

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

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

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

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

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

Returns:

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

Return type:

torch.Tensor

nvalchemiops.torch.spline.spline_gather_with_force(positions, charges, mesh, cell, spline_order=4, batch_idx=None, cell_inv_t=None)

Fused gather of scalar potential AND derivative-based force from one mesh.

Returns (output, forces) where:

• output[atom] = \(\sum_g \text{mesh}[g] \cdot w(\text{atom}, g)\) — raw potential per atom (the caller multiplies by charge in the PME corrections step).

• forces[atom] = \(-q_\text{atom} \sum_g \text{mesh}[g] \cdot C^{-T} \nabla w\) — Cartesian force.

This replaces spline_gather(...) followed by spline_gather_gradient(...) on the same mesh: each thread reads its stencil cell ONCE and accumulates both outputs. Halves the mesh DRAM traffic and shares the per-thread weight derivative work across both channels.

Parameters mirror spline_gather_gradient. For spline_order in the set the per-order kernels cover ({2, 3, 4, 5, 6}), both single-system and batched inputs use the fused kernel directly. For unsupported orders, batched inputs fall back to the two-kernel sequence (spline_gather + spline_gather_gradient).

Parameters:

• positions (Tensor)

• charges (Tensor)

• mesh (Tensor)

• cell (Tensor)

• spline_order (int)

• batch_idx (Tensor | None)

• cell_inv_t (Tensor | None)

Return type:

tuple[Tensor, Tensor]

Composition

Combine several models (e.g. an MLIP plus a long-range electrostatics term) into a single potential; see the models user guide for wiring and neighbor-list sharing.

PipelineModelWrapper	Compose multiple models via a grouped pipeline.
PipelineStep	Wraps a model with an output rename mapping.
PipelineGroup	A group of steps that share a derivative computation strategy.