DFT-D3(BJ) Dispersion Correction#

Dispersion corrections account for van der Waals interactions that standard DFT functionals underestimate. ALCHEMI Toolkit-Ops provides GPU-accelerated DFT-D3 with Becke-Johnson damping via NVIDIA Warp, supporting batched computation, periodic systems, and bindings for both PyTorch and JAX.

Important

The current implementation computes two-body terms only (C6 and C8). Three-body Axilrod-Teller-Muto (ATM/C9) contributions are not included.

Why Dispersion Correction Matters#

Standard DFT functionals systematically underestimate long-range dispersion due to their local or semi-local nature. This leads to significant errors in:

Binding energies of weakly bound complexes
Molecular crystal structures and lattice energies
Conformational energies of flexible molecules
Adsorption energies on surfaces

DFT-D3 adds an empirical pairwise correction with environment-dependent C6 coefficients that adapt based on coordination numbers. The Becke-Johnson damping variant provides improved short-range behavior for molecular geometries.

Data Formats#

Neighbor Representations#

ALCHEMI Toolkit-Ops supports two neighbor formats that produce identical results:

Neighbor Matrix (default) : Dense array of shape [num_atoms, max_neighbors]. Each row contains neighbor indices for that atom, padded with fill_value (typically num_atoms). Use with neighbor_matrix and neighbor_matrix_shifts arguments.

Neighbor List (COO) : Sparse array of shape [2, num_pairs] where row 0 contains source indices and row 1 contains target indices. No padding needed. Use with neighbor_list and unit_shifts arguments.

When to Use Each Format#

We refer the reader for a more in-depth discussion in the neighbor list documentation, but briefly:

Neighbor Matrix is preferred when:

Using torch.compile or jax.jit (fixed memory layout avoids graph breaks)
Systems have dense, uniform neighbor distributions

Neighbor List (COO) is preferred when:

Integrating with graph neural network libraries (PyG, DGL)
Systems are sparse with highly variable neighbors per atom
Memory efficiency is critical

Note

The neighbor representation should be symmetric (bidirectional): if atom j is a neighbor of atom i, then atom i should also be a neighbor of atom j.

Units#

While the DFT-D3 kernels themselves are unit-agnostic, the reference parameters themselves typically use atomic units. The table below lists quantities and their conversions from conventional units more commonly encountered in computational chemistry and materials science:

Quantity	Unit	Common conversions
Positions	Bohr	\((\text{Å} \times 1.8897259886)\)
Energy (output)	Hartree	\((\times 27.211 \rightarrow \text{eV})\)
Forces (output)	Hartree/Bohr	\((\times 51.422 \rightarrow \text{eV/Å})\)
Virial (output)	Hartree	\((\times 27.211 \rightarrow \text{eV})\)
`a2` parameter	Bohr	Same as positions
Cutoffs	Bohr	Same as positions
C6 coefficients	(Hartree \(\cdot\) Bohr\(^6\))	—

Coordination numbers are dimensionless.

Tip

Use scipy.constants for precise conversions: scipy.constants.value('atomic unit of length') for Bohr, scipy.constants.value('Hartree energy in eV') for Hartree.

Virial Convention#

The DFT-D3 virial follows the project-wide convention:

\[W_{ab} = -\partial E / \partial \varepsilon_{ab}\]

Tensile-positive Cauchy stress is obtained via \(\sigma = -W / V\) where \(V = |\det(\mathbf{C})|\).

See Conventions for the full project-wide definitions.

Data Types#

Tensor	Dtype	Notes
Positions	`float32` or `float64`	FP64 used for distance vectors only
Cell	`float32` or `float64`	Same format as Positions
Atomic numbers	`int32`
Neighbor indices	`int32`
Energy (output)	`float32`	Always
Forces (output)	`float32`	Always
Virial (output)	`float32`	Always
Coordination numbers (output)	`float32`	Always

Padding Atoms#

Atoms with atomic number 0 are treated as padding atoms and are excluded from all computations:

Padding atoms contribute zero to energy, forces, virial, and coordination numbers.
Neighbor entries pointing to padding atoms (or from padding atoms) are skipped.
Index 0 in the reference parameter arrays (rcov, r4r2, c6ab, cn_ref) is reserved for the padding element and should be set to 0.0.

# Example: batch of two systems with 3 and 2 atoms, padded to 3
numbers = torch.tensor([8, 1, 1, 6, 8, 0], dtype=torch.int32, device="cuda")
#                       ^^^system 0^^^  ^^system 1^^  ^pad^

Note

The reference parameter files from the Grimme group already follow this convention (index 0 is unused). No special handling is needed when loading standard parameters.

