MLIP Network: Base Class

class mlip.models.mlip_network.MLIPNetwork(config: MLIPNetworkConfig, dataset_info: DatasetInfo, parent: Module | Scope | _Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Base class for GNN node-wise energy models.

Energy models deriving from this class return node-wise contributions to the total energy, from the edge vectors of a graph, the atomic species of the nodes, and the edges themselves passed as senders and receivers indices.

Our MLIP models are validated with Pydantic, and hold a reference to their .Config class describing the set of hyperparameters.

__call__(graph: Graph) → Graph

Calculate node-wise energy contributions and update the input Graph.

Parameters:

graph – Input Graph object containing edge vectors, node species, senders, receivers, and node mask.

Returns:

The updated Graph with per-node energy contributions stored in the “energy” node feature.