nvalchemiops.interactions.dispersion: Dispersion Corrections#

Core warp-based dispersion correction implementations.

This module provides framework-agnostic warp launchers for dispersion corrections. For PyTorch bindings, use nvalchemiops.torch.interactions.dispersion instead.

Warp-Level Interface#

DFT-D3(BJ) Dispersion Corrections#

Tip

This is the low-level Warp interface that operates on warp.array objects. For PyTorch tensor support, see nvalchemiops.torch.interactions.dispersion: 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.

Non-Periodic Systems#

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

Launch DFT-D3(BJ) dispersion calculation using neighbor matrix format (non-periodic).

This is a framework-agnostic warp launcher for non-periodic (non-PBC) systems that accepts warp arrays directly and orchestrates the multi-pass kernel execution for DFT-D3(BJ) dispersion energy, forces, and coordination number calculations. Framework-specific wrappers (PyTorch, JAX) handle tensor-to-warp conversion and call this function.

For periodic systems, use dftd3_matrix_pbc() instead.

Multi-Pass Algorithm#
  1. Pass 1: Compute coordination numbers using geometric counting function

  2. Pass 2: Compute direct forces, energy, and accumulate dE/dCN

  3. Pass 3: Add CN-dependent force contribution using chain rule

param positions:

Atomic coordinates in consistent distance units (typically Bohr). Supports both float32 (vec3f) and float64 (vec3d) precision.

type positions:

wp.array(dtype=vec3f or vec3d), shape [num_atoms]

param numbers:

Atomic numbers

type numbers:

wp.array(dtype=int32), shape [num_atoms]

param neighbor_matrix:

Neighbor indices. Padding entries have values >= fill_value.

type neighbor_matrix:

wp.array2d(dtype=int32), shape [num_atoms, max_neighbors]

param covalent_radii:

Covalent radii indexed by atomic number, in same units as positions

type covalent_radii:

wp.array(dtype=float32), shape [max_Z+1]

param r4r2:

<r⁴>/<r²> expectation values for C8 computation (dimensionless)

type r4r2:

wp.array(dtype=float32), shape [max_Z+1]

param c6_reference:

C6 reference values in energy x distance^6 units

type c6_reference:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param coord_num_ref:

CN reference grid (dimensionless)

type coord_num_ref:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param a1:

Becke-Johnson damping parameter 1 (functional-dependent, dimensionless)

type a1:

float

param a2:

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

type a2:

float

param s8:

C8 term scaling factor (functional-dependent, dimensionless)

type s8:

float

param coord_num:

OUTPUT: Coordination numbers (dimensionless). Must be pre-allocated and zeroed.

type coord_num:

wp.array(dtype=float32), shape [num_atoms]

param forces:

OUTPUT: Atomic forces in energy/distance units. Must be pre-allocated and zeroed.

type forces:

wp.array(dtype=vec3f), shape [num_atoms]

param energy:

OUTPUT: Dispersion energy in energy units. Must be pre-allocated and zeroed.

type energy:

wp.array(dtype=float32), shape [num_systems]

param virial:

OUTPUT: Virial tensor (not computed for non-periodic systems). Must be pre-allocated but will not be modified.

type virial:

wp.array(dtype=mat33f), shape [num_systems]

param batch_idx:

Batch indices mapping each atom to its system index.

type batch_idx:

wp.array(dtype=int32), shape [num_atoms]

param cartesian_shifts:

SCRATCH: Pre-allocated buffer for Cartesian shift vectors. Values are not used for non-periodic systems, but the array must still be provided with shape matching neighbor_matrix. Must be pre-allocated by caller.

type cartesian_shifts:

wp.array(dtype=vec3f or vec3d), shape [num_atoms, max_neighbors]

param dE_dCN:

SCRATCH: Pre-allocated buffer for chain rule dE/dCN intermediate. Must be pre-allocated and zeroed by caller.

type dE_dCN:

wp.array(dtype=float32), shape [num_atoms]

param wp_dtype:

Warp scalar dtype (wp.float32 or wp.float64) matching positions dtype.

type wp_dtype:

type

param device:

Warp device string (e.g., ‘cuda:0’, ‘cpu’).

type device:

str

param k1:

CN counting function steepness parameter, in inverse distance units (typically 16.0 1/Bohr). Default: 16.0

type k1:

float, optional

param k3:

