Simulation Module

Simulation Runner

Simulation runner for executing MD simulations with OpenMM.

This module handles running equilibration and production phases with configurable parameters, reporters, and checkpoint management.

class polyzymd.simulation.runner.SimulationRunner(topology, system, positions, working_dir, platform='CUDA')[source]

Bases: object

Runner for executing OpenMM molecular dynamics simulations.

This class manages: - Equilibration and production phase execution - Checkpoint saving and state data reporting - Energy minimization - Unique force group assignment for energy decomposition

Example

>>> runner = SimulationRunner(
...     topology=omm_topology,
...     system=omm_system,
...     positions=omm_positions,
...     working_dir="output/",
... )
>>> runner.minimize()
>>> runner.run_equilibration(temperature=300, duration_ns=0.5)
>>> runner.run_production(temperature=300, duration_ns=100)
__init__(topology, system, positions, working_dir, platform='CUDA')[source]

Initialize the SimulationRunner.

Parameters:
  • topology (Any) – OpenMM Topology.

  • system (openmm.System) – OpenMM System.

  • positions (Any) – Initial positions with units.

  • working_dir (str | Path) – Working directory for output files.

  • platform (str) – Compute platform (CUDA, OpenCL, CPU).

property simulation: openmm.app.Simulation | None

Get the current OpenMM Simulation object.

property working_dir: Path

Get the working directory path.

property history: Dict[str, Any]

Get the simulation history.

minimize(max_iterations=1000, tolerance=10.0)[source]

Run energy minimization.

Parameters:
  • max_iterations (int) – Maximum iterations (0 = until convergence).

  • tolerance (float) – Energy tolerance in kJ/mol/nm.

Returns:

Final potential energy in kJ/mol.

Return type:

float

run_equilibration(temperature, duration_ns=None, num_samples=10, timestep_fs=2.0, friction=1.0, output_prefix='equilibration', config=None)[source]

Run equilibration phase.

Supports two modes: 1. Parameter-based (legacy): Pass duration_ns, num_samples, etc. 2. Config-based: Pass config for automatic mode selection

When config is provided and uses staged equilibration, position restraints and temperature ramping are handled automatically. Component information is derived from the topology’s chain IDs.

Parameters:
  • temperature (float) – Temperature in Kelvin

  • duration_ns (Optional[float]) – Duration in nanoseconds (required for legacy mode)

  • num_samples (int) – Number of trajectory frames to save

  • timestep_fs (float) – Time step in femtoseconds

  • friction (float) – Friction coefficient in 1/ps

  • output_prefix (str) – Prefix for output files

  • config (Optional['SimulationPhasesConfig']) – SimulationPhasesConfig for config-based dispatch

Returns:

Dictionary with equilibration results

Raises:

ValueError – If neither config nor duration_ns is provided

Return type:

Dict[str, Any]

run_equilibration_stage(stage, reference_positions, atom_group_resolver, stage_index, default_timestep=2.0, default_friction=1.0)[source]

Run a single equilibration stage with optional position restraints.

This method runs one stage of a multi-stage equilibration protocol. It supports: - Position restraints on predefined atom groups - Temperature ramping (simulated annealing) - NVT or NPT ensembles

Parameters:
  • stage (EquilibrationStageConfig) – EquilibrationStageConfig with stage settings

  • reference_positions (Any) – Positions to restrain atoms to (typically post-minimization)

  • atom_group_resolver (AtomGroupResolver) – Resolver for predefined atom group names

  • stage_index (int) – Index of this stage (for output naming)

  • default_timestep (float) – Default time step in fs if not specified in stage

  • default_friction (float) – Default friction coefficient in 1/ps

Returns:

Dictionary with stage results

Return type:

Dict[str, Any]

run_staged_equilibration(stages, atom_group_resolver, target_temperature)[source]

Run complete multi-stage equilibration protocol.

This method executes a sequence of equilibration stages, each with potentially different: - Temperature (constant or ramping) - Position restraints on different atom groups - Thermodynamic ensemble (NVT/NPT)

Positions carry over between stages, and restraint forces are added/removed as needed.

Parameters:
  • stages (List['EquilibrationStageConfig']) – List of EquilibrationStageConfig objects

  • atom_group_resolver (AtomGroupResolver) – Resolver for predefined atom group names

  • target_temperature (float) – Final target temperature (for logging)

Returns:

Dictionary with all stage results and summary

Return type:

Dict[str, Any]

run_production(temperature, duration_ns, num_samples=2500, timestep_fs=2.0, friction=1.0, pressure=1.0, barostat_frequency=25, output_prefix='production', segment_index=0)[source]

Run NPT production simulation.

Parameters:
  • temperature (float) – Temperature in Kelvin.

  • duration_ns (float) – Duration in nanoseconds.

  • num_samples (int) – Number of trajectory frames to save.

  • timestep_fs (float) – Time step in femtoseconds.

  • friction (float) – Friction coefficient in 1/ps.

  • pressure (float) – Pressure in atmospheres.

  • barostat_frequency (int) – Barostat update frequency.

  • output_prefix (str) – Prefix for output files.

  • segment_index (int) – Segment index for daisy-chaining.

Returns:

Dictionary with phase results.

Return type:

Dict[str, Any]

save_history(path=None)[source]

Save simulation history to JSON.

Parameters:

path (str | Path | None) – Output path (defaults to working_dir/simulation_history.json).

load_checkpoint(checkpoint_path)[source]

Load state from a checkpoint file.

Parameters:

checkpoint_path (str | Path) – Path to checkpoint file.

Continuation Manager

Continuation manager for resuming MD simulations from checkpoints.

This module handles loading simulation state from previous segments and continuing production simulations for daisy-chain workflows.

polyzymd.simulation.continuation.quantity_from_dict(qdict)[source]

Convert serialized quantity dictionary back to OpenMM Quantity.

Parameters:

qdict (Dict[str, Any]) – Dictionary with __values__ containing value and unit.

Returns:

OpenMM Quantity with appropriate units.

Return type:

openmm.unit.Quantity

class polyzymd.simulation.continuation.ContinuationManager(working_dir, segment_index)[source]

Bases: object

Manager for continuing MD simulations from previous segments.

This class handles loading state from previous production segments and continuing the simulation, primarily for daisy-chain workflows on HPC clusters.

Example

>>> manager = ContinuationManager(
...     working_dir="simulation_output/",
...     segment_index=2,  # Continuing to segment 2
... )
>>> manager.load_previous_state()
>>> manager.run_segment(duration_ns=20.0, num_samples=250)
__init__(working_dir, segment_index)[source]

Initialize the ContinuationManager.

Parameters:
  • working_dir (str | Path) – Working directory containing simulation outputs.

  • segment_index (int) – Current segment index (1-based).

property working_dir: Path

Get the working directory.

property segment_index: int

Get the current segment index.

property simulation: openmm.app.Simulation | None

Get the OpenMM Simulation object.

load_previous_state()[source]

Load state from the previous segment.

This loads the system, topology, and parameters from the previous production segment.

Raises:

FileNotFoundError – If required files are missing.

run_segment(duration_ns, num_samples=250, timestep_fs=2.0)[source]

Run the continuation segment.

Parameters:
  • duration_ns (float) – Duration of this segment in nanoseconds.

  • num_samples (int) – Number of trajectory frames to save.

  • timestep_fs (float) – Time step in femtoseconds.

Returns:

Dictionary with segment results.

Return type:

Dict[str, Any]

polyzymd.simulation.continuation.main()[source]

Main entry point for continuation script.

Returns:

Exit code (0 for success, 1 for failure).

Return type:

int