Metadynamics Simulation Engine

class mlip.simulation.metadynamics.jax_md_metad_engine.JaxMDMetadynamicsSimulationEngine(atoms: Atoms | list[Atoms], force_field: ForceField, config: JaxMDMetadynamicsSimulationConfig)

Simulation engine for well-tempered metadynamics using the JAX-MD backend.

Extends JaxMDSimulationEngine to periodically deposit Gaussian hills along one or two collective variables (CVs), optionally applying well-tempered rescaling, wall potentials, and positional restraints.

Only non-batched simulations are supported. The metadynamics state (hill centers and heights) is carried inside MetadynamicsSystemState and injected into the Graph globals at each step so bias potentials can access it.

__init__(atoms: Atoms | list[Atoms], force_field: ForceField, config: JaxMDMetadynamicsSimulationConfig) → 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.