API Reference

This reference is curated around the stable public objects that are useful when embedding Equitrain in Python code. The narrative pages remain the primary documentation for end-to-end workflows.

Torch Wrapper Interface

Bases: Module, ABC

Common interface exposed by all torch wrappers.

forward(*args) abstractmethod

Implement the model forward pass.

Implementations must return a mapping with at least an 'energy' entry of shape [batch, 1] or [batch]. Wrappers that produce forces or stresses should also return 'forces' with shape [num_atoms, 3] and 'stress' with shape [batch, 3, 3]. Additional observables such as dipoles or virials can be included as extra keys; they are forwarded to the loss and metrics stack unchanged.

Torch Model Wrappers

Bases: AbstractWrapper

Adapt a Torch MACE model to Equitrain's graph-batch interface.

The wrapper accepts Torch Geometric Data/Batch objects produced by Equitrain's Torch data loaders and returns a prediction dictionary with energy and, when enabled by loss weights, forces and stress.

atomic_numbers property

Return the atomic-number table supported by the wrapped MACE model.

atomic_energies property

Return per-element atomic reference energies from the MACE model.

r_max property writable

Return the MACE cutoff radius used for graph construction.

forward(*args)

Run the wrapped MACE model on one Torch Geometric batch.

Bases: AbstractWrapper

Wrapper for TorchANI models to be used with Equitrain.

This wrapper integrates the Atomic Neural Network (ANI) potential from the TorchANI library into the Equitrain framework. It supports energy-only training workflows and can optionally compute forces using autograd.

atomic_numbers property

Return the atomic-number table supported by the ANI model.

atomic_energies property

Return ANI self atomic energies when available.

r_max property writable

Return the ANI radial cutoff used by the AEV computer.

forward(*args)

Run the wrapped TorchANI model and normalize output keys.

Bases: AbstractWrapper

Wrapper for ORB (Orbital Materials) models.

orb-models>=0.5 expects the input as an AtomGraphs object. The torch backend operates on Torch Geometric Data/Batch instances, so this wrapper performs the conversion and exposes a uniform dictionary with energy, forces, and stress tensors to the shared training pipeline.

atomic_numbers property

Return the atomic-number table supported by the ORB model.

atomic_energies property

Return ORB atomic reference energies when available.

r_max property writable

Return the ORB neighbor cutoff radius.

forward(*args)

Run the wrapped ORB model and return energy/forces/stress tensors.

Bases: AbstractWrapper

Adapt a SevenNet Torch model to Equitrain's graph-batch interface.

atomic_numbers property

Return the atomic-number table supported by the SevenNet model.

atomic_energies property

Return atomic reference energies when provided by the model.

r_max property writable

Return the SevenNet cutoff radius used for graph construction.

forward(input)

Run SevenNet on a Torch graph batch and normalize output keys.

get_edge_vectors_and_lengths(positions, edge_index, shifts) classmethod

Return edge displacement vectors and lengths for a graph batch.

batch_voigt_to_tensor(voigts) classmethod

Convert batched Voigt stress vectors to full 3x3 stress tensors.

M3GNet is covered in Model Wrappers. Its generated API reference is omitted because the MatGL wrapper requires the optional MatGL runtime package at import time.

Fine-Tuning Wrappers

Bases: AbstractWrapper

Wrap a :class:~equitrain.backends.torch_wrappers.AbstractWrapper instance with additive residual parameters for L^2-SP fine-tuning.

The wrapped base model is frozen at its pre-trained starting point. Each base parameter is mirrored by a zero-initialized delta parameter, and the forward pass evaluates the effective parameter theta = theta_0 + delta. Optimizer weight decay on trainable deltas therefore corresponds to the L^2-SP penalty on distance from the starting weights.

When freeze_layers freezes selected semantic delta layers, Equitrain calls the configuration targeted L^2-SP (L^2-TSP): L^2-SP is applied only to the remaining trainable delta layers, while frozen layers keep delta = 0.

delta_layer_names property

Return semantic delta layer names in MACE forward order.

delta_parameters()

Iterate over all delta parameters.

named_delta_parameters()

Iterate over all named delta parameters using original parameter names.

freeze_delta_layers(freeze_layers=None)

Freeze selected semantic delta layers; None keeps all deltas trainable.

Bases: AbstractWrapper

Wrap a torch model wrapper and fine-tune selected base parameters directly.

Unlike delta or LoRA fine-tuning, this wrapper does not add adapter weights. It only controls which base parameters have gradients enabled, so export can write the updated base model without a merge step.

freeze_layer_names property

Return semantic layer names in MACE forward order.

freeze_model_layers(freeze_layers=None)

Freeze selected semantic layers; None keeps all layers trainable.

fine_tune_parameters()

Iterate over currently trainable base parameters.

named_fine_tune_parameters()

Iterate over currently trainable base parameters with original names.

Bases: AbstractWrapper

Wrap a torch model wrapper with LoRA adapters on eligible weight tensors.

Eligible tensors are parameters named *.weight with ndim >= 2. Higher-order weights are flattened to (shape[0], prod(shape[1:])) for the low-rank update and reshaped back to the original tensor shape.

JAX fine-tuning wrappers are described in Fine-Tuning. They are kept out of this generated page so the documentation build does not require the optional JAX runtime.