Force Field¶

class mlip.models.force_field.ForceField(predictor: ForceFieldPredictor, params: dict[str, dict[str, Array | dict]], inference_context: InferenceContext | None = None)¶

An initialized force field, wrapping a ForceFieldPredictor with 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.predictors.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]]

inference_context¶

Optional context used by MoE multi-head models to configure graph-level routing globals (e.g. charge, dataset index) at inference time. When set, prepare_experts_for_inference() uses it to contract MOE expert parameters into a single-expert model.

Type:: mlip.models.inference_context.InferenceContext | None

__call__(graph: Graph) → Prediction¶

Predict physical properties of an input graph from current parameters.

See documentation of the predict() method for more information.

Parameters:: graph – The input graph.
Returns:: The predicted properties as a Prediction object extracted from the graph.

classmethod from_mlip_network(mlip_network: MLIPNetwork, required_properties: Properties | None = None, seed: int = 42, inference_context: InferenceContext | None = None) → 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.
required_properties – The properties that are required from this model, which means these will end up in the Prediction objects that are returned from the __call__ method. Error is raised if the MLIPNetwork is not able to compute these required properties. Default is None which means that a default Properties object will be constructed (i.e., energy and forces).
seed – The initialization seed for the parameters. Default is 42.
inference_context – Additional context that will be placed on the graph during inference, such as dataset_idx.

Returns:

The initialized instance of the force field.

classmethod init(predictor: ForceFieldPredictor, seed: int = 42, inference_context: InferenceContext | None = None) → Self¶

Initialize force field parameters using a random seed.

Parameters:

predictor – The force field predictor.
seed – The seed to use for generating initial random parameters.
inference_context – Additional context that will be placed on the graph during inference, such as dataset_idx.

Returns:

The initialized instance of the force field with random parameters.

replace_config(**kwargs) → Self¶

Returns an updated force field object with updated fields in the model config.

Updates fields in the config according to the passed kwargs.

Returns:: An updated force field instance.

replace_required_properties(required_properties: Properties) → Self¶

Returns an updated force field object with updated required properties.

This function re-evaluates the energy head and predictor class based on the new required properties, so be aware that these might have changed for the returned force field instance.

Parameters:: required_properties – The required properties to set.
Returns:: An updated force field instance.

replace_inference_context(inference_context: InferenceContext) → Self¶

Return a copy of this force field with inference_context attached.

The context is resolved against the model’s DatasetInfo so callers can pass a partial spec (e.g. just dataset_name) and have complementary fields populated.

Parameters:: inference_context – The context to attach.
Returns:: A new force field instance with the resolved context attached.

calculate(graph: Graph) → Graph¶

Computes a forward pass of the predictor and returns the updated graph.

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

Parameters:: graph – The input graph.
Returns:: The graph updated with the calculated properties.

predict(graph: Graph) → Prediction¶

Computes a forward pass of the predictor and returns predicted properties.

Parameters:: graph – The input graph.
Returns:: The predicted properties as a Prediction object extracted from the graph.

prepare_experts_for_inference() → Self¶

Contract MoE expert parameters for a fixed inference context.

Only models that advertise is_moe_model support this operation. Routes through the model’s MoE router using the context’s globals and linearly combines expert kernels into standard dense kernels. The returned force field has config.moe=None, no expert_kernel parameters, and routing coefficients baked in.

load_model_from_zip and load_trained_force_field call this method for MoE models when an inference_context is provided.

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: MLIPNetworkConfig¶: Return configuration of the underlying MLIP model.

property dataset_info: DatasetInfo¶: Return dataset info stored in the MLIP network.

classmethod get_energy_head(config: MLIPNetworkConfig, required_properties: Properties) → Callable[[Graph], Array]¶

Returns the appropriate energy computation head function.

A model trained with use_coulomb_term=True uses the Coulomb head, otherwise the standard energy head is used.

Parameters:

config – The configuration of the model.
required_properties – The properties required by the predictor.

Returns:

The selected function for computing the energy from a graph object.

classmethod get_predictor_class(required_properties: Properties, config: MLIPNetworkConfig) → type[ForceFieldPredictor]¶

Returns the appropriate predictor class based on the required properties.

This method can be used to customize which predictor class is instantiated, depending on the properties needed for the force field.

Parameters:

required_properties – The properties required by the predictor.
config – The configuration of the model.

Returns:

The selected predictor class.

classmethod validate_properties(required_properties: Properties, mlip_available_properties: Properties) → None¶

Validates that the mlip network provides all properties required by the user.

Raises an error if any required property is not available in the mlip network.

Parameters:

required_properties – The set of properties needed.
mlip_available_properties – The set of properties the mlip network supports.

Raises:

ValueError – If a required property is unavailable.