xnn.common.models.base.InteratomicPotential

class xnn.common.models.base.InteratomicPotential(*args, **kwargs)

Bases: Module

Abstract base class every interatomic-potential model must subclass.

A model is an torch.nn.Module that maps an AtomicGraph to a dict of tensors with at least the keys "node_energy" (shape (N,), per-atom energies) and "energy" (shape (B,), per-structure energies). Models do not compute forces or stress themselves; those are added uniformly by xnn.common.models.outputs.ForceStressOutput via autograd, so every model gets correct conservative forces for free.

Notes

The radial cutoff (in Angstrom, used to build neighbor lists) is set as the instance attribute cutoff in each subclass __init__. It is intentionally not declared as a bare class annotation here, so subclasses remain compatible with torch.jit.script().

abstractmethod forward(data)

Compute per-atom and per-structure energies for a graph.

Parameters:

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

Returns:

A mapping containing at least "node_energy" of shape (N,) (per-atom energies) and "energy" of shape (B,) (per-structure energies), where N is the number of atoms and B the number of structures in the batch.

Return type:

dict of str to torch.Tensor

Raises:

NotImplementedError – Always, unless overridden by a concrete subclass.

abstractmethod classmethod from_config(cfg)

Construct a model instance from a configuration dataclass.

Parameters:

cfg (ModelConfig) – The model configuration dataclass describing hyperparameters.

Returns:

A newly constructed model instance.

Return type:

InteratomicPotential

Raises:

NotImplementedError – Always, unless overridden by a concrete subclass.

aggregate_energy(node_energy, data)

Sum per-atom energies into per-structure energies (locality).

Parameters:

• node_energy (torch.Tensor) – Per-atom energies of shape (N,).

• data (AtomicGraph) – The batched graph, providing num_graphs and the batch vector that maps each atom to its structure index.

Returns:

Per-structure total energies of shape (B,), where B == data.num_graphs.

Return type:

torch.Tensor