Parameter Setup#

DFT-D3 requires reference parameters that must be explicitly provided. Three options are available:

Option 1: D3Parameters Dataclass (Recommended)#

PyTorch

from nvalchemiops.torch.interactions.dispersion import D3Parameters

d3_params = D3Parameters(
    rcov=covalent_radii,      # [max_Z+1] in Bohr
    r4r2=r4r2_values,         # [max_Z+1] <r^4>/<r^2> expectation values
    c6ab=c6_reference,        # [max_Z+1, max_Z+1, 5, 5] C6 grid
    cn_ref=coord_num_ref,     # [max_Z+1, max_Z+1, 5, 5] CN grid
)

energy, forces, coord_num = dftd3(
    positions, numbers, neighbor_matrix,
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
)

JAX

from nvalchemiops.jax.interactions.dispersion import D3Parameters

d3_params = D3Parameters(
    rcov=covalent_radii,      # [max_Z+1] in Bohr
    r4r2=r4r2_values,         # [max_Z+1] <r^4>/<r^2> expectation values
    c6ab=c6_reference,        # [max_Z+1, max_Z+1, 5, 5] C6 grid
    cn_ref=coord_num_ref,     # [max_Z+1, max_Z+1, 5, 5] CN grid
)

energy, forces, coord_num = dftd3(
    positions, numbers, neighbor_matrix,
    a1=0.3981, a2=4.4211, s8=0.7875,
    d3_params=d3_params,
)

Option 2: Dictionary#

d3_params = {
    "rcov": covalent_radii,
    "r4r2": r4r2_values,
    "c6ab": c6_reference,
    "cn_ref": coord_num_ref,
}

Option 3: Individual Parameters#

energy, forces, coord_num = dftd3(
    positions, numbers, neighbor_matrix,
    a1=0.4289, a2=4.4407, s8=0.7875,
    covalent_radii=covalent_radii,
    r4r2=r4r2_values,
    c6_reference=c6_reference,
    coord_num_ref=coord_num_ref,
)

Obtaining Parameters#

Parameter files are available from the Grimme group website. See examples/dispersion/utils.py for loading utilities.

Functional-Specific Damping Parameters#

Common BJ damping parameters (a1, a2, s8):

Functional	a1	a2 (Bohr)	s8
PBE	0.4289	4.4407	0.7875
PBE0	0.4145	4.8593	1.2177
B3LYP	0.3981	4.4211	1.9889

See the DFT-D3 BJ-damping parameters for a complete list.

Performance Tuning#

Key Parameters#

s5_smoothing_on / s5_smoothing_off : Enable smooth cutoff with the S5 switching function. The function provides \(C^2\) continuity at both boundaries, preventing force discontinuities.

JIT Compilation#

Both neighbor formats support JIT optimization in PyTorch and JAX:

PyTorch

compiled_dftd3 = torch.compile(dftd3)

energy, forces, coord_num = compiled_dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_list=neighbor_list_coo,
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
    num_systems=num_systems  # will introduce CUDA graph break if not provided
)

JAX

from functools import partial
import jax

jit_dftd3 = jax.jit(partial(
    dftd3,
    a1=0.3981, a2=4.4211, s8=0.7875,
    num_systems=num_systems,  # must be a static integer for jax.jit
))

energy, forces, coord_num = jit_dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_list=neighbor_list_coo,
    neighbor_ptr=neighbor_list_ptr,
    d3_params=d3_params,
)

Note

When using jax.jit, provide num_systems explicitly as a static integer. If omitted, the code attempts to infer it from cell or batch_idx, which may fail during JIT tracing.

Precision Notes#

Input positions and unit cell can be float64 for large unit cells or systems far from origin
Intermediate calculations use float32 (sufficient for dispersion accuracy)
All outputs are float32

Usage Patterns#

Single Molecule (Non-Periodic)#

PyTorch

import torch
from nvalchemiops.torch.interactions.dispersion import dftd3
from nvalchemiops.torch.neighbors import neighbor_list

ANGSTROM_TO_BOHR = 1.8897259886
HARTREE_TO_EV = 27.211386245981

# Water molecule (Angstrom -> Bohr)
positions_ang = torch.tensor([
    [0.0000,  0.0000,  0.1173],  # O
    [0.0000,  0.7572, -0.4692],  # H
    [0.0000, -0.7572, -0.4692],  # H
], dtype=torch.float32, device="cuda")
positions = positions_ang * ANGSTROM_TO_BOHR

numbers = torch.tensor([8, 1, 1], dtype=torch.int32, device="cuda")

