Model wrappers

These classes are used to wrap existing models and provide additional functionalities, such as new outputs or modified behavior.

class metatomic.torch.heat_flux.HeatFlux(model: AtomisticModel)[source]

HeatFlux is a wrapper around an AtomisticModel that computes the heat flux of a system using the unfolded system approach.

The unfolded system is generated by creating replica atoms for those near the cell boundaries within the interaction range of the model wrapped. The wrapper adds the heat flux to the model’s outputs under the key “heat_flux”.

For more details on the heat flux calculation, see Langer, M. F., et al., Heat flux for semilocal machine-learning potentials. (2023). Physical Review B, 108, L100302.

Parameters:

model (AtomisticModel) – the AtomisticModel to wrap, which should be able to compute atomic energies and their gradients with respect to positions

static wrap(model: AtomisticModel) → AtomisticModel[source]

Wrap an existing model able to compute atomic energies (i.e. model with a per-atom "energy" output, or any energy variants like "energy/pbe"), creating a new model that can compute the heat flux in addition to the original model’s outputs.

The returned model will have a "heat_flux[/variant]" output for each energy variant model can compute, where the same variant is used for the heat flux calculation. For example, if the original model can compute both an "energy" output and an "energy/pbe" output, the wrapped model will have both a "heat_flux" output (using the default energy output) and a "heat_flux/pbe" output (using the "energy/pbe" output).

Parameters:

model (AtomisticModel) – the AtomisticModel to wrap

Return type:

AtomisticModel

class metatomic.torch.dftd3.DFTD3(model: AtomisticModel, *, damping_params: Dict[str, Dict[str, float]], d3_params: Dict[str, Tensor] | None = None, cutoff: float | None = None, cn_cutoff: float | None = None, excluded_atom_types: List[int] | None = None)[source]

DFTD3 wraps an AtomisticModel and adds a DFT-D3(BJ) dispersion correction to its energy output(s). The three-body correction term is not included.

The wrapper can correct multiple output variants at once, each with its own damping parameters. For every energy output key (e.g. "energy" or "energy/pbe") listed in damping_params, the wrapper adds the D3 energy as a differentiable tensor: E_corrected = E_base + E_D3.

The D3 energy is implemented in pure PyTorch and is naturally differentiable: torch.autograd flows from the corrected energy back to positions and cell through the neighbor list distances.

damping_params can also be specified for non-conservative outputs such as "non_conservative_force/<variant>" and "non_conservative_stress/<variant>". Non-conservative force/stress outputs are corrected only when their output keys are explicitly listed in damping_params; energy damping parameters are not inferred for these outputs. The D3 contribution to non-conservative force/stress outputs is computed analytically with respect to the neighbor-vector values, including the coordination-number dependence of the interpolated C6 coefficients.

selected_atoms is supported with the usual domain-decomposition convention: the D3 environment is computed with all atoms in each System, while pair energies are split equally between the two pair endpoints and only the shares belonging to selected atoms are added. Per-atom energy outputs use the same half-pair split.

excluded_atom_types can be used to disable D3 pair energies involving specific atom types, while keeping all atoms in the D3 coordination-number environment. This is useful for systems where D3 should not be applied to pairs involving selected species, such as common cations.

The D3 reference tables (d3_params) are shared across variants, matching the convention that the Grimme reference data is functional independent. The reference tables include four sets of parameters, rcov (dimension of length, Bohr), r4r2 (dimension of length, Bohr), c6 (dimension of energy * length^6, Hartree * Bohr^6), and cn_ref (dimensionless). Damping parameters (a1, a2, s8, …) are provided per variant. Among these damping parameters, only a2 has a dimension of length (Bohr), the others are dimensionless. All D3 tables and damping parameters must be passed in atomic units. The wrapper converts the final D3 energy into the wrapped model’s energy unit of the corresponding output.

