MLIP Network: Base Class

class mlip.models.mlip_network.MLIPNetwork(config: BaseModel, 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__(edge_vectors: Array, node_species: Array, senders: Array, receivers: Array) → Array

Compute node-wise energy summands. This function must be overridden by the implementation of MLIPNetwork.