# Non-periodic: use large box with pbc=False
cell = torch.eye(3, dtype=torch.float32, device="cuda").unsqueeze(0) * 100.0
pbc = torch.tensor([False, False, False], device="cuda")

neighbor_matrix, num_neighbors, _ = neighbor_list(
    positions, cutoff=50.0, cell=cell, pbc=pbc
)

# d3_params loaded from file (see examples/dispersion/utils.py)
energy, forces, coord_num = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
)

print(f"Dispersion energy: {energy[0].item():.6f} Ha")
print(f"                   {energy[0].item() * HARTREE_TO_EV:.6f} eV")

JAX

import jax.numpy as jnp
from nvalchemiops.jax.interactions.dispersion import dftd3
from nvalchemiops.jax.neighbors import neighbor_list

ANGSTROM_TO_BOHR = 1.8897259886
HARTREE_TO_EV = 27.211386245981

# Water molecule (Angstrom -> Bohr)
positions_ang = jnp.array([
    [0.0000,  0.0000,  0.1173],  # O
    [0.0000,  0.7572, -0.4692],  # H
    [0.0000, -0.7572, -0.4692],  # H
], dtype=jnp.float32)
positions = positions_ang * ANGSTROM_TO_BOHR

numbers = jnp.array([8, 1, 1], dtype=jnp.int32)

# Non-periodic: use large box with pbc=False
cell = jnp.eye(3, dtype=jnp.float32)[None, ...] * 100.0
pbc = jnp.array([[False, False, False]])

neighbor_matrix, num_neighbors, _ = neighbor_list(
    positions, cutoff=50.0, cell=cell, pbc=pbc
)

# d3_params loaded from file (see examples/interactions/utils.py)
energy, forces, coord_num = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    a1=0.3981, a2=4.4211, s8=0.7875,
    d3_params=d3_params,
)

print(f"Dispersion energy: {energy[0].item():.6f} Ha")
print(f"                   {energy[0].item() * HARTREE_TO_EV:.6f} eV")

Batched Periodic Crystals#

PyTorch

import torch
from nvalchemiops.torch.interactions.dispersion import dftd3
from nvalchemiops.torch.neighbors import neighbor_list

ANGSTROM_TO_BOHR = 1.8897259886

# Batch of 4 systems, 8 atoms each
BATCH_SIZE = 4
atoms_per_system = 8
total_atoms = BATCH_SIZE * atoms_per_system

positions_ang = torch.randn(total_atoms, 3, device="cuda")
positions = positions_ang * ANGSTROM_TO_BOHR
numbers = torch.randint(1, 20, (total_atoms,), dtype=torch.int32, device="cuda")

# Per-system cells (Angstrom -> Bohr)
cells = torch.eye(3, device="cuda").unsqueeze(0).repeat(BATCH_SIZE, 1, 1) * 10.0
cells = cells * ANGSTROM_TO_BOHR

pbc = torch.tensor([True, True, True], device="cuda").unsqueeze(0).repeat(BATCH_SIZE, 1)

batch_idx = torch.repeat_interleave(
    torch.arange(BATCH_SIZE, dtype=torch.int32, device="cuda"),
    atoms_per_system
)
batch_ptr = torch.arange(0, total_atoms + 1, atoms_per_system, dtype=torch.int32, device="cuda")

# Build neighbors
neighbor_matrix, num_neighbors, shifts = neighbor_list(
    positions, cutoff=20.0 * ANGSTROM_TO_BOHR,
    cell=cells, pbc=pbc, batch_idx=batch_idx, batch_ptr=batch_ptr,
    method="batch_cell_list"
)

energy, forces, coord_num, virial = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=shifts,
    cell=cells,
    batch_idx=batch_idx,
    a1=0.4145, a2=4.8593, s8=1.2177,  # PBE0
    d3_params=d3_params,
    compute_virial=True
)

print(f"Energies per system (Ha): {energy}")

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.interactions.dispersion import dftd3
from nvalchemiops.jax.neighbors import neighbor_list

ANGSTROM_TO_BOHR = 1.8897259886

# Batch of 4 systems, 8 atoms each
BATCH_SIZE = 4
atoms_per_system = 8
total_atoms = BATCH_SIZE * atoms_per_system

key = jax.random.PRNGKey(0)
positions_ang = jax.random.normal(key, (total_atoms, 3), dtype=jnp.float32)
positions = positions_ang * ANGSTROM_TO_BOHR
numbers = jnp.array([1, 6, 7, 8, 1, 6, 7, 8] * BATCH_SIZE, dtype=jnp.int32)

