xnn.dnn.models.base.DescriptorPotential#

class xnn.dnn.models.base.DescriptorPotential(featurizer, species, hidden=(64, 64), activation='silu', bias=True, atomic_energies=None)[source]#

Bases: InteratomicPotential

Shared body for descriptor-based potentials (HDNNP, ANI).

Composition: featurizer (AtomicGraph -> per-atom descriptor) + per-element atomic networks (+ optional per-species self atomic energies). Reuse by passing any invariant Featurizer; HDNNP and ANI are thin subclasses that differ in the featurizer and per-element architecture.

Parameters:

  • featurizer (Featurizer) – Invariant featurizer mapping an AtomicGraph to a per-atom descriptor of shape (N, featurizer.output_dim). Its cutoff sets the model’s neighbour-list cutoff.

  • species (sequence of int) – Atomic numbers to build per-element networks for.

  • hidden (sequence of int or dict[int, sequence of int], optional) – Hidden-layer widths of each per-element MLP (shared sequence or per-Z dict), by default (64, 64).

  • activation (str or torch.nn.Module, optional) – Hidden-layer activation, by default "silu".

  • bias (bool, optional) – Whether the linear layers carry a bias, by default True.

  • atomic_energies (sequence of float or None, optional) – Per-species self atomic energy added to each atom’s contribution (aligned with species). None (default) adds nothing.

Variables:

  • featurizer (Featurizer) – The composed featurizer.

  • cutoff (float) – Neighbour-list cutoff, taken from featurizer.cutoff.

  • species (list[int]) – The elements handled.

  • element_nets (_ElementNetworks) – Per-element atomic MLPs.

forward(data)[source]#

Compute per-atom and total energies for a batch of structures.

Parameters:

data (AtomicGraph) – Batched atomic graph passed to the featurizer.

Returns:

Dictionary with "node_energy" (per-atom energy including the self-energy shift, shape (N,)), "energy" (per-structure total from aggregate_energy) and "node_features" (the descriptor).

Return type:

dict[str, Tensor]