Force Field predictor

class dipm.models.predictor.ForceFieldPredictor(*args: Any, **kwargs: Any)

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.

force_model

The MLIP network to use in this force field.

predict_stress

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

__call__(graph: GraphsTuple, rngs: Rngs | None = None, ctx: dict | None = None) → Prediction

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

Parameters:

• graph – The input graph.

• rngs (optional) – The random number generator for dropout. None for eval.

• ctx (optional) – The context dictionary obtained from pre-calling the force model.

Returns:

The properties as a Prediction object including “energy” and “forces”. It may also include “stress” and “pressure” when predict_stress=True.