# Per-system cells (Angstrom -> Bohr)
cells = jnp.stack([jnp.eye(3, dtype=jnp.float32) * 10.0] * BATCH_SIZE)
cells = cells * ANGSTROM_TO_BOHR

pbc = jnp.array([[True, True, True]] * BATCH_SIZE)

batch_idx = jnp.repeat(
    jnp.arange(BATCH_SIZE, dtype=jnp.int32), atoms_per_system
)
batch_ptr = jnp.arange(0, total_atoms + 1, atoms_per_system, dtype=jnp.int32)

# Build neighbors
neighbor_matrix, num_neighbors, shifts = neighbor_list(
    positions, cutoff=20.0 * ANGSTROM_TO_BOHR,
    cell=cells, pbc=pbc, batch_idx=batch_idx, batch_ptr=batch_ptr,
    method="batch_cell_list"
)

energy, forces, coord_num, virial = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=shifts,
    cell=cells,
    batch_idx=batch_idx,
    a1=0.4145, a2=4.8593, s8=1.2177,  # PBE0
    d3_params=d3_params,
    compute_virial=True,
    num_systems=BATCH_SIZE
)

print(f"Energies per system (Ha): {energy}")

Smooth Cutoff#

energy, forces, coord_num = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
    s5_smoothing_on=40.0,   # Start transition at 40 Bohr
    s5_smoothing_off=50.0,  # Complete cutoff at 50 Bohr
)

Theory Background#

This section summarizes the DFT-D3 method [1] with Becke-Johnson damping [2]. For complete derivations and benchmarks, see the original publications.

Dispersion Energy Expression#

The DFT-D3 energy with Becke-Johnson damping:

\[E_{\text{disp}} = -\frac{1}{2} \sum_{i \neq j} S_w(r_{ij}) \left[ \frac{s_6 C_6^{ij}}{r_{ij}^6 + f_{\text{damp},6}^6} + \frac{s_8 C_8^{ij}}{r_{ij}^8 + f_{\text{damp},8}^8} \right]\]

where:

\(r_{ij}\): interatomic distance
\(C_6^{ij}, C_8^{ij}\): dispersion coefficients
\(s_6, s_8\): functional-dependent scaling factors
\(f_{\text{damp}}\): Becke-Johnson damping function
\(S_w(r)\): optional smooth switching function

Coordination-Dependent C6#

A key innovation in DFT-D3 [1] is environment-dependent C6 coefficients. The C6 coefficient adapts to the local chemical environment via Gaussian interpolation over reference coordination numbers:

\[C_6(\text{CN}_i, \text{CN}_j) = \frac{\sum_{pq} C_6^{\text{ref}}[p,q] \cdot L_{pq}}{\sum_{pq} L_{pq}}\]

Coordination numbers use a geometric counting function:

\[\text{CN}_i = \sum_{j \neq i} \frac{1}{1 + \exp\left[k_1 \left(\frac{r_{\text{cov}}}{r_{ij}} - 1\right)\right]}\]

Becke-Johnson Damping#

The BJ damping function [2] prevents unphysical short-range behavior while providing improved accuracy for equilibrium geometries:

\[f_{\text{damp}} = a_1 \sqrt{C_8^{ij}/C_6^{ij}} + a_2\]

Force Calculation#

Forces include both direct and coordination-dependent contributions via chain rule:

\[\mathbf{F}_i = -\sum_j \left[\left.\frac{\partial E_{ij}}{\partial \mathbf{r}_i} \right|_{\text{CN}} + \left(\sum_k \frac{\partial E_{ik}}{\partial \text{CN}_i}\right) \frac{\partial \text{CN}_i}{\partial \mathbf{r}_i}\right]\]

This decomposition into direct and chain-rule terms is central to the multi-pass kernel architecture described in Implementation Details.

Implementation Details#

Why Multi-Pass Kernels?#

The force calculation requires computing the chain-rule contribution:

\[\mathbf{F}_{\text{chain},i} = -\left(\sum_k \frac{\partial E_{ik}}{\partial \text{CN}_i}\right) \frac{\partial \text{CN}_i}{\partial \mathbf{r}_i}\]

The term \(\sum_k \partial E_{ik}/\partial \text{CN}_i\) must be accumulated over all pairs involving atom \(i\) before computing the CN-dependent forces. This creates a data dependency that prevents computing direct and chain-rule forces in a single pass. The multi-pass architecture separates these computations:

Pass 0 (PBC only) : Convert integer unit cell shifts to Cartesian shifts using the lattice vectors. This is performed once and reused in subsequent passes.

