Force Field predictor

class mlip.models.predictor.ForceFieldPredictor(mlip_network: ~flax.linen.module.Module, predict_stress: bool, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Flax module for a force field predictor.

The apply function of this predictor returns the force field function used basically everywhere in the rest of the code base. This module is initialized from an already constructed MLIP model network module and a boolean whether to predict stress properties.

mlip_network

The MLIP network.

Type:

flax.linen.module.Module

predict_stress

Whether to predict stress properties. If false, only energies and forces are computed.

Type:

bool

__call__(graph: GraphsTuple) → Prediction

Returns a Prediction dataclass of properties based on an input graph.

Note: The stress-related properties and “pressure” have not yet been thoroughly tested by us, so see them as an experimental feature for now.

Parameters:

graph – The input graph.

Returns:

The properties as a Prediction object that may include “energy”, “forces”, “stress”, “stress_cell”, “stress_forces”, and “pressure”. Only the first two exist if predict_stress=False is set for this module.