FourierD3: Periodic Particle-Mesh Dispersion#

The DFT-D3 dispersion correction decays as \(1/r^6\). Summed over three dimensions that leaves a truncation error decaying only as \(1/r^3\), so converging it in real space needs a neighbour list far larger than a machine-learned force field’s own, and building that list comes to dominate the simulation step.

FourierD3 evaluates the dispersion sum on a particle mesh instead, so there is no pair cutoff on the dispersion at all. It still needs one real-space cutoff — the short coordination-number list the force field already builds — but that radius is fixed by the coordination-number model rather than by how far the dispersion reaches.

FourierD3 is periodic by construction: the mesh sum runs over the infinite lattice, so there is no open-boundary path.

In this example you will learn:

  • How to turn the published DFT-D3 tables into FourierD3Parameters

  • How to evaluate the correction on a CsCl crystal with fourier_dftd3()

  • How to check that your mesh is fine enough

Important

This script is intended as an API demonstration. Do not use this script for performance benchmarking; refer to the benchmarks folder instead.

Setup#

The D3 reference parameters are in atomic units, so every length here is in Bohr.

import os
from pathlib import Path

import torch

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

BOHR_TO_ANGSTROM = 0.529177210544
HARTREE_TO_EV = 27.211386245981

# The Torch binding runs on either device; CPU is far slower but agrees to round-off.
device = "cuda" if torch.cuda.is_available() else "cpu"
Warp UserWarning: Kernel cache artifacts from a previous Warp version were found in '/tmp/warp-cache'. These will be ignored. You can safely delete them.

Reference parameters#

The published Grimme tables, the same ones the real-space dftd3() uses. utils.py downloads and parses them on first run and caches the result.

param_file = (
    Path(os.path.expanduser("~")) / ".cache" / "nvalchemiops" / "dftd3_parameters.pt"
)
if param_file.exists():
    tables = torch.load(param_file, weights_only=True)
else:
    from utils import extract_dftd3_parameters, save_dftd3_parameters

    tables = extract_dftd3_parameters()
    save_dftd3_parameters(tables)

A CsCl crystal#

Caesium chloride is simple cubic with Cs at the corner and Cl at the body centre. The supercell below is wide enough that the coordination-number cutoff fits inside it.

LATTICE_A = 4.119 / BOHR_TO_ANGSTROM  # Bohr
REPEATS = 4

basis_fractional = torch.tensor([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]], dtype=torch.float64)
basis_numbers = torch.tensor([55, 17], dtype=torch.int32)  # Cs, Cl

offsets = torch.cartesian_prod(*[torch.arange(REPEATS, dtype=torch.float64)] * 3)
fractional = ((offsets[:, None, :] + basis_fractional[None, :, :]) / REPEATS).reshape(
    -1, 3
)
cell = torch.eye(3, dtype=torch.float64) * (LATTICE_A * REPEATS)

positions = (fractional @ cell).to(device)
numbers = basis_numbers.repeat(REPEATS**3).to(device)
cell = cell.to(device)

print(f"atoms      : {positions.shape[0]}")
print(f"cell edge  : {float(cell[0, 0]) * BOHR_TO_ANGSTROM:.2f} Angstrom")
atoms      : 128
cell edge  : 16.48 Angstrom

Selecting the species#

FourierD3Parameters decomposes Grimme’s reference tensor into separable factors, once, for the species actually present. That decomposition is what lets an environment-dependent pair coefficient be evaluated on a mesh at all, and its cost grows with the number of distinct elements — so pass only the ones in your system.

It depends on the tables, the species and the requested tolerance, but not on the damping parameters, so one instance serves every functional.

species = sorted(set(numbers.tolist()))
params = FourierD3Parameters.from_tables(
    tables["rcov"].to(device=device, dtype=torch.float64),
    tables["r4r2"].to(device=device, dtype=torch.float64),
    tables["c6ab"].to(device=device, dtype=torch.float64),
    tables["cn_ref"].to(device=device, dtype=torch.float64),
    species=species,
    device=device,
    dtype=torch.float64,
)
print(f"\nspecies              : {species}")
print(f"decomposition rank   : {params.rank}")
print(f"mesh channels needed : {params.n_species * params.rank}")
print(f"reconstruction error : {params.max_relative_error:.2e}")
species              : [17, 55]
decomposition rank   : 3
mesh channels needed : 6
reconstruction error : 1.56e-06