Pass 1 : Compute coordination numbers \(\text{CN}_i\) for all atoms using the geometric counting function. These are needed for C6 interpolation in Pass 2.

Pass 2 : For each atom pair: interpolate C6 coefficients, compute damped dispersion energy, compute direct forces \(\mathbf{F}_{\text{direct}}\), and accumulate \(\partial E/\partial \text{CN}_i\) into a per-atom buffer.

Pass 3 : Using the accumulated \(\partial E/\partial \text{CN}\) values from Pass 2, compute the chain-rule force contribution and add it to the direct forces.

Neighbor Format Dispatch#

The dftd3 function dispatches to different kernel implementations based on the neighbor representation format:

Neighbor matrix (neighbor_matrix argument): Dispatches to _dftd3_matrix_op, which launches kernels that iterate over a dense [num_atoms, max_neighbors] array.
Neighbor list (neighbor_list + neighbor_ptr arguments): Dispatches to _dftd3_op, which launches kernels that use CSR (Compressed Sparse Row) format for memory-efficient sparse traversal.

Both paths execute the same four-pass algorithm and produce identical results. The choice affects memory layout and access patterns but not numerical output.

Precision Handling#

The kernels support both float32 and float64 input positions through Warp’s overload mechanism. At module load time, kernel overloads are registered for each precision:

float32 positions use wp.vec3f vectors and wp.mat33f matrices
float64 positions use wp.vec3d vectors and wp.mat33d matrices

Distance vectors are computed at the input precision, but all dispersion calculations (C6 interpolation, damping, energy/force accumulation) use float32 for efficiency. Outputs are always float32.

Numerical Stability#

Log-sum-exp trick: The Gaussian-weighted C6 interpolation computes \(\sum_k w_k C_6^{(k)} / \sum_k w_k\) where \(w_k = \exp(\text{arg}_k)\). To prevent overflow, all exponentials are computed relative to the maximum argument: \(\exp(\text{arg}_k - \text{arg}_{\max})\).
Threshold skipping: Contributions with \(\exp(\text{arg}) < 10^{-12}\) (relative to max) are skipped to avoid unnecessary computation.

References#

API Reference#

For detailed API documentation, see the PyTorch API, JAX API, and Warp API references.

FourierD3: Periodic Particle-Mesh Dispersion#

fourier_dftd3 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. It has no pair cutoff on the dispersion sum; the one real-space cutoff that remains is the short coordination-number list, which a machine-learned force field already builds for its own descriptors.

Because the coordination number sets the environment-dependent C6 coefficients, that modification means fourier_dftd3 is not expected to reproduce dftd3 exactly. The comparison below sets out where the two differ.

Choosing between `dftd3` and `fourier_dftd3`#

The two are not interchangeable, and neither is uniformly better.

	`dftd3`	`fourier_dftd3`
Boundary conditions	periodic or open	periodic only
Dispersion cutoff	yours to choose, and to converge	none
Cost with accuracy	grows as the cutoff cubed	flat, set by the mesh
Coordination-number function	standard D3	modified, decays to zero at the list cutoff
Best for	molecules in vacuum, small cells, cheap approximate corrections	condensed phases, and anything where the truncation error matters

The reason to reach for fourier_dftd3 is convergence rather than raw speed. A 1/r^6 interaction summed over three dimensions leaves a truncation error decaying only as 1/r^3, so converging dftd3 to sub-meV/atom needs a cutoff far beyond an MLFF’s own, and the neighbour list then dominates the step. fourier_dftd3 removes the cutoff instead of enlarging it.

Important

The two do not produce identical numbers, and are not expected to. fourier_dftd3 uses a modified coordination-number function that decays to zero at the list cutoff, where the standard D3 function tends to a non-zero constant. That difference is what makes the coordination numbers independent of the list used to build them; without it, no choice of cutoff converges.

In practice the gap is small. On diamond with the published tables, fourier_dftd3 reading only a 6 Angstrom list agrees to 5.9e-06 relative with dftd3 extrapolated to an infinite cutoff — where dftd3 at that same 6 Angstrom cutoff is still 10% short. Treat a difference of that order as the model difference; treat a large one as a bug report.

Note

rcov follows the same convention as dftd3. The shipped table already folds in Grimme’s 4/3 scale, so the counting function crosses one half at a separation equal to the sum of the two tabulated radii. Pass both functions the same table; do not rescale it for one of them.

Quick Start#

import torch
from nvalchemiops.torch.interactions.dispersion import (
    FourierD3Parameters, fourier_dftd3,
)
from nvalchemiops.torch.neighbors import neighbor_list