Parameters:

  • model (AtomisticModel) – the AtomisticModel to wrap

  • damping_params (Dict[str, Dict[str, float]]) – a mapping from an output key to a mapping of damping parameters for that output. Keys can be "energy[/<variant>]", "non_conservative_force[/<variant>]" or "non_conservative_stress[/<variant>]". Each damping map must provide a1, a2 and s8; s6 is optional (defaults to 1.0).

  • d3_params (Dict[str, Tensor] | None) – shared DFT-D3 reference tables with keys "rcov" (shape (Z,)), "r4r2" (shape (Z,)), "c6" (shape (Z, Z, M, M)) and "cn_ref" (shape (Z, M) — per-element CN reference grid, with -1 marking absent slots), where Z-1 is the maximum atomic number supported by the wrapped model (because of the 0-based indexing), and M is the number of CN reference points. Tables must be in D3 atomic units. If None, the packaged Grimme D3 reference tables are used.

  • cutoff (float | None) – dispersion-pair cutoff in the wrapped model’s length unit. If None, defaults to the standard Grimme value of 50 Bohr converted into the model’s length unit.

  • cn_cutoff (float | None) – coordination-number cutoff in the wrapped model’s length unit. If None, defaults to 25 Bohr converted into the model’s length unit.

  • excluded_atom_types (List[int] | None) – optional atom types for which D3 pair energies should be disabled. Any pair where either endpoint has one of these types contributes zero D3 energy. Coordination numbers are still computed with all atoms. The wrapped model’s atomic types must be real atomic numbers; these are used directly to index the D3 parameter tables.

static wrap(model: AtomisticModel, *, damping_params: Dict[str, Dict[str, float]], d3_params: Dict[str, Tensor] | None = None, cutoff: float | None = None, cn_cutoff: float | None = None, excluded_atom_types: List[int] | None = None) → AtomisticModel[source]

Wrap model with a differentiable DFT-D3(BJ) energy correction.

The returned AtomisticModel has the same outputs as the input, but each output listed in damping_params is corrected by the corresponding D3 contribution. The correction is differentiable so the standard autograd path produces D3-corrected conservative forces and stress.

Parameters:
Return type:

AtomisticModel

class metatomic.torch.weighted_sum.WeightedSum(model: ModelInterface, output_name: str, weights: Dict[str, float])[source]

Wrap a model, adding a new variant computing a fixed weighted sum of several existing variants.

The variants combined in the sum (for example "energy/pbe", "energy/r2scan" and "energy/lda") are all computed by the wrapped model in a single call. This does not detach or otherwise interrupt the autograd graph, so the weighted sum output stays connected to the same computational graph as each individual variant.

Parameters:

  • model (ModelInterface) – underlying ModelInterface. The wrap() method obtains this module from AtomisticModel.module.

  • output_name (str) – name of the new output computing the weighted sum, e.g. "energy/mix".

  • weights (Dict[str, float]) – mapping from the name of a existing output variants of model (e.g. "energy/pbe") to its fixed coefficient in the weighted sum.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

static wrap(model: AtomisticModel, output_name: str, weights: Dict[str, float], normalize_coefficients: bool = False) → AtomisticModel[source]

Wrap an existing model, adding a new output_name output computing the fixed weighted sum of the outputs named in weights.

The returned model retains every output declared by model under its original name, and adds the new weighted sum output. The original metadata, requested inputs, neighbor lists, and compatible capabilities are preserved.

In particular, the individual variants entering the sum stay accessible, and a model declaring both a default quantity and variants of it (for example "energy" and "energy/pbe") keeps both. At the same time, both a quantity and variants of it (for example "energy" and "energy/pbe") can be used as inputs in the same weighted sum call. If output_name collides with an existing output of model, a ValueError is raised.

Parameters:

  • model (AtomisticModel) – the AtomisticModel to wrap

  • output_name (str) – name of the new output to add, e.g. "energy/mix"

  • weights (Dict[str, float]) – mapping from the name of an existing output of model (e.g. "energy/pbe") to its fixed coefficient in the weighted sum

  • normalize_coefficients (bool) – if True, rescale weights so they sum to one.

Return type:

AtomisticModel