<!-- markdownlint-disable MD013 -->

(dispersion_userguide)=

# 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](https://nvidia.github.io/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.

## Quick Start

```{important}
**DFT-D3 parameters must be explicitly provided** via `d3_params` (a
{class}`~nvalchemiops.torch.interactions.dispersion.D3Parameters` instance or
dictionary). See [Parameter Setup](#parameter-setup) for details.

**Unit consistency is required**: Standard D3 parameters use atomic units---
Bohr for lengths, Hartree for energies. Unit mismatches may cause the neighbor list
calculation to run out of memory (e.g. cutoff in Bohr, but positions are in
Angstroms) or yield an unconverged estimate of the dispersion corrections (i.e.
cutoff in Angstrom, positions in Bohr). See [Units](dispersion_units) for guidance
on units for each parameter.
```

:::::::{tab-set}

::::::{tab-item} Neighbor Matrix (Dense)
:sync: matrix

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

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

# Build neighbor matrix; 50 Bohr cutoff
neighbor_matrix, num_neighbors, shifts = neighbor_list(
    positions, cutoff=50.0, cell=cell, pbc=pbc
)

# Compute dispersion correction (PBE functional)
energy, forces, coord_num = dftd3(
    positions=positions,           # [num_atoms, 3] in Bohr
    numbers=numbers,               # [num_atoms] atomic numbers
    neighbor_matrix=neighbor_matrix,
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
)
```

:::

:::{tab-item} JAX
:sync: jax

```python
from nvalchemiops.jax.interactions.dispersion import dftd3
from nvalchemiops.jax.neighbors import neighbor_list

# Build neighbor matrix; 50 Bohr cutoff
neighbor_matrix, num_neighbors, shifts = neighbor_list(
    positions, cutoff=50.0, cell=cell, pbc=pbc
)

# Compute dispersion correction (PBE functional)
energy, forces, coord_num = dftd3(
    positions=positions,           # [num_atoms, 3] in Bohr
    numbers=numbers,               # [num_atoms] atomic numbers
    neighbor_matrix=neighbor_matrix,
    a1=0.3981, a2=4.4211, s8=0.7875,
    d3_params=d3_params,
)
```

:::

::::

::::::

::::::{tab-item} Neighbor List (Sparse COO)
:sync: coo

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

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

# Build neighbor list in COO format; 50 Bohr cutoff
neighbor_list_coo, neighbor_list_ptr, unit_shifts = neighbor_list(
    positions, cutoff=50.0, cell=cell, pbc=pbc, return_neighbor_list=True
)

# Compute dispersion correction (PBE functional)
energy, forces, coord_num = dftd3(
    positions=positions,           # [num_atoms, 3] in Bohr
    numbers=numbers,               # [num_atoms] atomic numbers
    neighbor_list=neighbor_list_coo,  # [2, num_pairs]
    neighbor_ptr=neighbor_list_ptr,
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
)
```

:::

:::{tab-item} JAX
:sync: jax

```python
from nvalchemiops.jax.interactions.dispersion import dftd3
from nvalchemiops.jax.neighbors import neighbor_list

# Build neighbor list in COO format; 50 Bohr cutoff
neighbor_list_coo, neighbor_list_ptr, unit_shifts = neighbor_list(
    positions, cutoff=50.0, cell=cell, pbc=pbc, return_neighbor_list=True
)

# Compute dispersion correction (PBE functional)
energy, forces, coord_num = dftd3(
    positions=positions,           # [num_atoms, 3] in Bohr
    numbers=numbers,               # [num_atoms] atomic numbers
    neighbor_list=neighbor_list_coo,  # [2, num_pairs]
    neighbor_ptr=neighbor_list_ptr,
    a1=0.3981, a2=4.4211, s8=0.7875,
    d3_params=d3_params,
)
```

:::

::::

::::::

:::::::

### Periodic Boundary Conditions

For periodic systems, provide the cell and shift arrays:

:::::::{tab-set}

::::::{tab-item} Neighbor Matrix (Dense)
:sync: matrix

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

