xnn.common.models.dispersion.DispersionCorrection#

class xnn.common.models.dispersion.DispersionCorrection(term, model=None)[source]#

Bases: InteratomicPotential

A dispersion term as an xnn potential, standalone or wrapped around a model.

Standalone (model=None) the potential is the dispersion energy alone; given a short-range model it adds the dispersion energy to that model’s prediction: the wrapper’s cutoff is the larger of the model’s and the term’s cutoffs (the neighbor-list radius the data pipeline uses), and the wrapped model only ever sees the edges within its own cutoff. D4Dispersion and D3Dispersion are the two instances; enable either from a config with model.extra["dispersion"] (see build_model()).

Parameters:

  • term (torch.nn.Module) – The dispersion evaluator (DFTD4 or DFTD3): exposes cutoff (Angstrom) and evaluate(atomic_numbers, pos, edge_index, edge_vec, batch, num_graphs, cell, pbc, total_charge) returning at least "node_energy" (eV), "energy_2body", "energy_3body", "coordination_numbers" and "node_features".

  • model (InteratomicPotential or None, optional) – The short-range model to correct; None for pure dispersion.

Variables:

  • term (torch.nn.Module) – The dispersion evaluator.

  • model (InteratomicPotential or None) – The wrapped model.

  • cutoff (float) – Neighbor-list radius (Angstrom): max(model.cutoff, term.cutoff).

  • node_feature_dim (int) – The wrapped model’s feature width, or the term’s own per-atom descriptors standalone (so LES can wrap a pure dispersion model too).

Notes

forward returns the combined "energy" / "node_energy" and adds "energy_sr" (the wrapped model’s energy), "energy_disp", "energy_2body", "energy_3body" (B,) and the term’s per-atom quantities ("coordination_numbers", and for D4 "eeq_charges", "polarizabilities", "dynamic_polarizabilities"; for D3 "c6_matrix"). The total charge of each structure is read from data.total_charge ((B,); zero when absent).

inner_graph(data)[source]#

data restricted to the edges within the wrapped model’s cutoff.

Parameters:

data (AtomicGraph)

Return type:

AtomicGraph

dispersion(data)[source]#

Evaluate the dispersion term on a (batched) graph; see evaluate_on_graph().

Parameters:

data (AtomicGraph)

Return type:

Dict[str, Tensor]

forward(data)[source]#

Wrapped-model prediction plus the dispersion energy.

Parameters:

data (AtomicGraph) – The batched graph, built with this wrapper’s cutoff.

Returns:

See the class notes.

Return type:

dict of str to Tensor