Force Field

class mlip.models.force_field.ForceField(predictor: ForceFieldPredictor, params: dict[str, dict[str, Array | dict]])

An initialized force field, wrapping a ForceFieldPredictor with parameters.

ForceField instances can be called on an input graph to obtain the predictions of the force field. Internally, this means we compose the predictor.apply method inherited from flax.linen.Module with the current dictionary of learnable parameters.

Only the cutoff_distance and allowed_atomic_numbers properties are subject to duck-typing in the simulation engine. Users are therefore free to provide any other force field callable that provides this simple interface.

predictor

The ForceFieldPredictor which derives forces and stress from the underlying MLIPNetwork energy model.

Type:

mlip.models.predictor.ForceFieldPredictor

params

The dictionary of learnable parameters. If an integer is passed instead, it will be used as seed for the random number generator to initialize model parameters.

Type:

dict[str, dict[str, jax.Array | dict]]

__call__(graph: GraphsTuple) → Prediction

Predict physical properties of an input graph from current parameters.

See documentation of the __call__() method of ForceFieldPredictor for more information on the returned object.

classmethod from_mlip_network(mlip_network: MLIPNetwork, predict_stress: bool = False, seed: int = 42) → Self

Initializes a force field from an MLIPNetwork instance with random parameters.

This is an alternative constructor to this dataclass, but the preferred one in a typical MLIP pipeline.

Parameters:

  • mlip_network – The MLIP network to use in this force field.

  • predict_stress – Whether to predict stress properties with this force field. Default is False.

  • seed – The initialization seed for the parameters. Default is 42.

Returns:

The initialized instance of the force field.

classmethod init(predictor: ForceFieldPredictor, seed: int = 42)

Initialize force field parameters from random number generator seed.

Parameters:

  • predictor – The force field predictor.

  • seed – The integer seed to generate initial random parameters from.

Returns:

The initialized instance of the force field (with random parameters).

property cutoff_distance: float

Cutoff distance in Angstrom the model was built for.

property allowed_atomic_numbers: set[int]

Set of atomic numbers supported by the model.

property config: BaseModel

Return configuration of the underlying MLIP model.

property dataset_info: DatasetInfo

Return dataset info stored in the MLIP network.