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

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

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

nvalchemiops.torch.interactions.dispersion: Dispersion Corrections#

High-Level Interface#

DFT-D3(BJ) Dispersion Corrections#

Data Structures#

Internal Custom Operators#

Non-Periodic Systems#

Periodic Boundary Conditions (PBC)#

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