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

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

PyTorch bindings for dispersion corrections.

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.

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

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:

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.

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)

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. Both neighbour formats are supported, as above.

Two differences from dftd3() are worth noting before use:

cell is required. FourierD3 evaluates a periodic sum; there is no open-boundary path. Molecules in vacuum should continue to use dftd3().
cutoff has no default, and must equal the radius the neighbour list was built with. The coordination-number function is constructed to reach zero exactly there.

nvalchemiops.torch.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, batch_idx=None, compute_virial=False, num_systems=None, rank_chunk_size=None, setup=None, device=None)[source]#

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

The dispersion sum itself is untruncated. cutoff applies only to the coordination-number neighbour list, which a machine-learned force field already builds for its own descriptors.

Parameters:

positions (torch.Tensor, shape (N, 3)) – Atomic positions.
numbers (torch.Tensor, shape (N,)) – Atomic numbers. Zero marks a padding atom.
a1 (float) – Becke-Johnson damping parameters for the exchange-correlation functional in use.
a2 (float) – Becke-Johnson damping parameters for the exchange-correlation functional in use.
s8 (float) – Becke-Johnson damping parameters for the exchange-correlation functional in use.
fourier_d3_params (FourierD3Parameters) – Separable coefficients covering every species present. Coverage is a caller precondition and is not checked; an uncovered element is dropped from the sum.
cell (torch.Tensor, 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. No default, because the DFT-D3 tables are conventionally atomic units and a value meant as 6 Angstrom would silently act as 6 Bohr. This must equal the radius the neighbour list was built with: the counting function is constructed to reach zero exactly there, and a mismatch reintroduces the truncation discontinuity it exists to remove.
mesh_dimensions (tuple[int, int, int], optional) – Mesh size. Exactly one of this and mesh_spacing must be given, unless a setup is passed – it already carries the mesh it was built for, and then both are optional.
mesh_spacing (float, optional) – Target spacing, in the same unit as cell. Sized from the largest cell in a batch, then rounded up to a size whose only prime factors are 2, 3, 5 and 7, so the transform stays on cuFFT’s radix kernels rather than falling back to Bluestein’s algorithm. The resulting mesh is therefore never coarser than the spacing asked for, and may be a little finer. An explicit mesh_dimensions is used exactly as given. Reads cell lengths into Python integers, so pass explicit mesh_dimensions when tracing.
neighbor_matrix (torch.Tensor, optional) – Dense padded neighbour indices and their lattice images.
neighbor_matrix_shifts (torch.Tensor, optional) – Dense padded neighbour indices and their lattice images.
neighbor_list (torch.Tensor, optional) –
CSR neighbour list and its lattice images. Exactly one format must be supplied.

Whichever format is used 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: detecting it costs a reduction over the whole list on every call, and the builders’ default already satisfies it.
neighbor_ptr (torch.Tensor, optional) –
CSR neighbour list and its lattice images. Exactly one format must be supplied.

Whichever format is used 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: detecting it costs a reduction over the whole list on every call, and the builders’ default already satisfies it.
unit_shifts (torch.Tensor, optional) –
CSR neighbour list and its lattice images. Exactly one format must be supplied.

Whichever format is used 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: detecting it costs a reduction over the whole list on every call, 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; unity for every common parametrisation.
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.
batch_idx (torch.Tensor, shape (N,), optional) – System index per atom. Atoms must be grouped by system.
setup (FourierD3Setup, optional) – Cell- and mesh-derived quantities from FourierD3Setup.build(), reused across steps. A performance cache only: it saves a matrix inversion and a set of spline moduli per call, and the call works without it under eager, torch.compile and CUDA graph capture alike. Batch size, species count, precision, device and mesh are checked; keeping it consistent with the cell is yours to do, since comparing cell values would read device memory. It also supplies the mesh and the spline order, so mesh_dimensions and mesh_spacing may be omitted and the usual “exactly one of them” rule does not apply. A mesh_dimensions that disagrees with the setup raises, and so does any mesh_spacing, rather than either being silently ignored. spline_order is the one exception: it has a default, so a value passed alongside a setup cannot be told apart from the default and the setup’s own order is used.
compute_virial (bool, default=False) – Whether to return the virial.
num_systems (int, optional) – Number of systems, inferred from cell when omitted.
rank_chunk_size (int, optional) – Number of rank slots to hold on the mesh at once. None, the default, processes 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. Must be a host-side Python integer, since it determines the number of kernel launches.
device (str, optional) – Warp device string. Inferred from positions when omitted.

Returns:

energy (torch.Tensor, shape (num_systems,))
forces (torch.Tensor, shape (N, 3))
virial (torch.Tensor, 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\), with deformation applied as \(R' = R(I + u)\) and \(C' = C(I + u)\). The tensile-positive Cauchy stress is \(\sigma = -W/V\). Matches dftd3().

Return type:

tuple[Tensor, …]

Notes

The returned energy is not differentiable. The kernels are launched with enable_backward=False and no torch.library.register_autograd rule is registered, so it comes back with requires_grad=False and grad_fn=None, detached from positions. torch.autograd.grad on it raises rather than reproducing forces. Use the returned forces and virial, which are analytic derivatives of the same energy. This matches dftd3(), and it is what keeps the pipeline capturable into a CUDA graph.

Energies are reduced with atomic adds, whose summation order varies between launches, so repeated identical calls can differ in the last bit.

Examples

>>> energy, forces = fourier_dftd3(
...     positions, numbers, a1=0.4289, a2=4.4407, s8=0.7875,
...     fourier_d3_params=params, cell=cell, cutoff=11.34,
...     mesh_dimensions=(32, 32, 32),
...     neighbor_list=pairs, neighbor_ptr=pointer, unit_shifts=shifts,
... )

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.torch.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 (torch.Tensor) – Covalent radii [max_Z+1] as float32 or float64. Units should be consistent with position coordinates. Index 0 is reserved for padding; valid atomic numbers are 1 to max_Z.
r4r2 (torch.Tensor) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values [max_Z+1] as float32 or float64. Dimensionless ratio used for computing \(C_8\) coefficients from \(C_6\) values.
c6ab (torch.Tensor) – \(C_6\) reference coefficients [max_Z+1, max_Z+1, interp_mesh, interp_mesh] as float32 or float64. Units are energy \(\times\) distance\(^6\). Indexed by atomic numbers and coordination number reference indices.
cn_ref (torch.Tensor) – Coordination number reference grid [max_Z+1, max_Z+1, interp_mesh, interp_mesh] as float32 or float64. 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 torch.Tensor 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 can be float32 or float64; they will be converted to float32 during computation for efficiency

Examples

Create parameters from individual tensors:

>>> params = D3Parameters(
...     rcov=torch.rand(95),  # 94 elements + padding
...     r4r2=torch.rand(95),
...     c6ab=torch.rand(95, 95, 5, 5),
...     cn_ref=torch.rand(95, 95, 5, 5),
... )

Create from a dictionary (e.g., loaded from file):

>>> state_dict = torch.load("dftd3_parameters.pt", weights_only=True)
>>> params = D3Parameters(
...     rcov=state_dict["rcov"],
...     r4r2=state_dict["r4r2"],
...     c6ab=state_dict["c6ab"],
...     cn_ref=state_dict["cn_ref"],
... )

c6ab: Tensor#

cn_ref: Tensor#

property device: device#: Device where parameters are stored.

interp_mesh: int = 5#

property max_z: int#: Maximum atomic number supported by these parameters.

r4r2: Tensor#

rcov: Tensor#

to(device=None, dtype=None)[source]#

Move all parameters to the specified device and/or convert to specified dtype.

Parameters:

device (str or torch.device or None, optional) – Target device (e.g., ‘cpu’, ‘cuda’, ‘cuda:0’). If None, keeps current device.
dtype (torch.dtype or None, optional) – Target dtype (e.g., torch.float32, torch.float64). If None, keeps current dtype.

Returns:

New instance with parameters on the target device and/or dtype

Return type:

D3Parameters

Examples

Move to GPU:

>>> params_gpu = params.to(device='cuda')

Convert to float32:

>>> params_f32 = params.to(dtype=torch.float32)

Move to GPU and convert to float32:

>>> params_gpu_f32 = params.to(device='cuda', dtype=torch.float32)

class nvalchemiops.torch.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.

Holds the low-rank factors of Grimme’s reference tensor together with the element data the evaluation needs. Every field is a function of the reference tables, the species present and the requested tolerance.

Notably absent are the damping parameters. They are call-time arguments to fourier_dftd3() and are not stored here, so a single instance is valid for every functional parametrisation and there is no way for a stored copy to disagree with the values actually used.

Parameters:

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

rcov#

Covalent radii indexed by atomic number.

Type:: torch.Tensor, shape (max_z + 1,)

sqrt_q#

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

Type:: torch.Tensor, shape (n_species,)

cn_ref#

Reference coordination numbers, negative in unused slots.

Type:: torch.Tensor, shape (n_species, n_ref)

v_q#

Eigenvectors of the decomposed reference tensor.

Type:: torch.Tensor, shape (n_species, n_ref, rank)

eigs#

Eigenvalues, which may be negative.

Type:: torch.Tensor, shape (rank,)

species_map#

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

Type:: torch.Tensor, shape (max_z + 1,), dtype=int32

max_relative_error#

Largest relative error the truncated decomposition makes on the reference tensor.

Type:: float

property device: device#: Device the parameters live on.

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

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

Parameters:

rcov (torch.Tensor, shape (max_z + 1,)) – Covalent radii and the quadrupole-to-dipole ratios, indexed by atomic number.
r4r2 (torch.Tensor, shape (max_z + 1,)) – Covalent radii and the quadrupole-to-dipole ratios, indexed by atomic number.
c6ab (torch.Tensor, shape (max_z + 1, max_z + 1, n_ref, n_ref)) – Reference dispersion coefficients and coordination numbers.
cn_ref (torch.Tensor, shape (max_z + 1, max_z + 1, n_ref, n_ref)) – Reference dispersion coefficients and coordination numbers.
species (Sequence[int]) – Atomic numbers present. Order and duplicates are ignored.
tol (float, default=1e-4) – Target maximum relative error of the reconstructed reference tensor.
max_rank (int, optional) – Ceiling on the retained rank. When it prevents tol from being met the result is still returned and max_relative_error reports what was achieved.
device (optional) – Device for the returned tensors. Defaults to that of rcov.
dtype (torch.dtype, default=torch.float64) – Floating dtype for the returned tensors.

Return type:

FourierD3Parameters

property n_species: int#: Number of species channels.

property rank: int#: Number of retained rank slots, and so of mesh channels per species.

to(device=None, dtype=None)[source]#

Return a copy on the given device and floating dtype.

species_map stays integral regardless of dtype.

Return type:: FourierD3Parameters

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 in any execution mode, because reading the answer back would synchronise on every step. An uncovered element is grouped with the padding and dropped from the dispersion sum, so the energy comes back finite and wrong rather than raising.

Call this once when the parameters or the composition change — at setup, or in a test — rather than inside a dynamics loop. It synchronises, and behaves the same eagerly, under torch.compile and under JAX, so a check written against it does not shift with the execution mode.

Parameters:: numbers (torch.Tensor, 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

Examples

>>> missing = params.uncovered_species(numbers)
>>> if missing:
...     raise ValueError(f"rebuild the decomposition to cover {missing}")

FourierD3Setup is a different kind of object: an optional cache of cell- and mesh-derived quantities, saving a matrix inversion and the spline moduli per call. It is not required for any execution mode. Batch size, species count, precision, device and mesh are checked on every call; keeping it consistent with the cell is the caller’s responsibility, as with the PME and multipole caches.

class nvalchemiops.torch.interactions.dispersion.FourierD3Setup(cell_inv_grouped, volumes, k_matrix, moduli_x, moduli_y, moduli_z, mesh_dimensions, spline_order)[source]#

Cell- and mesh-derived quantities that do not change from step to step.

An optional performance cache: it saves a matrix inversion and the spline moduli per call. fourier_dftd3 works without it, including under torch.compile and CUDA graph capture.

Keeping it consistent with the cell is the caller’s responsibility, as with the PME and multipole caches. The derived quantities stand in for the cell everywhere except the Cartesian image shifts, which still come from the cell passed to the call, so a stale setup mixes two cells in one evaluation with nothing raised. Rebuild it whenever the cell changes. validate_for() checks only metadata – batch size, species count, dtype, device and mesh – which is cheap and works under compilation; it does not compare cell values.

Parameters:

cell_inv_grouped (Tensor)
volumes (Tensor)
k_matrix (Tensor)
moduli_x (Tensor)
moduli_y (Tensor)
moduli_z (Tensor)
mesh_dimensions (tuple[int, int, int])
spline_order (int)

cell_inv_grouped#

Transpose of the inverse cell, repeat-interleaved across species channels.

Type:: torch.Tensor, shape (B * n_species, 3, 3)

volumes#

Cell volume per system.

Type:: torch.Tensor, shape (B,)

k_matrix#

2 * pi * inverse(cell) per system.

Type:: torch.Tensor, shape (B, 3, 3)

moduli_x, moduli_y, moduli_z

B-spline attenuation per mesh axis.

Type:: torch.Tensor

mesh_dimensions#

The mesh these were built for.

Type:: tuple[int, int, int]

spline_order#

The spline order these were built for.

Type:: int

classmethod build(cell, n_species, mesh_dimensions, spline_order=4)[source]#

Derive the reusable quantities from a cell and a mesh.

Parameters:

cell (torch.Tensor, shape (3, 3), (1, 3, 3) or (B, 3, 3)) – Lattice vectors as rows.
n_species (int) – Number of species channels, from FourierD3Parameters.n_species.
mesh_dimensions (tuple[int, int, int]) – Mesh size.
spline_order (int, default=4) – B-spline order.

Return type:

FourierD3Setup

validate_for(cells, n_species, mesh_dimensions, mesh_spacing=None)[source]#

Check the metadata this setup was built for: batch, species, dtype, device, mesh.

Not a full check. Cell values are never compared – that reads device memory, which graph capture forbids – so a setup built for a different cell passes here. The derived quantities stand in for the cell everywhere except the Cartesian image shifts, which come from the call’s cell, so a stale setup mixes two cells in one evaluation and returns a finite, incoherent answer. Rebuild it whenever the cell changes; the same contract as the PME and multipole caches.

Internal Custom Operators#

These are low-level custom operators that wrap the Warp kernels. For most use cases, prefer the high-level dftd3() wrapper above. These operators are exposed for advanced users who need fine-grained control or torch.compile compatibility.

Non-Periodic Systems#

nvalchemiops.torch.interactions.dispersion._dftd3._dftd3_matrix_op(positions, numbers, neighbor_matrix, covalent_radii, r4r2, c6_reference, coord_num_ref, a1, a2, s8, energy, forces, coord_num, virial, k1=16.0, k3=-4.0, s6=1.0, s5_smoothing_on=1e10, s5_smoothing_off=1e10, fill_value=None, batch_idx=None, device=None)[source]#

Internal custom op for DFT-D3(BJ) dispersion energy and forces computation (non-PBC, neighbor matrix format).

This is a low-level custom operator that performs DFT-D3(BJ) dispersion calculations using Warp kernels for non-periodic systems with neighbor matrix format. Output tensors must be pre-allocated by the caller and are modified in-place. For most use cases, prefer the higher-level dftd3() wrapper function instead of calling this method directly.

This function is torch.compile compatible.

Parameters:

positions (torch.Tensor, shape (num_atoms, 3)) – Atomic coordinates as float32 or float64, in consistent distance units (conventionally Bohr)
numbers (torch.Tensor, shape (num_atoms), dtype=int32) – Atomic numbers
neighbor_matrix (torch.Tensor, shape (num_atoms, max_neighbors), dtype=int32) – Neighbor indices 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). Padding atoms (numbers[i] == 0) are skipped.
covalent_radii (torch.Tensor, shape (max_Z+1), dtype=float32) – Covalent radii indexed by atomic number, in same units as positions
r4r2 (torch.Tensor, shape (max_Z+1), dtype=float32) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values for C8 computation (dimensionless)
c6_reference (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – C6 reference values in energy x distance^6 units
coord_num_ref (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – CN reference grid (dimensionless)
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) – C8 term scaling factor (functional-dependent, dimensionless)
energy (torch.Tensor, shape (num_systems,), dtype=float32) – OUTPUT: Total dispersion energy. Must be pre-allocated. Units are energy (Hartree when using standard D3 parameters).
forces (torch.Tensor, shape (num_atoms, 3), dtype=float32) – OUTPUT: Atomic forces. Must be pre-allocated. Units are energy/distance (Hartree/Bohr when using standard D3 parameters).
coord_num (torch.Tensor, shape (num_atoms,), dtype=float32) – OUTPUT: Coordination numbers (dimensionless). Must be pre-allocated.
virial (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – OUTPUT: Virial tensor (remains zeros for non-PBC). Must be pre-allocated.
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) – C6 term scaling factor (typically 1.0, dimensionless)
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
fill_value (int | None, optional) – Value indicating padding in neighbor_matrix. If None, defaults to num_atoms.
batch_idx (torch.Tensor, shape (num_atoms,), dtype=int32, optional) – Batch indices. If None, all atoms are in a single system (batch 0).
device (str, optional) – Warp device string (e.g., ‘cuda:0’, ‘cpu’). If None, inferred from positions.

Returns:

None
Modifies input tensors in-place (energy, forces, coord_num, virial. Output)
buffers are reset to zero on every successful call, including empty
systems. The non-PBC virial remains zero.

Return type:

None

Notes

All input tensors should use consistent units. Standard D3 parameters use atomic units (Bohr for distances, Hartree for energy).
Float32 or float64 precision for positions; outputs always float32
Padding atoms indicated by numbers[i] == 0
Two-body only: Computes pairwise C6 and C8 dispersion terms; three-body Axilrod-Teller-Muto (ATM/C9) terms are not included
For PBC calculations, use _dftd3_matrix_pbc_op() instead

See also

dftd3(): Higher-level wrapper that handles allocation
_dftd3_matrix_pbc_op(): PBC variant with neighbor matrix format

nvalchemiops.torch.interactions.dispersion._dftd3._dftd3_op(positions, numbers, idx_j, neighbor_ptr, covalent_radii, r4r2, c6_reference, coord_num_ref, a1, a2, s8, energy, forces, coord_num, virial, k1=16.0, k3=-4.0, s6=1.0, s5_smoothing_on=1e10, s5_smoothing_off=1e10, batch_idx=None, device=None)[source]#

Internal custom op for DFT-D3(BJ) using CSR neighbor list format (non-PBC).

This is a low-level custom operator that performs DFT-D3(BJ) dispersion calculations using CSR (Compressed Sparse Row) neighbor list format with idx_j (destination indices) and neighbor_ptr (row pointers) for non-periodic systems. Output tensors must be pre-allocated by the caller and are modified in-place. For most use cases, prefer the higher-level dftd3() wrapper function instead of calling this method directly.

This function is torch.compile compatible.

Parameters:

positions (torch.Tensor, shape (num_atoms, 3)) – Atomic coordinates as float32 or float64
numbers (torch.Tensor, shape (num_atoms), dtype=int32) – Atomic numbers
idx_j (torch.Tensor, shape (num_edges,), dtype=int32) – Destination atom indices (flattened neighbor list in CSR format)
neighbor_ptr (torch.Tensor, shape (num_atoms+1,), dtype=int32) – CSR row pointers where neighbor_ptr[i]:neighbor_ptr[i+1] gives neighbors of atom i
covalent_radii (torch.Tensor, shape (max_Z+1), dtype=float32) – Covalent radii indexed by atomic number
r4r2 (torch.Tensor, shape (max_Z+1), dtype=float32) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values
c6_reference (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – C6 reference values
coord_num_ref (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – CN reference grid
a1 (float) – Becke-Johnson damping parameter 1
a2 (float) – Becke-Johnson damping parameter 2
s8 (float) – C8 term scaling factor
energy (torch.Tensor, shape (num_systems,), dtype=float32) – OUTPUT: Total dispersion energy
forces (torch.Tensor, shape (num_atoms, 3), dtype=float32) – OUTPUT: Atomic forces
coord_num (torch.Tensor, shape (num_atoms,), dtype=float32) – OUTPUT: Coordination numbers
virial (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – OUTPUT: Virial tensor (remains zeros for non-PBC). Must be pre-allocated.
k1 (float, optional) – CN counting function steepness parameter
k3 (float, optional) – CN interpolation Gaussian width parameter
s6 (float, optional) – C6 term scaling factor
s5_smoothing_on (float, optional) – Distance where S5 switching begins
s5_smoothing_off (float, optional) – Distance where S5 switching completes
batch_idx (torch.Tensor, shape (num_atoms,), dtype=int32, optional) – Batch indices
device (str, optional) – Warp device string

Returns:

None
Modifies input tensors in-place (energy, forces, coord_num, virial. Output)
buffers are reset to zero on every successful call, including empty systems
and CSR graphs with no edges. The non-PBC virial is reset but not computed.

Return type:

None

Notes

All input tensors should use consistent units. Standard D3 parameters use atomic units (Bohr for distances, Hartree for energy).
Float32 or float64 precision for positions; outputs always float32
Padding atoms indicated by numbers[i] == 0
Two-body only: Computes pairwise C6 and C8 dispersion terms; three-body Axilrod-Teller-Muto (ATM/C9) terms are not included
For PBC calculations, use _dftd3_pbc_op() instead

See also

dftd3(): Higher-level wrapper that handles allocation
_dftd3_pbc_op(): PBC variant with CSR neighbor list format

Periodic Boundary Conditions (PBC)#

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

Internal custom op for DFT-D3(BJ) dispersion energy and forces computation (PBC, neighbor matrix format).

This is a low-level custom operator that performs DFT-D3(BJ) dispersion calculations using Warp kernels for periodic systems with neighbor matrix format. Output tensors must be pre-allocated by the caller and are modified in-place. For most use cases, prefer the higher-level dftd3() wrapper function instead of calling this method directly.

This function is torch.compile compatible.

Parameters:

positions (torch.Tensor, shape (num_atoms, 3)) – Atomic coordinates as float32 or float64, in consistent distance units (conventionally Bohr)
numbers (torch.Tensor, shape (num_atoms), dtype=int32) – Atomic numbers
neighbor_matrix (torch.Tensor, shape (num_atoms, max_neighbors), dtype=int32) – Neighbor indices 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). Padding atoms (numbers[i] == 0) are skipped.
cell (torch.Tensor, shape (num_systems, 3, 3), dtype=float32 or float64) – Unit cell lattice vectors for PBC, in same dtype and units as positions.
neighbor_matrix_shifts (torch.Tensor, shape (num_atoms, max_neighbors, 3), dtype=int32) – Integer unit cell shifts for PBC.
covalent_radii (torch.Tensor, shape (max_Z+1), dtype=float32) – Covalent radii indexed by atomic number, in same units as positions
r4r2 (torch.Tensor, shape (max_Z+1), dtype=float32) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values for C8 computation (dimensionless)
c6_reference (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – C6 reference values in energy x distance^6 units
coord_num_ref (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – CN reference grid (dimensionless)
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) – C8 term scaling factor (functional-dependent, dimensionless)
energy (torch.Tensor, shape (num_systems,), dtype=float32) – OUTPUT: Total dispersion energy. Must be pre-allocated. Units are energy (Hartree when using standard D3 parameters).
forces (torch.Tensor, shape (num_atoms, 3), dtype=float32) – OUTPUT: Atomic forces. Must be pre-allocated. Units are energy/distance (Hartree/Bohr when using standard D3 parameters).
coord_num (torch.Tensor, shape (num_atoms,), dtype=float32) – OUTPUT: Coordination numbers (dimensionless). Must be pre-allocated.
virial (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – OUTPUT: Virial tensor. Must be pre-allocated. Units are energy (Hartree when using standard D3 parameters).
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) – C6 term scaling factor (typically 1.0, dimensionless)
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
fill_value (int | None, optional) – Value indicating padding in neighbor_matrix. If None, defaults to num_atoms.
batch_idx (torch.Tensor, shape (num_atoms,), dtype=int32, optional) – Batch indices. If None, all atoms are in a single system (batch 0).
compute_virial (bool, optional) – If True, compute virial tensor. Default: False
device (str, optional) – Warp device string (e.g., ‘cuda:0’, ‘cpu’). If None, inferred from positions.

Returns:

None
Modifies input tensors in-place (energy, forces, coord_num, virial. Output)
buffers are reset to zero on every successful call, including empty
systems. compute_virial=True populates virial; otherwise it remains
reset to zero.

Return type:

None

Notes

All input tensors should use consistent units. Standard D3 parameters use atomic units (Bohr for distances, Hartree for energy).
Float32 or float64 precision for positions and cell; outputs always float32
Padding atoms indicated by numbers[i] == 0
Two-body only: Computes pairwise C6 and C8 dispersion terms; three-body Axilrod-Teller-Muto (ATM/C9) terms are not included
Bulk stress tensor can be obtained by dividing virial by system volume.
For non-PBC calculations, use _dftd3_matrix_op() instead

See also

dftd3(): Higher-level wrapper that handles allocation
_dftd3_matrix_op(): Non-PBC variant with neighbor matrix format

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

Internal custom op for DFT-D3(BJ) using CSR neighbor list format (PBC).

This is a low-level custom operator that performs DFT-D3(BJ) dispersion calculations using CSR (Compressed Sparse Row) neighbor list format with idx_j (destination indices) and neighbor_ptr (row pointers) for periodic systems. Output tensors must be pre-allocated by the caller and are modified in-place. For most use cases, prefer the higher-level dftd3() wrapper function instead of calling this method directly.

This function is torch.compile compatible.

Parameters:

positions (torch.Tensor, shape (num_atoms, 3)) – Atomic coordinates as float32 or float64
numbers (torch.Tensor, shape (num_atoms), dtype=int32) – Atomic numbers
idx_j (torch.Tensor, shape (num_edges,), dtype=int32) – Destination atom indices (flattened neighbor list in CSR format)
neighbor_ptr (torch.Tensor, shape (num_atoms+1,), dtype=int32) – CSR row pointers where neighbor_ptr[i]:neighbor_ptr[i+1] gives neighbors of atom i
cell (torch.Tensor, shape (num_systems, 3, 3), dtype=float32 or float64) – Unit cell lattice vectors for PBC, in same dtype and units as positions.
unit_shifts (torch.Tensor, shape (num_edges, 3), dtype=int32) – Integer unit cell shifts for PBC
covalent_radii (torch.Tensor, shape (max_Z+1), dtype=float32) – Covalent radii indexed by atomic number
r4r2 (torch.Tensor, shape (max_Z+1), dtype=float32) – \(\langle r^4 \rangle / \langle r^2 \rangle\) expectation values
c6_reference (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – C6 reference values
coord_num_ref (torch.Tensor, shape (max_Z+1, max_Z+1, 5, 5), dtype=float32) – CN reference grid
a1 (float) – Becke-Johnson damping parameter 1
a2 (float) – Becke-Johnson damping parameter 2
s8 (float) – C8 term scaling factor
energy (torch.Tensor, shape (num_systems,), dtype=float32) – OUTPUT: Total dispersion energy
forces (torch.Tensor, shape (num_atoms, 3), dtype=float32) – OUTPUT: Atomic forces
coord_num (torch.Tensor, shape (num_atoms,), dtype=float32) – OUTPUT: Coordination numbers
virial (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – OUTPUT: Virial tensor. Must be pre-allocated. Units are energy (Hartree when using standard D3 parameters).
k1 (float, optional) – CN counting function steepness parameter
k3 (float, optional) – CN interpolation Gaussian width parameter
s6 (float, optional) – C6 term scaling factor
s5_smoothing_on (float, optional) – Distance where S5 switching begins
s5_smoothing_off (float, optional) – Distance where S5 switching completes
batch_idx (torch.Tensor, shape (num_atoms,), dtype=int32, optional) – Batch indices
compute_virial (bool, optional) – If True, compute virial tensor. Default: False
device (str, optional) – Warp device string

Returns:

None
Modifies input tensors in-place (energy, forces, coord_num, virial. Output)
buffers are reset to zero on every successful call, including empty systems
and CSR graphs with no edges. compute_virial=True populates virial;
otherwise it remains reset to zero.

Return type:

None

Notes

All input tensors should use consistent units. Standard D3 parameters use atomic units (Bohr for distances, Hartree for energy).
Float32 or float64 precision for positions and cell; outputs always float32
Padding atoms indicated by numbers[i] == 0
Two-body only: Computes pairwise C6 and C8 dispersion terms; three-body Axilrod-Teller-Muto (ATM/C9) terms are not included
Bulk stress tensor can be obtained by dividing virial by system volume.
For non-PBC calculations, use _dftd3_op() instead

See also

dftd3(): Higher-level wrapper that handles allocation
_dftd3_op(): Non-PBC variant with CSR neighbor list format

FourierD3 Operators#

nvalchemiops.torch.interactions.dispersion._fourier_dftd3._fd3_prologue_op(positions, numbers, species_index, cartesian_shifts, covalent_radii, cnref, v_q, cutoff, coord_num, c6, dc6_dcn, neighbor_list=None, neighbor_ptr=None, neighbor_matrix=None, fill_value=None, device=None)[source]#

Internal op for the real-space passes: coordination numbers and coefficients.

Parameters:

positions (Tensor)
numbers (Tensor)
species_index (Tensor)
cartesian_shifts (Tensor)
covalent_radii (Tensor)
cnref (Tensor)
v_q (Tensor)
cutoff (float)
coord_num (Tensor)
c6 (Tensor)
dc6_dcn (Tensor)
neighbor_list (Tensor | None)
neighbor_ptr (Tensor | None)
neighbor_matrix (Tensor | None)
fill_value (int | None)
device (str | None)

Return type:

None

nvalchemiops.torch.interactions.dispersion._fourier_dftd3._fd3_spread_op(positions, c6, group_idx, cell_inv_t, spline_order, rank, mesh, device=None)[source]#

Internal op spreading the separable coefficients onto the mesh.

Parameters:

positions (Tensor)
c6 (Tensor)
group_idx (Tensor)
cell_inv_t (Tensor)
spline_order (int)
rank (int)
mesh (Tensor)
device (str | None)

Return type:

None

nvalchemiops.torch.interactions.dispersion._fourier_dftd3._fd3_kspace_op(mesh_fft, k_matrix, moduli_x, moduli_y, moduli_z, volumes, sqrt_q, eigs, s6, s8, a1, a2, mesh_nx, mesh_ny, mesh_nz, n_species, rank, energy, cotangent, virial, compute_virial=False, device=None)[source]#

Internal op contracting the transformed mesh against the dispersion kernel.

Parameters:

mesh_fft (Tensor)
k_matrix (Tensor)
moduli_x (Tensor)
moduli_y (Tensor)
moduli_z (Tensor)
volumes (Tensor)
sqrt_q (Tensor)
eigs (Tensor)
s6 (float)
s8 (float)
a1 (float)
a2 (float)
mesh_nx (int)
mesh_ny (int)
mesh_nz (int)
n_species (int)
rank (int)
energy (Tensor)
cotangent (Tensor)
virial (Tensor)
compute_virial (bool)
device (str | None)

Return type:

None

nvalchemiops.torch.interactions.dispersion._fourier_dftd3._fd3_gather_op(potential, positions, group_idx, cell_inv_t, c6, spline_order, rank, d_energy_d_c6, forces, device=None)[source]#

Internal op for the mesh gather: the coefficient derivative and the direct force.

This is the only stage that touches the mesh, so it is the only one that runs per rank chunk. Both outputs are accumulated into rather than overwritten, because a chunk contributes a slice of d_energy_d_c6 and a share of the force.

Parameters:

potential (Tensor)
positions (Tensor)
group_idx (Tensor)
cell_inv_t (Tensor)
c6 (Tensor)
spline_order (int)
rank (int)
d_energy_d_c6 (Tensor)
forces (Tensor)
device (str | None)

Return type:

None

nvalchemiops.torch.interactions.dispersion._fourier_dftd3._fd3_finalise_op(positions, numbers, species_index, batch_idx, c6, dc6_dcn, cartesian_shifts, covalent_radii, sqrt_q, eigs, s6, s8, a1, a2, cutoff, d_energy_d_c6, d_energy_d_cn, energy, forces, virial, neighbor_list=None, neighbor_ptr=None, neighbor_matrix=None, fill_value=None, compute_virial=False, device=None)[source]#

Internal op for the self-energy and the coordination chain rule.

Neither stage touches the mesh and both need every rank slot at once, so this runs once after the rank chunks rather than inside them. The chain rule walks the whole neighbour list, which is why running it per chunk would be the expensive mistake.

The self-energy is folded into d_energy_d_c6 before the chain rule contracts it, because it is quadratic in the coefficients and so reaches the forces through the coordination numbers. Applying it afterwards as a scalar would drop that contribution.

Parameters:

positions (Tensor)
numbers (Tensor)
species_index (Tensor)
batch_idx (Tensor)
c6 (Tensor)
dc6_dcn (Tensor)
cartesian_shifts (Tensor)
covalent_radii (Tensor)
sqrt_q (Tensor)
eigs (Tensor)
s6 (float)
s8 (float)
a1 (float)
a2 (float)
cutoff (float)
d_energy_d_c6 (Tensor)
d_energy_d_cn (Tensor)
energy (Tensor)
forces (Tensor)
virial (Tensor)
neighbor_list (Tensor | None)
neighbor_ptr (Tensor | None)
neighbor_matrix (Tensor | None)
fill_value (int | None)
compute_virial (bool)
device (str | None)

Return type:

None

nvalchemiops.torch.interactions.dispersion: Dispersion Corrections#

High-Level Interface#

DFT-D3(BJ) Dispersion Corrections#

FourierD3: Particle-Mesh DFT-D3#

Data Structures#

Internal Custom Operators#

Non-Periodic Systems#

Periodic Boundary Conditions (PBC)#

FourierD3 Operators#

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