CN interpolation Gaussian width parameter (typically -4.0, dimensionless). Default: -4.0

type k3:

float, optional

param s6:

C6 term scaling factor (typically 1.0, dimensionless). Default: 1.0

type s6:

float, optional

param s5_smoothing_on:

Distance where S5 switching begins, in same units as positions. Default: 0.0

type s5_smoothing_on:

float, optional

param s5_smoothing_off:

Distance where S5 switching completes, in same units as positions. Default: 0.0

type s5_smoothing_off:

float, optional

param fill_value:

Value indicating padding in neighbor_matrix. If None, inferred from num_atoms.

type fill_value:

int or None, optional

returns:

All outputs are written to pre-allocated arrays (coord_num, forces, energy).

rtype:

None

Notes

  • All output arrays must be pre-allocated and zeroed by the caller

  • Supports float32 and float64 positions; outputs always float32

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

  • Two-body only: Three-body Axilrod-Teller-Muto terms not included

  • Unit consistency required: standard D3 parameters use atomic units (Bohr for distances, Hartree for energy)

  • Virial is NOT computed for non-periodic systems (use dftd3_matrix_pbc for PBC)

See also

dftd3_matrix_pbc

Neighbor matrix format with PBC support

dftd3

Neighbor list (CSR) format, non-periodic

dftd3_pbc

Neighbor list (CSR) format with PBC support

Parameters:
  • positions (array)

  • numbers (array)

  • neighbor_matrix (array)

  • covalent_radii (array)

  • r4r2 (array)

  • c6_reference (array)

  • coord_num_ref (array)

  • a1 (float)

  • a2 (float)

  • s8 (float)

  • coord_num (array)

  • forces (array)

  • energy (array)

  • virial (array)

  • batch_idx (array)

  • cartesian_shifts (array)

  • dE_dCN (array)

  • wp_dtype (type)

  • device (str)

  • k1 (float)

  • k3 (float)

  • s6 (float)

  • s5_smoothing_on (float)

  • s5_smoothing_off (float)

  • fill_value (int | None)

Return type:

None

nvalchemiops.interactions.dispersion._dftd3.dftd3(positions, numbers, idx_j, neighbor_ptr, covalent_radii, r4r2, c6_reference, coord_num_ref, a1, a2, s8, coord_num, forces, energy, virial, batch_idx, cartesian_shifts, dE_dCN, wp_dtype, device, k1=16.0, k3=-4.0, s6=1.0, s5_smoothing_on=0.0, s5_smoothing_off=0.0)[source]#

Launch DFT-D3(BJ) dispersion calculation using neighbor list (CSR) format (non-periodic).

This is a framework-agnostic warp launcher for non-periodic (non-PBC) systems that accepts warp arrays directly and orchestrates the multi-pass kernel execution for DFT-D3(BJ) dispersion energy, forces, and coordination number calculations using CSR (Compressed Sparse Row) neighbor list format. Framework-specific wrappers (PyTorch, JAX) handle tensor-to-warp conversion and call this function.

For periodic systems, use dftd3_pbc() instead.

Multi-Pass Algorithm#
  1. Pass 1: Compute coordination numbers using geometric counting function

  2. Pass 2: Compute direct forces, energy, and accumulate dE/dCN

  3. Pass 3: Add CN-dependent force contribution using chain rule

param positions:

Atomic coordinates in consistent distance units (typically Bohr). Supports both float32 (vec3f) and float64 (vec3d) precision.

type positions:

wp.array(dtype=vec3f or vec3d), shape [num_atoms]

param numbers:

Atomic numbers

type numbers:

wp.array(dtype=int32), shape [num_atoms]

param idx_j:

Destination atom indices in CSR format

type idx_j:

wp.array(dtype=int32), shape [num_edges]

param neighbor_ptr:

CSR row pointers where neighbor_ptr[i]:neighbor_ptr[i+1] gives the range of neighbors for atom i

type neighbor_ptr:

wp.array(dtype=int32), shape [num_atoms+1]

param covalent_radii:

Covalent radii indexed by atomic number, in same units as positions

type covalent_radii:

wp.array(dtype=float32), shape [max_Z+1]

param r4r2:

<r⁴>/<r²> expectation values for C8 computation (dimensionless)

type r4r2:

wp.array(dtype=float32), shape [max_Z+1]

param c6_reference:

C6 reference values in energy x distance^6 units

