`nvalchemiops.jax.interactions.dispersion`: Dispersion Corrections#

The dispersion module provides JAX bindings for the GPU-accelerated implementations of dispersion interactions.

JAX dispersion interactions API.

This module provides JAX bindings for dispersion corrections (DFT-D3).

Tip

For the underlying framework-agnostic Warp kernels, see nvalchemiops.interactions.dispersion: Dispersion Corrections.

High-Level Interface#

DFT-D3(BJ) Dispersion Corrections#

The DFT-D3 implementation supports two neighbor representation formats:

Neighbor matrix (dense): [num_atoms, max_neighbors] with padding
Neighbor list (sparse CSR): Compressed sparse row format with idx_j and neighbor_ptr

Both formats produce identical results and support all features including periodic boundary conditions, batching, and smooth cutoff functions. The high-level wrapper automatically dispatches to the appropriate kernels based on which format is provided. The method should be jax.jit compatible.

nvalchemiops.jax.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)[source]#

Compute DFT-D3(BJ) dispersion energy and forces using Warp with JAX arrays.

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

D3Parameters dataclass: Supply a D3Parameters instance (recommended). Individual parameters can override dataclass values if both are provided.
Explicit parameters: Supply all four parameters individually: covalent_radii, r4r2, c6_reference, and coord_num_ref.
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.

Parameters:

