Ctrl+K

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 COO): [2, num_pairs] without padding

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

  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/interactions/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.neighborlist.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) – C8 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) – C6 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) – <r4>/<r2> expectation values [max_Z+1] as float32 for C8 computation (dimensionless). If provided, overrides the value in d3_params.

  • c6_reference (torch.Tensor | None, optional) – C6 reference values [max_Z+1, max_Z+1, 5, 5] as float32 in energy × 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. See module docstring for details on the format. Padding entries have values >= fill_value. 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 C6 and C8 dispersion terms; three-body Axilrod-Teller-Muto (ATM/C9) 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 produce identical results and support PBC

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_nm_op()

Internal custom operator for neighbor matrix format

_dftd3_nl_op()

Internal custom operator for neighbor list format

_compute_cartesian_shifts()

Pass 0 - Converts unit cell shifts to Cartesian

_cn_kernel_nm()

Pass 1 - Computes coordination numbers (neighbor matrix)

_cn_kernel_nl()

Pass 1 - Computes coordination numbers (neighbor list)

_direct_forces_and_dE_dCN_kernel_nm()

Pass 2 - neighbor matrix format

_direct_forces_and_dE_dCN_kernel_nl()

Pass 2 - neighbor list format

_cn_forces_contrib_kernel_nm()

Pass 3 - neighbor matrix format

_cn_forces_contrib_kernel_nl()

Pass 3 - neighbor list format

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.interactions.dispersion.dftd3.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) – <r⁴>/<r²> expectation values [max_Z+1] as float32 or float64. Dimensionless ratio used for computing C8 coefficients from C6 values.

  • c6ab (torch.Tensor) – C6 reference coefficients [max_Z+1, max_Z+1, interp_mesh, interp_mesh] as float32 or float64. Units are energy x 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 5x5 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 x Bohr^6 for C6 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")
>>> 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)
On this page
Show Source