Force Field predictor

class mlip.models.predictors.predictor.ForceFieldPredictor(mlip_network: ~mlip.models.mlip_network.MLIPNetwork, required_properties: ~mlip.typing.properties.Properties, energy_head: ~typing.Callable[[~mlip.graph.graph.Graph], ~jax.Array] | None = None, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Base class for all force field predictors.

mlip_network

The MLIP network to use as a backbone for property prediction.

Type:

mlip.models.mlip_network.MLIPNetwork

required_properties

The properties that the predictor is required to compute.

Type:

mlip.typing.properties.Properties

energy_head

The energy head to use for aggregating MLIP outputs into energies, and/or other outputs. If None, we use a default standard energy head that sums node energies into a total energy per graph.

Type:

Callable[[mlip.graph.graph.Graph], jax.Array] | None

abstractmethod __call__(graph: Graph) → Graph

Compute required properties of the graph.