FEP Simulation Sampler¶

class mlip.simulation.fep.sampler.FEPSimulationSampler(atoms: Atoms | list[Atoms], force_field: ForceField, config: FEPSimulationSamplerConfig, checkpoint_dir_upload_fun: Callable[[Path], None] | None = None)¶

Orchestrator for running Free Energy Perturbation (FEP) simulations.

Runs multiple simulation engines for different lambda values in parallel across local devices, and also supports multi-host deployments on TPU.

Note that for multi-host deployment, jax.distributed.initialize() must be called on every host before any JAX operations and before constructing this sampler.

__init__(atoms: Atoms | list[Atoms], force_field: ForceField, config: FEPSimulationSamplerConfig, checkpoint_dir_upload_fun: Callable[[Path], None] | None = None) → None¶

Initialize the FEP sampler.

Parameters:

atoms – Starting configuration(s). Pass a list with one entry per lambda to use different starting configurations per lambda; pass a single Atoms object to use the same starting configuration for all lambdas.
force_field – The force field to use for running each alchemical simulation. Depending on config.use_alchemical_mlip, the force field is either converted to an Alchemical MLIP, which scales alchemical edges in its update equations, or the alchemical potential is computed by summing predictions of the base model on the two endstates.
config – Simulation configuration.
checkpoint_dir_upload_fun – An optional function to upload the checkpoints directory to remote storage after each checkpoint save.

Raises:

NotImplementedError – If the input force field was trained with use_coulomb_term=True, which is not supported for FEP simulations.

run() → None¶: Run all lambda windows for FEP with multi-device parallelism.

attach_logger(logger: Callable[[FEPSimulationState, int | None], None]) → None¶

Adds a logger to the list of loggers of the sampler.

The logger must take two arguments: an FEPSimulationState and an integer referencing which engine the state refers to (or None if global state).

Parameters:: logger – The logger to add.

restore_checkpoint(checkpoint_dir: Path) → None¶

Restore complete sampler state from a local checkpoint directory.

Restores per-lambda jax_md_state leaves, the replica exchange PRNG key, and infers episode_offset from the saved episode count.

Note that this method must read from a local directory. If the checkpoint directory is in remote storage, download it before calling this method.

Parameters:: checkpoint_dir – Local directory written by _save_checkpoint.

property engine_states: tuple[FEPSimulationState, ...]¶

The state of each lambda window’s engine, in schedule order.

Used to extract per-engine outputs after running an FEP simulation. Each engine state contains a per_lambda_energies array, which is required for computing a free energy estimate with BAR or MBAR.

property replica_exchange_log: list[dict[str, float | None]] | None¶

The log of attempted replica exchange swaps throughout the simulation.

Used to evaluate swap frequency and replica exchange paths after running an FEP simulation.

State for a single simulation engine in a multi-engine FEP simulation.

Inherits all the fields from SimulationState and adds the ones below.

per_lambda_energies¶

The per-lambda energies along the simulation.

Type:: numpy.ndarray | None

replica_exchange_log¶

The log of the replica exchange attempts. This is populated only for a shared global state stored by the sampler.

Type:: list[dict[str, float | None]] | None

final_positions¶

The final positions of the simulation.

Type:: numpy.ndarray | None

final_velocities¶

The final velocities of the simulation.

Type:: numpy.ndarray | None

final_cell¶

The final cell of the simulation.

Type:: numpy.ndarray | None

__init__(atomic_numbers: ndarray | list[ndarray] | None = None, positions: ndarray | list[ndarray] | None = None, forces: ndarray | list[ndarray] | None = None, velocities: ndarray | list[ndarray] | None = None, temperature: ndarray | list[ndarray] | None = None, kinetic_energy: ndarray | list[ndarray] | None = None, cell: ndarray | list[ndarray] | None = None, potential_energy: ndarray | list[ndarray] | None = None, partial_charges: ndarray | list[ndarray] | None = None, step: int = 0, compute_time_seconds: float = 0.0, per_lambda_energies: ndarray | None = None, replica_exchange_log: list[dict[str, float | None]] | None = None, final_positions: ndarray | None = None, final_velocities: ndarray | None = None, final_cell: ndarray | None = None) → None¶