Conservative Predictor

class mlip.models.predictors.conservative_predictor.ConservativePredictor(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)

Implementation of the conservative force field predictor.

The conservative predictor computes forces and stress as first-order derivatives of the energy with respect to atomic positions and strains, respectively.

__call__(graph: Graph) → Graph

Evaluates the predictor on a given graph.

Computes the required properties, and updates the input graph with these quantities. If only energy is required, auto-differentiation is skipped to minimize computational cost.

Parameters:

graph – The input graph.

Returns:

An updated graph containing all predicted properties.

compute_forces_and_stress(positions: Array, strains: Array, graph: Graph) → tuple[Array, Graph]

Computes forces and stress of the input graph.

This method takes in atomic positions, strains, and the graph, computes the gradients of the energy (forces), and, if required, also computes the stress tensor and pressure.

Parameters:

• positions – The positions of the nodes.

• strains – The strains to apply.

• graph – The input graph.

Returns:

Computed forces and the updated graph containing the predicted properties.