positions (jax.Array) – Atomic coordinates [num_atoms, 3] as float32 or float64, in consistent distance units (conventionally Bohr when using standard D3 parameters)
numbers (jax.Array) – 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). Default: 16.0
k3 (float, optional) – CN interpolation Gaussian width parameter (typically -4.0, dimensionless). Default: -4.0
s6 (float, optional) – \(C_6\) term scaling factor (typically 1.0, dimensionless). Default: 1.0
s5_smoothing_on (float, optional) – Distance where S5 switching begins, in same units as positions. 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. Default: None
d3_params (D3Parameters | dict[str, jax.Array] | 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 (jax.Array | 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 (jax.Array | 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 (jax.Array | 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 (jax.Array | 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 (jax.Array or None, optional) – Batch indices [num_atoms] as int32. If None, all atoms are assumed to be in a single system (batch 0). Default: None
cell (jax.Array or None, optional) – 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 (jax.Array | None, optional) – Neighbor indices [num_atoms, max_neighbors] as int32. Each row i contains indices of atom i’s neighbors, padded with fill_value for unused slots. Mutually exclusive with neighbor_list. Default: None
neighbor_matrix_shifts (jax.Array 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. Mutually exclusive with unit_shifts. Default: None
neighbor_list (jax.Array 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. Default: None
neighbor_ptr (jax.Array 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. Default: None
unit_shifts (jax.Array or None, optional) – Integer unit cell shifts [num_pairs, 3] as int32 for PBC with neighbor_list format. If None, non-periodic calculation. Mutually exclusive with neighbor_matrix_shifts. Default: None
compute_virial (bool, optional) – If True, compute and return virial tensor. Default: False
num_systems (int, optional) – Number of systems in batch. If None, inferred from cell or from batch_idx (introduces device synchronization overhead). Default: None

Returns:

energy (jax.Array) – Total dispersion energy [num_systems] as float32. Units are energy (Hartree when using standard D3 parameters).
forces (jax.Array) – Atomic forces [num_atoms, 3] as float32. Units are energy/distance (Hartree/Bohr when using standard D3 parameters).
coord_num (jax.Array) – Coordination numbers [num_atoms] as float32 (dimensionless)
virial (jax.Array, optional) – Virial tensor [num_systems, 3, 3] as float32. Only returned if compute_virial=True.

Return type:

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

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).
Float32 or float64 precision for positions and cell; outputs always float32
Neighbor formats: Supports both neighbor_matrix (dense) and neighbor_list (sparse) 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.

Raises:

ValueError – If neighbor format is invalid or PBC requirements are not met
RuntimeError – If DFT-D3 parameters are not provided

Parameters:

positions (Array)
numbers (Array)
a1 (float)
a2 (float)
s8 (float)
k1 (float)
k3 (float)
s6 (float)
s5_smoothing_on (float)
s5_smoothing_off (float)
fill_value (int | None)
d3_params (D3Parameters | dict[str, Array] | None)
covalent_radii (Array | None)
r4r2 (Array | None)
c6_reference (Array | None)
coord_num_ref (Array | None)
batch_idx (Array | None)
cell (Array | None)
neighbor_matrix (Array | None)
neighbor_matrix_shifts (Array | None)
neighbor_list (Array | None)
neighbor_ptr (Array | None)
unit_shifts (Array | None)
compute_virial (bool)
num_systems (int | None)

Return type:

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

Examples

Using neighbor matrix format:

>>> energy, forces, coord_num = dftd3(
...     positions, numbers,
...     neighbor_matrix=neighbor_matrix,
...     a1=0.3981, a2=4.4211, s8=1.9889,
...     d3_params=params,
... )

Using neighbor list format with PBC:

>>> energy, forces, coord_num, virial = dftd3(
...     positions, numbers,
...     neighbor_list=neighbor_list,
...     neighbor_ptr=neighbor_ptr,
...     a1=0.3981, a2=4.4211, s8=1.9889,
...     d3_params=params,
...     cell=cell,
...     unit_shifts=unit_shifts,
...     compute_virial=True,
... )

Data Structures#

This data structure is not necessarily required to use the kernels, however is provided for convenience—the dataclass will validate shapes and keys for parameters required by the kernels.

class nvalchemiops.jax.interactions.dispersion.D3Parameters(rcov, r4r2, c6ab, cn_ref, interp_mesh=5)[source]#

DFT-D3 reference parameters for dispersion correction calculations.

This dataclass encapsulates all element-specific parameters required for DFT-D3 dispersion corrections. The main purpose for this structure is to provide validation, ensuring the correct shapes, dtypes, and keys are present and complete. These parameters are used by dftd3().

Parameters:

rcov (jax.Array) – Covalent radii [max_Z+1] as float32. Units should be consistent with position coordinates. Index 0 is reserved for padding; valid atomic numbers are 1 to max_Z.
r4r2 (jax.Array) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values [max_Z+1] as float32. Dimensionless ratio used for computing \(C_8\) coefficients from \(C_6\) values.
c6ab (jax.Array) – \(C_6\) reference coefficients [max_Z+1, max_Z+1, interp_mesh, interp_mesh] as float32. Units are energy \(\times\) distance\(^6\). Indexed by atomic numbers and coordination number reference indices.
cn_ref (jax.Array) – Coordination number reference grid [max_Z+1, max_Z+1, interp_mesh, interp_mesh] as float32. Dimensionless CN values for Gaussian interpolation.
interp_mesh (int, optional) – Size of the coordination number interpolation mesh. Default: 5 (standard DFT-D3 uses a \(5 \times 5\) grid)

Raises:

ValueError – If parameter shapes are inconsistent or invalid
TypeError – If parameters are not jax.Array or have invalid dtypes

Notes

Parameters should use consistent units matching your coordinate system. Standard D3 parameters from the Grimme group use atomic units (Bohr for distances, Hartree \(\times\) Bohr\(^6\) for \(C_6\) coefficients).
Index 0 in all arrays is reserved for padding atoms (atomic number 0)
Valid atomic numbers range from 1 to max_z
The standard DFT-D3 implementation supports elements 1-94 (H to Pu)
Parameters should be float32 for efficiency

Examples

Create parameters from individual arrays:

>>> params = D3Parameters(
...     rcov=jnp.array([...]),
...     r4r2=jnp.array([...]),
...     c6ab=jnp.array([...]),
...     cn_ref=jnp.array([...]),
... )

c6ab: Array#

cn_ref: Array#

interp_mesh: int = 5#

property max_z: int#

Return the maximum atomic number supported by these parameters.

Returns:: Largest atomic number Z for which parameters are available; equal to rcov.shape[0] - 1 (index 0 is reserved for padding).
Return type:: int

r4r2: Array#

rcov: Array#

Internal Implementation#

These are low-level implementation functions that wrap the Warp kernels for JAX. For most use cases, prefer the high-level dftd3() wrapper above.

Neighbor Matrix Implementation#

nvalchemiops.jax.interactions.dispersion._dftd3._dftd3_nm_impl(positions, numbers, neighbor_matrix, covalent_radii, r4r2, c6_reference, coord_num_ref, a1, a2, s8, k1=16.0, k3=-4.0, s6=1.0, s5_smoothing_on=1e10, s5_smoothing_off=1e10, fill_value=None, batch_idx=None, cell=None, neighbor_matrix_shifts=None, compute_virial=False, num_systems=None)[source]#

Internal implementation for neighbor matrix format using jax_kernel wrappers.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates, float32 or float64.
numbers (jax.Array, shape (N,)) – Atomic numbers, int32.
neighbor_matrix (jax.Array, shape (N, max_neighbors)) – Neighbor indices in dense row format, int32. Row i lists the neighbor atom indices of atom i; unused slots are padded with values >= fill_value.
covalent_radii (jax.Array, shape (max_Z+1,)) – Covalent radii indexed by atomic number, float32.
r4r2 (jax.Array, shape (max_Z+1,)) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values, float32.
c6_reference (jax.Array, shape (max_Z+1, max_Z+1, interp_mesh, interp_mesh)) – C6 reference coefficients, float32.
coord_num_ref (jax.Array, shape (max_Z+1, max_Z+1, interp_mesh, interp_mesh)) – Coordination number reference grid, float32.
a1 (float) – Becke-Johnson damping parameter 1.
a2 (float) – Becke-Johnson damping parameter 2.
s8 (float) – C8 scaling factor.
k1 (float, optional) – CN counting steepness parameter. Default: 16.0.
k3 (float, optional) – CN interpolation Gaussian width parameter. Default: -4.0.
s6 (float, optional) – C6 scaling factor. Default: 1.0.
s5_smoothing_on (float, optional) – Distance where S5 switching begins. Default: 1e10.
s5_smoothing_off (float, optional) – Distance where S5 switching completes. Default: 1e10.
fill_value (int | None, optional) – Padding sentinel in neighbor_matrix. Defaults to num_atoms.
batch_idx (jax.Array | None, optional) – System index per atom, int32. Defaults to all zeros (single system).
cell (jax.Array | None, optional) – Unit cell lattice vectors [num_systems, 3, 3] for PBC.
neighbor_matrix_shifts (jax.Array | None, optional) – Integer unit cell shifts [num_atoms, max_neighbors, 3], int32, for PBC with neighbor-matrix format.
compute_virial (bool, optional) – If True, compute and return the virial tensor. Default: False.
num_systems (int | None, optional) – Number of systems in the batch. Required inside jax.jit when it cannot be inferred from cell or batch_idx.

Returns:

energy (jax.Array, shape (num_systems,)) – Per-system dispersion energy, float32.
forces (jax.Array, shape (N, 3)) – Atomic forces, float32.
coord_num (jax.Array, shape (N,)) – Coordination numbers, float32.
virial (jax.Array, shape (num_systems, 3, 3), optional) – Per-system virial tensor, float32. Returned only when compute_virial=True.

Return type:

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

Neighbor List Implementation#

nvalchemiops.jax.interactions.dispersion._dftd3._dftd3_nl_impl(positions, numbers, idx_j, neighbor_ptr, covalent_radii, r4r2, c6_reference, coord_num_ref, a1, a2, s8, k1=16.0, k3=-4.0, s6=1.0, s5_smoothing_on=1e10, s5_smoothing_off=1e10, batch_idx=None, cell=None, unit_shifts=None, compute_virial=False, num_systems=None)[source]#

Internal implementation for neighbor list format using jax_kernel wrappers.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic coordinates, float32 or float64.
numbers (jax.Array, shape (N,)) – Atomic numbers, int32.
idx_j (jax.Array, shape (num_pairs,)) – Destination atom indices in CSR neighbor list, int32.
neighbor_ptr (jax.Array, shape (N+1,)) – CSR row pointers indexing into idx_j, int32.
covalent_radii (jax.Array, shape (max_Z+1,)) – Covalent radii indexed by atomic number, float32.
r4r2 (jax.Array, shape (max_Z+1,)) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values, float32.
c6_reference (jax.Array, shape (max_Z+1, max_Z+1, interp_mesh, interp_mesh)) – C6 reference coefficients, float32.
coord_num_ref (jax.Array, shape (max_Z+1, max_Z+1, interp_mesh, interp_mesh)) – Coordination number reference grid, float32.
a1 (float) – Becke-Johnson damping parameter 1.
a2 (float) – Becke-Johnson damping parameter 2.
s8 (float) – C8 scaling factor.
k1 (float, optional) – CN counting steepness parameter. Default: 16.0.
k3 (float, optional) – CN interpolation Gaussian width parameter. Default: -4.0.
s6 (float, optional) – C6 scaling factor. Default: 1.0.
s5_smoothing_on (float, optional) – Distance where S5 switching begins. Default: 1e10.
s5_smoothing_off (float, optional) – Distance where S5 switching completes. Default: 1e10.
batch_idx (jax.Array | None, optional) – System index per atom, int32. Defaults to all zeros (single system).
cell (jax.Array | None, optional) – Unit cell lattice vectors [num_systems, 3, 3] for PBC.
unit_shifts (jax.Array | None, optional) – Integer unit cell shifts [num_pairs, 3], int32, for PBC with neighbor-list format.
compute_virial (bool, optional) – If True, compute and return the virial tensor. Default: False.
num_systems (int | None, optional) – Number of systems in the batch. Required inside jax.jit when it cannot be inferred from cell or batch_idx.

Returns:

energy (jax.Array, shape (num_systems,)) – Per-system dispersion energy, float32. Zero for empty systems and nonempty systems with no CSR edges.
forces (jax.Array, shape (N, 3)) – Atomic forces, float32. Shape (0, 3) when N == 0; for N > 0 with no CSR edges, shape (N, 3) and filled with zeros.
coord_num (jax.Array, shape (N,)) – Coordination numbers, float32. Shape (0,) when N == 0; for N > 0 with no CSR edges, shape (N,) and filled with zeros.
virial (jax.Array, shape (num_systems, 3, 3), optional) – Per-system virial tensor, float32. Returned only when compute_virial=True and zero for empty systems and zero-edge graphs.

Return type:

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

FourierD3: Particle-Mesh DFT-D3#

Evaluates a periodic particle-mesh variant of DFT-D3(BJ) in \(O(N \log N)\), using the same damped dispersion lattice sum with a modified coordination-number model, and with no pair cutoff on the dispersion sum.

Note

mesh_spacing reads the cell lengths into Python integers, so it cannot be used inside jax.jit. Pass mesh_dimensions explicitly when tracing.

nvalchemiops.jax.interactions.dispersion.fourier_dftd3(positions, numbers, a1, a2, s8, *, fourier_d3_params, cell, cutoff, mesh_dimensions=None, mesh_spacing=None, neighbor_matrix=None, neighbor_matrix_shifts=None, neighbor_list=None, neighbor_ptr=None, unit_shifts=None, fill_value=None, s6=1.0, spline_order=4, rank_chunk_size=None, batch_idx=None, compute_virial=False, num_systems=None)[source]#

Evaluate the DFT-D3(BJ) dispersion correction by particle-mesh summation.

The JAX counterpart of nvalchemiops.torch.interactions.dispersion.fourier_dftd3(); the arguments and their meanings are the same, minus device and the precomputed setup, which exists in the Torch binding to keep torch.linalg.inv out of a CUDA graph and has no counterpart here because jax.jit already hoists it.

Parameters:

positions (jax.Array, shape (N, 3)) – Atomic positions.
numbers (jax.Array, shape (N,)) – Atomic numbers. Zero marks a padding atom.
a1 (float) – Becke-Johnson damping parameters.
a2 (float) – Becke-Johnson damping parameters.
s8 (float) – Becke-Johnson damping parameters.
fourier_d3_params (FourierD3Parameters) – Separable coefficients covering every species present.
cell (jax.Array, shape (3, 3), (1, 3, 3) or (B, 3, 3)) – Lattice vectors as rows. Required: FourierD3 is periodic.
cutoff (float) – Coordination-number cutoff, in the same length unit as positions. Must equal the radius the neighbour list was built with.
mesh_dimensions (tuple[int, int, int] | None) – Exactly one is required. mesh_spacing reads cell lengths and so cannot be used inside jax.jit. A mesh derived from it is rounded up to a size whose only prime factors are 2, 3, 5 and 7, keeping the transform on cuFFT’s radix kernels; the result is never coarser than the spacing asked for. An explicit mesh_dimensions is used exactly as given.
mesh_spacing (float | None) – Exactly one is required. mesh_spacing reads cell lengths and so cannot be used inside jax.jit. A mesh derived from it is rounded up to a size whose only prime factors are 2, 3, 5 and 7, keeping the transform on cuFFT’s radix kernels; the result is never coarser than the spacing asked for. An explicit mesh_dimensions is used exactly as given.
neighbor_matrix –
Exactly one neighbour format, with its matching lattice images.

It must hold both directions of every pair, which is what the neighbour builders produce by default. FourierD3 accumulates each atom’s coordination number from its own row alone, so a list built with half_fill=True loses half of every atom’s coordination and yields wrong energies and non-conservative forces. This is not checked: the values are unreadable under jax.jit, which is the documented path, and the builders’ default already satisfies it.
neighbor_matrix_shifts –
Exactly one neighbour format, with its matching lattice images.

It must hold both directions of every pair, which is what the neighbour builders produce by default. FourierD3 accumulates each atom’s coordination number from its own row alone, so a list built with half_fill=True loses half of every atom’s coordination and yields wrong energies and non-conservative forces. This is not checked: the values are unreadable under jax.jit, which is the documented path, and the builders’ default already satisfies it.
neighbor_list –
Exactly one neighbour format, with its matching lattice images.

It must hold both directions of every pair, which is what the neighbour builders produce by default. FourierD3 accumulates each atom’s coordination number from its own row alone, so a list built with half_fill=True loses half of every atom’s coordination and yields wrong energies and non-conservative forces. This is not checked: the values are unreadable under jax.jit, which is the documented path, and the builders’ default already satisfies it.
neighbor_ptr –
Exactly one neighbour format, with its matching lattice images.

It must hold both directions of every pair, which is what the neighbour builders produce by default. FourierD3 accumulates each atom’s coordination number from its own row alone, so a list built with half_fill=True loses half of every atom’s coordination and yields wrong energies and non-conservative forces. This is not checked: the values are unreadable under jax.jit, which is the documented path, and the builders’ default already satisfies it.
unit_shifts –
Exactly one neighbour format, with its matching lattice images.

It must hold both directions of every pair, which is what the neighbour builders produce by default. FourierD3 accumulates each atom’s coordination number from its own row alone, so a list built with half_fill=True loses half of every atom’s coordination and yields wrong energies and non-conservative forces. This is not checked: the values are unreadable under jax.jit, which is the documented path, and the builders’ default already satisfies it.
fill_value (int, optional) – Padding sentinel for the dense format. Defaults to the atom count.
s6 (float, default=1.0) – Sixth-order scaling.
spline_order (int, default=4) – B-spline interpolation order, from 2 to 6. Accuracy at a fixed mesh improves with order: measured against a converged reference, orders 2 to 5 land roughly three orders of magnitude apart each way, so raising the order buys more than refining the mesh does. Order 3 is noticeably noisier than its neighbours; prefer an even order unless you have measured otherwise.
rank_chunk_size (int, optional) – Number of rank slots to hold on the mesh at once. None, the default, stages all rank slots in a single pass, which is fastest. The mesh and its transforms dominate the workspace and scale as num_systems * n_species * rank * nx * ny * nz, so a system with many species or a large retained rank can exceed device memory on a fine mesh. Setting this to k caps the resident slots at k and reduces that allocation by roughly rank / k, at the cost of one extra spread, forward and inverse transform, and gather per chunk. Only those reciprocal stages are chunked; the self-energy and the coordination chain rule need every slot at once and run once afterwards. The per-atom (N, rank) coefficient arrays stay at full width throughout. The result is unchanged to round-off: every reciprocal stage is a sum over slots with no coupling between them. Host-static, so it is safe under jax.jit – but the loop is unrolled at trace time, so the traced graph and the compile time grow linearly in rank / rank_chunk_size. Prefer the largest chunk that fits.
batch_idx (jax.Array, shape (N,), optional) – System index per atom.
compute_virial (bool, default=False) – Whether to return the virial.
num_systems (int, optional) – Number of systems. Inferred from cell when omitted.

Returns:

energy (jax.Array, shape (num_systems,))
forces (jax.Array, shape (N, 3))
virial (jax.Array, shape (num_systems, 3, 3)) – Returned only when compute_virial is set. Follows the repository convention in docs/userguide/about/conventions.md: the negative derivative of the energy with respect to the affine displacement, \(W = -\partial E/\partial u\). The tensile-positive Cauchy stress is \(\sigma = -W/V\).

Notes

fourier_d3_params must cover every element present. This is a caller precondition, not a checked one, as in dftd3(): verifying it would read device memory on every call. An uncovered element is grouped with the padding and dropped from the dispersion sum, so the result is finite and wrong rather than an error. Build the decomposition from the species actually present.

The returned energy is not differentiable. The kernels are launched with enable_backward=False and no VJP or JVP rule is registered on top of them, so jax.grad of the energy raises ValueError: ... cannot be differentiated rather than reproducing forces. Use the returned forces and virial, which are analytic derivatives of the same energy. This matches the real-space dftd3(), and it is what keeps the call traceable under jax.jit().

class nvalchemiops.jax.interactions.dispersion.FourierD3Parameters(rcov, sqrt_q, cn_ref, v_q, eigs, species_map, max_relative_error)[source]#

Separable dispersion coefficients for the species in a system.

The JAX counterpart of nvalchemiops.torch.interactions.dispersion.FourierD3Parameters, and like it it stores no damping parameters: those are call-time arguments, so one instance is valid for every functional parametrisation.

Parameters:

rcov (Array)
sqrt_q (Array)
cn_ref (Array)
v_q (Array)
eigs (Array)
species_map (Array)
max_relative_error (float)

rcov#

Covalent radii indexed by atomic number.

Type:: jax.Array, shape (max_z + 1,)

sqrt_q#

Square root of the quadrupole-to-dipole ratio, per species channel.

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

cn_ref#

Reference coordination numbers, negative in unused slots.

Type:: jax.Array, shape (n_species, n_ref)

v_q#

Eigenvectors of the decomposed reference tensor.

Type:: jax.Array, shape (n_species, n_ref, rank)

eigs#

Eigenvalues, which may be negative.

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

species_map#

Atomic number to channel index; -1 where a species is not covered.

Type:: jax.Array, shape (max_z + 1,), dtype=int32

max_relative_error#

Largest relative error the truncated decomposition makes.

Type:: float

classmethod from_tables(rcov, r4r2, c6ab, cn_ref, species, tol=1e-4, max_rank=None, dtype=jnp.float64)[source]#

Decompose Grimme’s reference tables for a given set of species.

Parameters mirror the Torch implementation; see that class for details.

Parameters:

tol (float)
max_rank (int | None)

Return type:

FourierD3Parameters

property n_species: int#: Number of species channels.

property rank: int#: Number of retained rank slots.

uncovered_species(numbers)[source]#

Atomic numbers in numbers that this decomposition does not cover.

Coverage is a caller precondition: fourier_dftd3 does not check it eagerly or under jax.jit, because reading the answer back is a device sync eagerly and impossible while tracing. An uncovered element is grouped with the padding and dropped from the dispersion sum, so the energy comes back finite and wrong.

Call this once when the parameters or the composition change, on concrete arrays – not on tracers, and not inside a jitted step. It matches the Torch binding’s method of the same name, so a check written against either behaves the same.

Parameters:: numbers (array_like, shape (N,)) – Atomic numbers of the system. Zero marks padding and is ignored.
Returns:: Sorted uncovered atomic numbers; empty when the parameters cover the system.
Return type:: list of int

nvalchemiops.jax.interactions.dispersion: Dispersion Corrections#

High-Level Interface#

DFT-D3(BJ) Dispersion Corrections#

Data Structures#

Internal Implementation#

Neighbor Matrix Implementation#

Neighbor List Implementation#

FourierD3: Particle-Mesh DFT-D3#

`nvalchemiops.jax.interactions.dispersion`: Dispersion Corrections#