nvalchemi.models.dftd3.DFTD3ModelWrapper

class nvalchemi.models.dftd3.DFTD3ModelWrapper(a1, a2, s8, cutoff=50.0, k1=16.0, k3=-4.0, s6=1.0, max_neighbors=128, auto_download=True, param_file=None)

DFT-D3(BJ) dispersion correction as a model wrapper.

Wraps the Warp-accelerated dftd3 kernel from nvalchemiops. Positions are expected in Å; the wrapper converts to Bohr internally and returns energies in eV, forces in eV/Å, and virials in eV.

Parameters:

• a1 (float) – BJ damping parameter \(a_1\) (dimensionless, functional-specific).

• a2 (float) – BJ damping parameter \(a_2\) in Bohr (functional-specific).

• s8 (float) – C8 coefficient scaling factor (dimensionless, functional-specific).

• cutoff (float, optional) – Interaction cutoff in Å. Defaults to 50.0 Å (≈95 Bohr), which covers virtually all D3 interactions.

• k1 (float, optional) – Steepness of the CN counting function (1/Bohr). Defaults to 16.0.

• k3 (float, optional) – Gaussian width for CN interpolation. Defaults to -4.0.

• s6 (float, optional) – C6 scaling factor. Defaults to 1.0.

• max_neighbors (int, optional) – Maximum neighbors per atom for the neighbor matrix. Defaults to 128.

• auto_download (bool, optional) – Automatically download D3 parameters if the cache file is missing. Defaults to True.

• param_file (Path or str or None, optional) – Explicit path to the D3 parameter .pt file. None uses the default cache at ~/.cache/nvalchemiops/dftd3_parameters.pt.

model_config

Mutable configuration controlling which outputs are computed. Set model.model_config.compute_stresses = True to enable virial computation for NPT/NPH simulations.

Type:

ModelConfig

rcov, r4r2, c6ab, cn_ref

D3 reference parameters registered as module buffers so they move with .to(device) calls.

Type:

nn.Buffer

adapt_input(data, **kwargs)

Collect required inputs from data without enabling gradients.

Forces are computed analytically by the Warp kernel; autograd is not required on positions.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

dict[str, Any]

adapt_output(model_output, data)

Adapt the model output to the framework output format.

Parameters:

• model_output (Any)

• data (AtomicData | Batch)

Return type:

OrderedDict[str, Float[Tensor, ‘B 1’] | Float[Tensor, ‘V 3’] | Float[Tensor, ‘V 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3’] | None]

compute_embeddings(data, **kwargs)

Compute embeddings is not meaningful for D3 models.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

AtomicData | Batch

property embedding_shapes: dict[str, tuple[int, ...]]

Retrieves the expected shapes of the node, edge, and graph embeddings.

export_model(path, as_state_dict=False)

Export model is not implemented for D3 models.

Parameters:

• path (Path)

• as_state_dict (bool)

Return type:

None

forward(data, **kwargs)

Run the DFT-D3(BJ) kernel and return a ModelOutputs dict.

Parameters:

• data (Batch) – Batch containing positions, atomic_numbers, neighbor_matrix, num_neighbors, and optionally cell / neighbor_shifts (populated by NeighborListHook).

• kwargs (Any)

Returns:

OrderedDict with keys "energies" (shape [B, 1], eV), "forces" (shape [N, 3], eV/Å), and optionally "stresses" (shape [B, 3, 3], eV — the physical virial +Σ r_ij ⊗ F_ij).

Return type:

ModelOutputs

input_data()

Return required input keys.

Overrides the base-class default to include atomic_numbers and the neighbor-matrix keys while omitting pbc (not needed for D3).

Return type:

set[str]

property model_card: ModelCard

Retrieves the model card for the model.

The model card is a Pydantic model that contains information about the model’s capabilities and requirements.

output_data()

Return the set of keys that the model produces.

Return type:

set[str]