# Decompose the reference tables once, for the species present. Independent of the
# functional, so one instance serves every damping parametrisation.
params = FourierD3Parameters.from_tables(
    rcov, r4r2, c6ab, cn_ref, species=[1, 6, 8], device="cuda",
)

cutoff = 11.34  # 6 Angstrom in Bohr; see Units below
neighbors, pointer, shifts = neighbor_list(
    positions, cutoff=cutoff, cell=cell, pbc=pbc,
    return_neighbor_list=True, method="cell_list",
)

energy, forces = fourier_dftd3(
    positions, numbers,
    a1=0.4289, a2=4.4407, s8=0.7875,      # PBE-D3(BJ)
    fourier_d3_params=params, cell=cell, cutoff=cutoff,
    mesh_dimensions=(32, 32, 32),
    neighbor_list=neighbors, neighbor_ptr=pointer, unit_shifts=shifts,
)

`cutoff` Must Match the Neighbour List#

Warning

cutoff has no default, and must equal the cutoff the neighbour list was built with. The modified coordination-number function is constructed to reach zero exactly at cutoff; if the list was truncated somewhere else, the discontinuity the modification exists to remove comes straight back.

There is no default because the D3 reference parameters are conventionally in atomic units, so a value written as 6.0 and meant as 6 Angstrom would silently act as 6 Bohr, roughly half the intended cutoff, with no error and no obvious symptom.

The Neighbour List Must Hold Both Directions#

Warning

Whichever neighbour format you pass must contain both directions of every pair. That is what neighbor_list and the dense builders produce by default; a list requested with half_fill=True does not.

FourierD3 accumulates each atom’s coordination number, and the chain rule from it, out of that atom’s own row alone — the reverse edge is walked by the other atom. A half-filled list therefore loses half of every atom’s coordination, which shifts the energy by around a percent and leaves the forces non-conservative, with nothing in the output to say so.

This is not checked at runtime. Detecting it means reducing over the whole neighbour list on every call, which measured about 0.13 ms at 8,000 atoms — roughly a sixth of the evaluation it guards — and it could never run on the JAX path at all, where jax.jit makes the values unreadable. Since both builders produce full lists by default, the requirement is only reachable by asking for half_fill=True and passing the result here.

Choosing a Spline Order#

spline_order runs from 2 to 6 and defaults to 4. Every order converges to the same energy — the lattice sum does not depend on how it was interpolated — but accuracy at a fixed mesh improves sharply with order. Measured on an 8-atom cell at a 96³ mesh, against a converged reference:

order	relative error
2	`2.8e-05`
3	`2.5e-07`
4	`2.0e-09`
5	`1.4e-10`

Raising the order is usually cheaper than refining the mesh, since the spread and gather cost order³ per atom while the transform cost grows with the mesh volume. Order 3 is noticeably noisier than its neighbours and its error does not fall cleanly with refinement, so prefer an even order unless you have measured otherwise.

Choosing a Mesh#

Note

A mesh derived from mesh_spacing is rounded up to a size whose only prime factors are 2, 3, 5 and 7. cuFFT has radix kernels for those and falls back to Bluestein’s algorithm otherwise, which is not a marginal difference — a prime or large-factor edge can cost several times what the next friendly size does. The rounded mesh is never coarser than the spacing asked for. An explicit mesh_dimensions is used exactly as given, including a poor size.

Warning

Every mesh axis must hold at least spline_order nodes. The interpolation stencil is that wide and wraps periodically, so a shorter axis makes two stencil points land on the same node and the interpolation is no longer the B-spline the gather differentiates. Equality is allowed — the stencil then covers each node exactly once — and the same minimum is applied to mesh_dimensions, to a mesh derived from mesh_spacing, and to FourierD3Setup.build.

Exactly one of mesh_dimensions or mesh_spacing is required; there is no accuracy-based default to fall back on, so neither and both are errors. This applies to calls that resolve a mesh themselves, which is every JAX call and any Torch call without a setup.

A Torch setup= argument carries the mesh and spline order it was built with, so the mesh is already fixed and both selectors become optional. A mesh_dimensions that disagrees with the setup raises, and so does any mesh_spacing — rounding it against the cell would read device memory, which is the cost a precomputed setup exists to avoid. Choose the mesh once, when building the setup. As a starting point, from the FourierD3 paper:

System size	Mesh
up to 200 atoms	16³
up to 2,000	32³
up to 20,000	64³
beyond	128³

To size the mesh from a target reciprocal cutoff instead, k_max = pi * N / L gives N = ceil(k_cut * L / pi).