type c6_reference:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param coord_num_ref:

CN reference grid (dimensionless)

type coord_num_ref:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param a1:

Becke-Johnson damping parameter 1 (functional-dependent, dimensionless)

type a1:

float

param a2:

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

type a2:

float

param s8:

C8 term scaling factor (functional-dependent, dimensionless)

type s8:

float

param coord_num:

OUTPUT: Coordination numbers (dimensionless). Must be pre-allocated and zeroed.

type coord_num:

wp.array(dtype=float32), shape [num_atoms]

param forces:

OUTPUT: Atomic forces in energy/distance units. Must be pre-allocated and zeroed.

type forces:

wp.array(dtype=vec3f), shape [num_atoms]

param energy:

OUTPUT: Dispersion energy in energy units. Must be pre-allocated and zeroed.

type energy:

wp.array(dtype=float32), shape [num_systems]

param virial:

OUTPUT: Virial tensor (not computed for non-periodic systems). Must be pre-allocated but will not be modified.

type virial:

wp.array(dtype=mat33f), shape [num_systems]

param batch_idx:

Batch indices mapping each atom to its system index.

type batch_idx:

wp.array(dtype=int32), shape [num_atoms]

param cartesian_shifts:

SCRATCH: Pre-allocated buffer for Cartesian shift vectors. Values are not used for non-periodic systems, but the array must still be provided with length matching idx_j. Must be pre-allocated by caller.

type cartesian_shifts:

wp.array(dtype=vec3f or vec3d), shape [num_edges]

param dE_dCN:

SCRATCH: Pre-allocated buffer for chain rule dE/dCN intermediate. Must be pre-allocated and zeroed by caller.

type dE_dCN:

wp.array(dtype=float32), shape [num_atoms]

param wp_dtype:

Warp scalar dtype (wp.float32 or wp.float64) matching positions dtype.

type wp_dtype:

type

param device:

Warp device string (e.g., ‘cuda:0’, ‘cpu’).

type device:

str

param k1:

CN counting function steepness parameter, in inverse distance units (typically 16.0 1/Bohr). Default: 16.0

type k1:

float, optional

param k3:

CN interpolation Gaussian width parameter (typically -4.0, dimensionless). Default: -4.0

type k3:

float, optional

param s6:

C6 term scaling factor (typically 1.0, dimensionless). Default: 1.0

type s6:

float, optional

param s5_smoothing_on:

Distance where S5 switching begins, in same units as positions. Default: 0.0

type s5_smoothing_on:

float, optional

param s5_smoothing_off:

Distance where S5 switching completes, in same units as positions. Default: 0.0

type s5_smoothing_off:

float, optional

returns:

All outputs are written to pre-allocated arrays (coord_num, forces, energy). Virial is not computed for non-periodic systems.

rtype:

None

Notes

  • All output arrays must be pre-allocated and zeroed by the caller

  • Supports float32 and float64 positions; outputs always float32

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

  • Two-body only: Three-body Axilrod-Teller-Muto terms not included

  • Unit consistency required: standard D3 parameters use atomic units (Bohr for distances, Hartree for energy)

  • CSR format is more memory-efficient for sparse neighbor lists

  • Virial is NOT computed for non-periodic systems (use dftd3_pbc for PBC)

See also

dftd3_pbc

Neighbor list (CSR) format with PBC support

dftd3_matrix

Neighbor matrix format, non-periodic

dftd3_matrix_pbc

Neighbor matrix format with PBC support

Parameters:
  • positions (array)

  • numbers (array)

  • idx_j (array)

  • neighbor_ptr (array)

  • covalent_radii (array)

  • r4r2 (array)

  • c6_reference (array)

  • coord_num_ref (array)

  • a1 (float)

  • a2 (float)

  • s8 (float)

  • coord_num (array)

  • forces (array)

  • energy (array)

  • virial (array)

  • batch_idx (array)

  • cartesian_shifts (array)

  • dE_dCN (array)

  • wp_dtype (type)

  • device (str)

  • k1 (float)

  • k3 (float)

  • s6 (float)

  • s5_smoothing_on (float)

  • s5_smoothing_off (float)

Return type:

None

Periodic Boundary Conditions (PBC)#

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

Launch DFT-D3(BJ) dispersion calculation using neighbor matrix format with PBC.

