JAX-MD Simulation Engine

class mlip.simulation.jax_md.jax_md_simulation_engine.JaxMDSimulationEngine(atoms: Atoms | list[Atoms], force_field: ForceField, config: JaxMDSimulationConfig)

Simulation engine handling simulations with the JAX-MD backend.

Batched MD simulations are supported. Just pass a list of ase.Atoms objects to the constructor. See deep-dive tutorials on simulations for more information.

For energy minimization, the FIRE algorithm is used (see here).

For MD, the integrator/ensemble can be set via the md_integrator attribute (see MDIntegrator), to use the NVT-Langevin algorithm (see here); the NPT-MC-Langevin algorithm, using Langevin dynamics with a Monte-Carlo Barostat (see here); or the NVE-Velocity-Verlet algorithm (see here).

__init__(atoms: Atoms | list[Atoms], force_field: ForceField, config: JaxMDSimulationConfig) → None

Constructor that initializes the simulation state and an empty list of loggers. Engine-specific initialization is then delegated to ._initialize()

Parameters:

  • atoms – The atoms of the system to simulate.

  • force_field – The force field to use in the simulation.

  • config – The configuration/settings of the simulation.

Raises:

ValueError – if input system has only one atom.

run() → None

See documentation of abstract parent class.

For the JAX-MD backend, the simulation run is divided into episodes to ensure usage of jitting of MD/minimization steps for optimal performance.

Important: The state of the simulation is updated and the loggers are called during this function.

attach_logger(logger: Callable[[SimulationState], None]) → None

Adds a logger to the list of loggers of the simulation engine.

The logger function must only take in a single argument, the simulation state, and it shall not return anything.

Parameters:

logger – The logger to add.