The mesh carries num_systems * n_species * rank channels, where rank comes from the decomposition and is typically 4 to 8. Memory grows with all three, so batching chemically dissimilar systems together costs more than batching similar ones.

Capping the Resident Mesh#

The mesh and its transforms are the dominant allocation, and they scale as num_systems * n_species * rank * nx * ny * nz. Many species or a large retained rank can therefore run a fine mesh out of device memory: seven species at rank 24 on a 128³ mesh needs several gigabytes before anything else is allocated.

rank_chunk_size caps how many rank slots are resident at once. Every stage after the coordination number is a sum over slots with no coupling between them, so a chunk of slots can be carried through spread, transform, contraction and gather on its own and its contribution added in:

energy, forces = fourier_dftd3(..., rank_chunk_size=4)

The result is unchanged to round-off. The cost is one extra spread, forward and inverse transform, and gather per chunk, so leave it at None when the mesh fits. Measured on a 300-atom, seven-species cell at rank 24 on a 64³ mesh in float64, peak allocation fell from 1066 MiB at None to 209 MiB at rank_chunk_size=4 and 79 MiB at 1, with energy, forces and virial agreeing to 1e-14 relative throughout.

Only the reciprocal stages are chunked: spread, forward transform, contraction, inverse transform and gather. The self-energy and the coordination-number chain rule need every slot at once and never touch the mesh, so they run once after the loop. That matters because the chain rule walks the whole neighbour list; running it per chunk cost 16% at 8,000 atoms with 2.4M edges.

The per-atom coefficient arrays stay at full (N, rank) width throughout — they are what the chain rule contracts at the end. They are per-atom rather than per-mesh-point, so they are a negligible part of the footprint this option exists to bound: at 8,000 atoms and rank 21 they are under 2 MiB, against a mesh measured in hundreds.

It is host-static in both bindings: it decides how many kernel launches happen, so it must be a Python integer and cannot be a tensor or a traced value.

Warning

Under jax.jit the chunk loop is unrolled at trace time, so the traced graph grows with the number of chunks. On the case above the jaxpr went from 395 equations at None to 538 at rank_chunk_size=4 and 1042 at 1, with compile time rising from 0.20 s to 0.33 s. That is cheap here, but it is paid per distinct shape signature and grows linearly in rank / rank_chunk_size, so prefer the largest chunk that fits rather than the smallest that works.

B-spline Deconvolution#

Interpolating onto a mesh attenuates each frequency, and dividing that attenuation out is what recovers the structure factor the mesh stands in for. FourierD3 uses sinc(m/N)**p, the same convention as the electrostatics PME path, via the shared compute_bspline_moduli_1d.

Forces are what this choice affects, so they are validated against central differences of a brute-force real-space lattice sum that uses no mesh, spline or transform. On the eight-atom test cell the max relative force error is 2.9e-05 at a 32-cubed mesh and 7.5e-06 at 64-cubed. The discrete B-spline modulus, the alternative convention, lands at 3.1e-05 and 6.6e-06 on the same reference: the two differ by far less than their common discretisation error, so the mesh, not the convention, sets force accuracy.

Making It Fast#

Two things matter, and neither is torch.compile on its own.

Precompute the cell-derived setup. FourierD3Setup.build derives the inverse cell, the volume, the reciprocal lattice and the B-spline moduli. None of them change while the cell and mesh are fixed, and deriving them costs a matrix inversion per call.

from nvalchemiops.torch.interactions.dispersion import FourierD3Setup

setup = FourierD3Setup.build(cell, params.n_species, mesh_dimensions=(32, 32, 32))
for step in trajectory:                      # constant-volume dynamics
    energy, forces = fourier_dftd3(..., setup=setup)

The saving is largest on small systems, where the per-call setup is a bigger share of the total, and shrinks as the mesh and spread work start to dominate.

The setup is a cache, not a requirement: every execution mode, including torch.compile(mode="reduce-overhead"), works without one.

Important

Keeping the setup consistent with the cell is yours to do, as with the PME and multipole caches. validate_for checks batch size, species count, precision, device and mesh — all metadata, so it works under compilation — but it does not compare cell values, because that would read device memory and graph capture forbids it. A setup built for a different cell mixes two cells in one evaluation with nothing raised. Rebuild it whenever the cell changes.

torch.compile adds a little more. It traces without graph breaks, but the runtime is dominated by Warp kernels and FFTs, which compilation cannot fuse into, so the remaining headroom is small. Measure before adding it to a workload.

To compare these on your own hardware, the shipped benchmark covers the first two:

python benchmarks/interactions/dispersion/benchmark_fourier_dftd3.py