```python
energy, forces, coord_num, virial = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=neighbor_matrix_shifts,   # [num_atoms, max_neighbors, 3]
    cell=cell,                                       # [num_systems, 3, 3]
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
    compute_virial=True                              # also compute virial
)
```

:::

:::{tab-item} JAX
:sync: jax

```python
energy, forces, coord_num, virial = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=neighbor_matrix_shifts,   # [num_atoms, max_neighbors, 3]
    cell=cell,                                       # [num_systems, 3, 3]
    a1=0.3981, a2=4.4211, s8=0.7875,
    d3_params=d3_params,
    compute_virial=True                              # also compute virial
)
```

:::

::::

::::::

::::::{tab-item} Neighbor List (Sparse COO)
:sync: coo

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

```python
energy, forces, coord_num = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_list=neighbor_list_coo,
    unit_shifts=unit_shifts,           # [num_pairs, 3]
    cell=cell,                         # [num_systems, 3, 3]
    a1=0.4289, a2=4.4407, s8=0.7875,
    d3_params=d3_params,
)
```

:::

:::{tab-item} JAX
:sync: jax

```python
energy, forces, coord_num = dftd3(
    positions=positions,
    numbers=numbers,
    neighbor_list=neighbor_list_coo,
    unit_shifts=unit_shifts,           # [num_pairs, 3]
    cell=cell,                         # [num_systems, 3, 3]
    a1=0.3981, a2=4.4211, s8=0.7875,
    d3_params=d3_params,
)
```

:::

::::

::::::

:::::::

Use `batch_idx` to specify multi-system batches, mapping each atom to its system.

## 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](./neighborlist.md), 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

(dispersion_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 {ref}`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`.

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

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

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

:::

:::{tab-item} JAX
:sync: jax

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

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

### Option 3: Individual Parameters

```python
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](https://www.chemie.uni-bonn.de/grimme/de/software/dft-d3/).
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](https://www.chemie.uni-bonn.de/grimme/de/software/dft-d3/bj_damping)
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:

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

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

:::

:::{tab-item} JAX
:sync: jax

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

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

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

:::

:::{tab-item} JAX
:sync: jax

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

::::{tab-set}

:::{tab-item} PyTorch
:sync: pytorch

```python
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}")
```

:::

:::{tab-item} JAX
:sync: jax

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

```python
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 [^grimme2010] with Becke-Johnson
damping [^grimme2011]. For complete derivations and benchmarks, see the
original publications.

### Dispersion Energy Expression

The DFT-D3 energy with Becke-Johnson damping:

```{math}
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 [^grimme2010] is environment-dependent C6 coefficients.
The C6 coefficient adapts to the local chemical environment via Gaussian
interpolation over reference coordination numbers:

```{math}
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:

```{math}
\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 [^grimme2011] prevents unphysical short-range behavior
while providing improved accuracy for equilibrium geometries:

```{math}
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:

```{math}
\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).

## Implementation Details

### Why Multi-Pass Kernels?

The force calculation requires computing the chain-rule contribution:

```{math}
\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

[^grimme2010]: Grimme, S.; Antony, J.; Ehrlich, S.; Krieg, H. "A consistent and
    accurate ab initio parametrization of density functional dispersion
    correction (DFT-D) for the 94 elements H-Pu." _J. Chem. Phys._ **2010**,
    _132_, 154104. [DOI: 10.1063/1.3382344](https://doi.org/10.1063/1.3382344)

[^grimme2011]: Grimme, S.; Ehrlich, S.; Goerigk, L. "Effect of the damping
    function in dispersion corrected density functional theory."
    _J. Comput. Chem._ **2011**, _32_, 1456-1465.
    [DOI: 10.1002/jcc.21759](https://doi.org/10.1002/jcc.21759)

## API Reference

For detailed API documentation, see the [PyTorch API](../../modules/torch/dispersion), [JAX API](../../modules/jax/dispersion), and [Warp API](../../modules/warp/dispersion) 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

```python
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:

```python
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.

```python
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:

```bash
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](../about/kernel-style-guide.md).

````{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.

```python
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](#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.
```