This is a framework-agnostic warp launcher for periodic boundary condition (PBC) systems that accepts warp arrays directly and orchestrates the multi-pass kernel execution for DFT-D3(BJ) dispersion energy, forces, virial, and coordination number calculations. Framework-specific wrappers (PyTorch, JAX) handle tensor-to-warp conversion and call this function.

For non-periodic systems, use dftd3_matrix() instead.

Multi-Pass Algorithm#
  1. Pass 0: Convert unit cell shifts to Cartesian coordinates

  2. Pass 1: Compute coordination numbers using geometric counting function

  3. Pass 2: Compute direct forces, energy, and accumulate dE/dCN

  4. Pass 3: Add CN-dependent force contribution using chain rule

param positions:

Atomic coordinates in consistent distance units (typically Bohr). Supports both float32 (vec3f) and float64 (vec3d) precision.

type positions:

wp.array(dtype=vec3f or vec3d), shape [num_atoms]

param numbers:

Atomic numbers

type numbers:

wp.array(dtype=int32), shape [num_atoms]

param neighbor_matrix:

Neighbor indices. Padding entries have values >= fill_value.

type neighbor_matrix:

wp.array2d(dtype=int32), shape [num_atoms, max_neighbors]

param cell:

Unit cell lattice vectors for PBC, in same dtype/units as positions. Convention: cell[s, i, :] is the i-th lattice vector for system s (row vectors).

type cell:

wp.array(dtype=mat33f or mat33d), shape [num_systems]

param neighbor_matrix_shifts:

Integer unit cell shifts for PBC. shift[i, k] is the shift for the k-th neighbor of atom i.

type neighbor_matrix_shifts:

wp.array2d(dtype=vec3i), shape [num_atoms, max_neighbors]

param covalent_radii:

Covalent radii indexed by atomic number, in same units as positions

type covalent_radii:

wp.array(dtype=float32), shape [max_Z+1]

param r4r2:

<r⁴>/<r²> expectation values for C8 computation (dimensionless)

type r4r2:

wp.array(dtype=float32), shape [max_Z+1]

param c6_reference:

C6 reference values in energy x distance^6 units

type c6_reference:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param coord_num_ref:

CN reference grid (dimensionless)

type coord_num_ref:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param a1:

Becke-Johnson damping parameter 1 (functional-dependent, dimensionless)

type a1:

float

param a2:

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

type a2:

float

param s8:

C8 term scaling factor (functional-dependent, dimensionless)

type s8:

float

param coord_num:

OUTPUT: Coordination numbers (dimensionless). Must be pre-allocated and zeroed.

type coord_num:

wp.array(dtype=float32), shape [num_atoms]

param forces:

OUTPUT: Atomic forces in energy/distance units. Must be pre-allocated and zeroed.

type forces:

wp.array(dtype=vec3f), shape [num_atoms]

param energy:

OUTPUT: Dispersion energy in energy units. Must be pre-allocated and zeroed.

type energy:

wp.array(dtype=float32), shape [num_systems]

param virial:

OUTPUT: Virial tensor in energy units. Must be pre-allocated and zeroed. Only computed if compute_virial=True.

type virial:

wp.array(dtype=mat33f), shape [num_systems]

param batch_idx:

Batch indices mapping each atom to its system index.

type batch_idx:

wp.array(dtype=int32), shape [num_atoms]

param cartesian_shifts:

SCRATCH: Pre-allocated buffer for Cartesian shift vectors. Populated by Pass 0 from unit cell shifts. Must be pre-allocated with shape matching neighbor_matrix.

type cartesian_shifts:

wp.array(dtype=vec3f or vec3d), shape [num_atoms, max_neighbors]

param dE_dCN:

SCRATCH: Pre-allocated buffer for chain rule dE/dCN intermediate. Must be pre-allocated and zeroed by caller.

type dE_dCN:

wp.array(dtype=float32), shape [num_atoms]

param wp_dtype:

Warp scalar dtype (wp.float32 or wp.float64) matching positions dtype.

type wp_dtype:

type

param device:

Warp device string (e.g., ‘cuda:0’, ‘cpu’).

type device:

str

param k1:

CN counting function steepness parameter, in inverse distance units (typically 16.0 1/Bohr). Default: 16.0

type k1:

float, optional

param k3:

CN interpolation Gaussian width parameter (typically -4.0, dimensionless). Default: -4.0

type k3:

float, optional

param s6:

