Inference and simulation helpers

mlipaudit.utils.inference.run_inference(atoms_list: list[Atoms], force_field: ForceField | Calculator, batch_size: int = 16) → list[Prediction | None]

Runs inference for a list of ase.Atoms objects.

If ForceField object is passed, run_batched_inference from the mlip library is used.

Parameters:

  • atoms_list – The list of ase.Atoms objects.

  • force_field – The force field.

  • batch_size – Batch size, default 16. Will only be used if force field is passed as a ForceField object.

Returns:

A list of Prediction or None objects, None when the model failed

to perform inference on a system for whatever reason.

Raises:

ValueError – If force field type is not compatible.

mlipaudit.utils.simulation.get_simulation_engine(atoms: Atoms, force_field: ForceField | Calculator, **kwargs) → JaxMDSimulationEngine | ASESimulationEngineWithCalculator | ASESimulationEngine

Returns the correct simulation engine based on the input force field type.

For MD simulations with mlip.models.ForceField objects, we return a JaxMDSimulationEngine. For energy minimizations with those objects, we return a ASESimulationEngine. For any type of simulations with ASE calculator objects, we return a ASESimulationEngineWithCalculator, which is a custom class of the mlipaudit library.

Parameters:

  • atoms – The ASE atoms.

  • force_field – The force field, either an mlip.models.ForceField or an ASE calculator.

  • **kwargs – Keyword arguments to be passed to the MD config object. Assumed to be JAX-MD based, will be modified automatically for ASE.

Returns:

The simulation engine.

Raises:

ValueError – If force field type is not compatible.

mlipaudit.utils.simulation.run_simulation(atoms: Atoms, force_field: ForceField | Calculator, **kwargs) → SimulationState | None

Run the simulation with the appropriate simulation engine based on the input force field type.

For MD simulations with mlip.models.ForceField objects, runs the simulation with JaxMDSimulationEngine. For energy minimizations with those objects, runs with an ASESimulationEngine. For any type of simulations with ASE calculator objects, runs with an ASESimulationEngineWithCalculator, which is a custom class of the mlipaudit library. If the simulation fails, the error will be caught and None will be returned.

Parameters:

  • atoms – The ASE atoms.

  • force_field – The force field, either an mlip.models.ForceField or an ASE calculator.

  • **kwargs – Keyword arguments to be passed to the MD config object. Assumed to be JAX-MD based, will be modified automatically for ASE.

Returns:

The simulation state or None if the simulation failed.

Raises:

ValueError – If force field type is not compatible.

class mlipaudit.utils.simulation.ASESimulationEngineWithCalculator(atoms: Atoms, ase_calculator: Calculator, config: ASESimulationConfig)

Class derived from mlip’s ASE simulation engine but allowing for a passed ASE calculator object.