The coordination-number neighbour list#

Only the coordination numbers need a list, so the cutoff is short. It must equal cutoff exactly: the counting function is built to reach zero there.

cutoff = 6.0 / BOHR_TO_ANGSTROM  # 6 Angstrom, a typical MLFF cutoff, in Bohr

pbc = torch.tensor([True, True, True], device=device)
neighbor_list, neighbor_ptr, unit_shifts = neighbors.neighbor_list(
    positions,
    cutoff=cutoff,
    cell=cell,
    pbc=pbc,
    return_neighbor_list=True,
)
print(f"\nneighbour cutoff : {cutoff * BOHR_TO_ANGSTROM:.1f} Angstrom")
print(f"directed edges   : {neighbor_list.shape[1]}")
neighbour cutoff : 6.0 Angstrom
directed edges   : 3328

Evaluating the correction#

cell and cutoff are both required. Exactly one of mesh_dimensions and mesh_spacing must be given; there is no accuracy-based default. Forces and the virial are returned directly rather than obtained by differentiating the energy.

DAMPING = dict(a1=0.4289, a2=4.4407, s8=0.7875)  # PBE-D3(BJ)

energy, forces, virial = fourier_dftd3(
    positions,
    numbers,
    **DAMPING,
    fourier_d3_params=params,
    cell=cell,
    cutoff=cutoff,
    mesh_dimensions=(32, 32, 32),
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    unit_shifts=unit_shifts,
    compute_virial=True,
)

n_atoms = positions.shape[0]
print(f"\nenergy         : {energy.item():.8f} Hartree")
print(f"per atom       : {energy.item() / n_atoms * HARTREE_TO_EV * 1e3:.4f} meV")
# Every site in an ideal CsCl lattice sits at a centre of symmetry, so the forces cancel.
# A value at round-off here is a check on the gather, not a null result.
print(
    f"max |force|    : {forces.abs().max().item():.3e} Hartree/Bohr (zero by symmetry)"
)
print(f"virial trace   : {virial[0].diagonal().sum().item():.6e} Hartree")
energy         : -0.97919696 Hartree
per atom       : -208.1665 meV
max |force|    : 1.233e-15 Hartree/Bohr (zero by symmetry)
virial trace   : -2.770629e+00 Hartree

Checking the mesh#

Refining the mesh with everything else fixed shows how sensitive the answer is to it. The change between successive meshes is an indication, not a certified error bar: it says the energy has stopped moving, not that it has stopped moving toward the right value.

The mesh is also not the only control. Spline order and the decomposition tolerance (tol in FourierD3Parameters.from_tables()) both shift the result, and energy converging says nothing on its own about forces or the virial, which converge more slowly.

print(f"\n{'mesh':>6} {'energy/atom (meV)':>20} {'|change|':>12}")
previous = None
for size in (16, 32, 48):
    value = fourier_dftd3(
        positions,
        numbers,
        **DAMPING,
        fourier_d3_params=params,
        cell=cell,
        cutoff=cutoff,
        mesh_dimensions=(size, size, size),
        neighbor_list=neighbor_list,
        neighbor_ptr=neighbor_ptr,
        unit_shifts=unit_shifts,
    )[0].item()
    per_atom = value / n_atoms * HARTREE_TO_EV * 1e3
    change = "" if previous is None else f"{abs(per_atom - previous):.2e}"
    print(f"{size:3d}^3  {per_atom:20.10f} {change:>12}")
    previous = per_atom
 mesh    energy/atom (meV)     |change|
16^3       -207.5825924070
32^3       -208.1664591207     5.84e-01
48^3       -208.1830526383     1.66e-02

Summary#

  • FourierD3Parameters.from_tables decomposes the published reference tensor once, for the species present, independently of the functional.

  • fourier_dftd3 needs a periodic cell and a coordination-number list, and the list’s cutoff must equal cutoff.

  • Refine the mesh to check sensitivity; also vary the spline order and the decomposition tolerance before trusting a number, and check forces as well as energy.

FourierD3 applies no cutoff to the dispersion sum, so its cost does not grow as the interaction range does. It is not a drop-in replacement for the real-space dftd3(): it uses a modified coordination-number function, so the two models differ. See the dispersion user guide for that distinction, and prefer dftd3 for open boundary conditions.

Total running time of the script: (0 minutes 6.826 seconds)

Gallery generated by Sphinx-Gallery