C6 term scaling factor (typically 1.0, dimensionless). Default: 1.0

type s6:

float, optional

param s5_smoothing_on:

Distance where S5 switching begins, in same units as positions. Default: 0.0

type s5_smoothing_on:

float, optional

param s5_smoothing_off:

Distance where S5 switching completes, in same units as positions. Default: 0.0

type s5_smoothing_off:

float, optional

param fill_value:

Value indicating padding in neighbor_matrix. If None, inferred from num_atoms.

type fill_value:

int or None, optional

param compute_virial:

If True, compute virial tensor. Default: False

type compute_virial:

bool, optional

returns:

All outputs are written to pre-allocated arrays (coord_num, forces, energy, virial).

rtype:

None

Notes

  • All output arrays must be pre-allocated and zeroed by the caller

  • Supports float32 and float64 positions/cell; outputs always float32

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

  • Two-body only: Three-body Axilrod-Teller-Muto terms not included

  • Unit consistency required: standard D3 parameters use atomic units (Bohr for distances, Hartree for energy)

  • Virial tensor is computed when compute_virial=True

See also

dftd3_matrix

Neighbor matrix format, non-periodic

dftd3

Neighbor list (CSR) format, non-periodic

dftd3_pbc

Neighbor list (CSR) format with PBC support

Parameters:
  • positions (array)

  • numbers (array)

  • neighbor_matrix (array)

  • cell (array)

  • neighbor_matrix_shifts (array)

  • covalent_radii (array)

  • r4r2 (array)

  • c6_reference (array)

  • coord_num_ref (array)

  • a1 (float)

  • a2 (float)

  • s8 (float)

  • coord_num (array)

  • forces (array)

  • energy (array)

  • virial (array)

  • batch_idx (array)

  • cartesian_shifts (array)

  • dE_dCN (array)

  • wp_dtype (type)

  • device (str)

  • k1 (float)

  • k3 (float)

  • s6 (float)

  • s5_smoothing_on (float)

  • s5_smoothing_off (float)

  • fill_value (int | None)

  • compute_virial (bool)

Return type:

None

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

Launch DFT-D3(BJ) dispersion calculation using neighbor list (CSR) format with PBC.

This is a framework-agnostic warp launcher for periodic boundary condition (PBC) systems that accepts warp arrays directly and orchestrates the multi-pass kernel execution for DFT-D3(BJ) dispersion energy, forces, virial, and coordination number calculations using CSR (Compressed Sparse Row) neighbor list format. Framework-specific wrappers (PyTorch, JAX) handle tensor-to-warp conversion and call this function.

For non-periodic systems, use dftd3() instead.

Multi-Pass Algorithm#
  1. Pass 0: Convert unit cell shifts to Cartesian coordinates

  2. Pass 1: Compute coordination numbers using geometric counting function

  3. Pass 2: Compute direct forces, energy, and accumulate dE/dCN

  4. Pass 3: Add CN-dependent force contribution using chain rule

param positions:

Atomic coordinates in consistent distance units (typically Bohr). Supports both float32 (vec3f) and float64 (vec3d) precision.

type positions:

wp.array(dtype=vec3f or vec3d), shape [num_atoms]

param numbers:

Atomic numbers

type numbers:

wp.array(dtype=int32), shape [num_atoms]

param idx_j:

Destination atom indices in CSR format

type idx_j:

wp.array(dtype=int32), shape [num_edges]

param neighbor_ptr:

CSR row pointers where neighbor_ptr[i]:neighbor_ptr[i+1] gives the range of neighbors for atom i

type neighbor_ptr:

wp.array(dtype=int32), shape [num_atoms+1]

param cell:

Unit cell lattice vectors for PBC, in same dtype/units as positions. Convention: cell[s, i, :] is the i-th lattice vector for system s (row vectors).

type cell:

wp.array(dtype=mat33f or mat33d), shape [num_systems]

param unit_shifts:

Integer unit cell shifts for PBC. shift[e] is the shift for edge e.

type unit_shifts:

wp.array(dtype=vec3i), shape [num_edges]

param covalent_radii:

Covalent radii indexed by atomic number, in same units as positions

type covalent_radii:

wp.array(dtype=float32), shape [max_Z+1]

param r4r2:

<r⁴>/<r²> expectation values for C8 computation (dimensionless)

type r4r2:

wp.array(dtype=float32), shape [max_Z+1]

