xnn.common.models.outputs.ForceStressOutput#

class xnn.common.models.outputs.ForceStressOutput(model, compute_forces=True, compute_stress=False)[source]#

Bases: Module

Wrap a model to add conservative forces and the stress tensor.

Wrapping any registered model augments its output dict with autograd-derived quantities, so forces and stress are computed uniformly in one place rather than in each model.

Forces are the negative gradient of the energy with respect to positions, F = -dE/dr. The stress is obtained with the symmetric-strain trick: a strain tensor eps of shape (B, 3, 3) initialised to zero is introduced, positions and cell are displaced by eps, and the stress is sigma = (1/V) dE/deps. This matches the NequIP/MACE convention expected by ASE and LAMMPS.

Parameters:

  • model (torch.nn.Module) – The wrapped model whose energy output is differentiated. Its cutoff attribute, if present, is exposed as self.cutoff.

  • compute_forces (bool, optional) – Whether to compute forces (default True).

  • compute_stress (bool, optional) – Whether to compute the stress tensor (default False). Stress is only computed for periodic structures (those with a non-None cell).

Variables:

  • model (torch.nn.Module) – The wrapped model.

  • compute_forces (bool) – Whether forces are computed.

  • compute_stress (bool) – Whether stress is computed.

  • cutoff (float or None) – The wrapped model’s radial cutoff, or None if it has none.

Examples

>>> model = ForceStressOutput(build_model(cfg.model), compute_stress=True)
>>> out = model(graph)   # out has energy, forces, (stress)
forward(data)[source]#

Run the wrapped model and add forces and stress via autograd.

Parameters:

data (AtomicGraph) – The batched atomic graph to evaluate.

Returns:

The wrapped model’s output dict, augmented with:

energy

Per-structure energies of shape (B,) (from the model).

forces

Present when compute_forces is set; per-atom forces of shape (N, 3) computed as F = -dE/dr (autograd with respect to positions).

stress

Present when compute_stress is set and the batch has a cell; the symmetric stress tensor of shape (B, 3, 3) computed as sigma = (1/V) dE/deps via the symmetric-strain trick, using the NequIP/MACE convention.

Return type:

dict of str to torch.Tensor

Notes

If the model’s energy is independent of positions or strain (e.g. a T = 0 MACE whose energy is only a per-atom reference), the corresponding gradient is unused and the associated forces/stress default to zeros instead of raising.