It reports fourier_dftd3 and fourier_dftd3_setup rows across system sizes, timing the evaluation only — the neighbour list is built outside the timed region, per the kernel style guide.

Warning

fourier_d3_params must cover every element in the system. This is a caller precondition, checked by neither binding in any execution mode — the same contract as dftd3, and for the same reason: verifying it means reading device memory on every step.

species_map marks both padding (atomic number 0) and uncovered elements with -1, and the mesh grouping treats -1 as padding. An uncovered atom is therefore dropped from the dispersion sum and contributes no coefficients, self-energy or force, while still shifting its covered neighbours’ coordination numbers. Nothing is raised and no NaN appears: the energy comes back finite and wrong.

Build FourierD3Parameters from the species actually present — sorted(set(numbers)) is enough — and rebuild it whenever the composition changes.

FourierD3Parameters.uncovered_species(numbers) names anything missing, so the precondition can be asserted where it is cheap to do so: at setup, on a composition change, or in a test. It synchronises, so keep it out of the dynamics loop. Both bindings expose it and both give the same answer.

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

Important

For the JAX binding, jax.jit is not optional. Unjitted, every operation dispatches separately and the cost is dominated by that dispatch overhead rather than by the system size, which makes it dramatically slower than the jitted path on anything worth computing. Always wrap the call in jax.jit.

What FourierD3 Returns#

energy and forces always, and virial when compute_virial=True. The virial follows the project-wide convention, \(W_{ab} = -\partial E / \partial u_{ab}\), the same as dftd3; see Virial Convention above.

Forces are an explicit output rather than something recovered by differentiating the energy, again matching dftd3.

Warning

The returned energy is not differentiable. Every Warp kernel is launched with enable_backward=False, and neither binding registers a reverse rule on top of it:

the Torch binding registers no torch.library.register_autograd, so the returned energy comes back with requires_grad=False and grad_fn=None, detached from the input positions;
the JAX binding registers no VJP or JVP rule, so the call cannot be transposed.

Do not expect torch.autograd.grad or jax.grad on the energy to reproduce forces. Both fail loudly rather than returning a wrong answer — Torch raises RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn, and JAX raises ValueError: The FFI call to '_fd3_cn_kernel_1' cannot be differentiated — but the failure is generic enough to be mistaken for a bug in your own code, so it is worth knowing in advance.

Use the forces and virial the call already returns. They are analytic derivatives of the same energy, hand-derived through the coordination number, the spread, the reciprocal contraction and the gather, and they are checked against finite differences in the test suite. The reason for the design is that enable_backward=False is what keeps the pipeline capturable into a CUDA graph and traceable under jax.jit.

Note

This means FourierD3 cannot sit inside a larger differentiated graph — an end-to-end loss back-propagated to positions through the dispersion energy, for example. If you need that, forces gives you \(-\partial E/\partial r\) directly and can be spliced in by hand with a torch.autograd.Function, but nothing in the library does it for you.

Note

Energies are reduced with atomic adds, whose summation order varies between launches, so two identical calls can differ in the last bit. Regression fixtures should compare with a tolerance rather than for exact equality.

DFT-D3(BJ) Dispersion Correction#

Why Dispersion Correction Matters#

Quick Start#

Periodic Boundary Conditions#

Data Formats#

Neighbor Representations#

When to Use Each Format#

Units#

Virial Convention#

Data Types#

Padding Atoms#

Parameter Setup#

Option 1: D3Parameters Dataclass (Recommended)#

Option 2: Dictionary#

Option 3: Individual Parameters#

Obtaining Parameters#

Functional-Specific Damping Parameters#

Performance Tuning#

Key Parameters#

JIT Compilation#

Precision Notes#

Usage Patterns#

Single Molecule (Non-Periodic)#

Batched Periodic Crystals#

Smooth Cutoff#

Theory Background#

Dispersion Energy Expression#

Coordination-Dependent C6#

Becke-Johnson Damping#

Force Calculation#

Implementation Details#

Why Multi-Pass Kernels?#

Neighbor Format Dispatch#

Precision Handling#

Numerical Stability#

References#

API Reference#

FourierD3: Periodic Particle-Mesh Dispersion#

Choosing between dftd3 and fourier_dftd3#

Quick Start#

cutoff Must Match the Neighbour List#

The Neighbour List Must Hold Both Directions#

Choosing a Spline Order#

Choosing a Mesh#

Capping the Resident Mesh#

B-spline Deconvolution#

Making It Fast#

What FourierD3 Returns#

Choosing between `dftd3` and `fourier_dftd3`#

`cutoff` Must Match the Neighbour List#