param c6_reference:

C6 reference values in energy x distance^6 units

type c6_reference:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param coord_num_ref:

CN reference grid (dimensionless)

type coord_num_ref:

wp.array4d(dtype=float32), shape [max_Z+1, max_Z+1, 5, 5]

param a1:

Becke-Johnson damping parameter 1 (functional-dependent, dimensionless)

type a1:

float

param a2:

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

type a2:

float

param s8:

C8 term scaling factor (functional-dependent, dimensionless)

type s8:

float

param coord_num:

OUTPUT: Coordination numbers (dimensionless). Must be pre-allocated and zeroed.

type coord_num:

wp.array(dtype=float32), shape [num_atoms]

param forces:

OUTPUT: Atomic forces in energy/distance units. Must be pre-allocated and zeroed.

type forces:

wp.array(dtype=vec3f), shape [num_atoms]

param energy:

OUTPUT: Dispersion energy in energy units. Must be pre-allocated and zeroed.

type energy:

wp.array(dtype=float32), shape [num_systems]

param virial:

OUTPUT: Virial tensor in energy units. Must be pre-allocated and zeroed. Only computed if compute_virial=True.

type virial:

wp.array(dtype=mat33f), shape [num_systems]

param batch_idx:

Batch indices mapping each atom to its system index.

type batch_idx:

wp.array(dtype=int32), shape [num_atoms]

param cartesian_shifts:

SCRATCH: Pre-allocated buffer for Cartesian shift vectors. Populated by Pass 0 from unit cell shifts. Must be pre-allocated with length matching idx_j.

type cartesian_shifts:

wp.array(dtype=vec3f or vec3d), shape [num_edges]

param dE_dCN:

SCRATCH: Pre-allocated buffer for chain rule dE/dCN intermediate. Must be pre-allocated and zeroed by caller.

type dE_dCN:

wp.array(dtype=float32), shape [num_atoms]

param wp_dtype:

Warp scalar dtype (wp.float32 or wp.float64) matching positions dtype.

type wp_dtype:

type

param device:

Warp device string (e.g., ‘cuda:0’, ‘cpu’).

type device:

str

param k1:

CN counting function steepness parameter, in inverse distance units (typically 16.0 1/Bohr). Default: 16.0

type k1:

float, optional

param k3:

CN interpolation Gaussian width parameter (typically -4.0, dimensionless). Default: -4.0

type k3:

float, optional

param s6:

C6 term scaling factor (typically 1.0, dimensionless). Default: 1.0

type s6:

float, optional

param s5_smoothing_on:

Distance where S5 switching begins, in same units as positions. Default: 0.0

type s5_smoothing_on:

float, optional

param s5_smoothing_off:

Distance where S5 switching completes, in same units as positions. Default: 0.0

type s5_smoothing_off:

float, optional

param compute_virial:

If True, compute virial tensor. Default: False

type compute_virial:

bool, optional

returns:

All outputs are written to pre-allocated arrays (coord_num, forces, energy, virial).

rtype:

None

Notes

  • All output arrays must be pre-allocated and zeroed by the caller

  • Supports float32 and float64 positions/cell; outputs always float32

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

  • Two-body only: Three-body Axilrod-Teller-Muto terms not included

  • Unit consistency required: standard D3 parameters use atomic units (Bohr for distances, Hartree for energy)

  • Virial tensor is computed when compute_virial=True

See also

dftd3

Neighbor list (CSR) format, non-periodic

dftd3_matrix

Neighbor matrix format, non-periodic

dftd3_matrix_pbc

Neighbor matrix format with PBC support

Parameters:
  • positions (array)

  • numbers (array)

  • idx_j (array)

  • neighbor_ptr (array)

  • cell (array)

  • unit_shifts (array)

  • covalent_radii (array)

  • r4r2 (array)

  • c6_reference (array)

  • coord_num_ref (array)

  • a1 (float)

  • a2 (float)

  • s8 (float)

  • coord_num (array)

  • forces (array)

  • energy (array)

  • virial (array)

  • batch_idx (array)

  • cartesian_shifts (array)

  • dE_dCN (array)

  • wp_dtype (type)

  • device (str)

  • k1 (float)

  • k3 (float)

  • s6 (float)

  • s5_smoothing_on (float)

  • s5_smoothing_off (float)

  • compute_virial (bool)

Return type:

None