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, config=sim_config.simulation_phases)
>>> 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, config)[source]

Run equilibration phase.

Staged equilibration is required. Position restraints and temperature ramping are handled automatically. Component information is derived from the topology’s chain IDs.

Parameters:

  • temperature (float) – Temperature in Kelvin

  • config (SimulationPhasesConfig) – SimulationPhasesConfig containing equilibration stages

Returns:

Dictionary with equilibration results

Return type:

Dict[str, Any]

run_equilibration_stage(stage, reference_positions, atom_group_resolver, stage_index, default_timestep=2.0, default_friction=1.0, resume_from_step=0, resume_temperature=None)[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 - Graceful interruption and mid-stage resume

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

  • resume_from_step (int) – Step to resume from within this stage (0 = start fresh). When resuming, the simulation must already be loaded from checkpoint.

  • resume_temperature (float | None) – Current temperature when resuming a temperature-ramping stage. Only used when resume_from_step > 0 and the stage uses temperature ramping.

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. If a previous run was interrupted, completed stages are detected on disk via their checkpoint files and skipped automatically.

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, report_interval=None, checkpoint_interval_s=60.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. Ignored when report_interval is provided.

  • 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 multi-segment production.

  • report_interval (int | None) – Fixed reporter interval in steps. When provided, this overrides the per-segment total_steps // num_samples calculation to keep frame spacing uniform across segments.

  • checkpoint_interval_s (float) – Wall-time interval in seconds between portable restart checkpoints. Also controls how frequently the loop checks for SLURM preemption signals. Set to 0 to disable wall-time checkpoints (reverts to legacy behaviour).

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 the simulation for self-resubmitting HPC workflows. Each segment runs until completion or interruption (wall-time / preemption), updates the progress tracker, and the SLURM script handles resubmission.

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. It integrates with the progress tracking system to enable self-resubmitting idempotent jobs.

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 or Path) – Working directory containing simulation outputs.

  • segment_index (int) – Current segment index (0-based for first continuation after initial production, incrementing from there).

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. Recovery prefers portable state XML files (production_N_state.xml, interrupted_state.xml, or restart_state.xml) over binary .chk checkpoints. Only falls back to loadCheckpoint() when no portable state XML is available (legacy interrupted segments, hard-killed segments).

Raises:

FileNotFoundError – If required files are missing.

run_segment(duration_ns, num_samples=250, timestep_fs=2.0, report_interval=None, checkpoint_interval_s=60.0)[source]

Run the continuation segment.

Runs the simulation for the specified duration, saving trajectory frames at regular intervals. On completion, updates the progress tracker. On interruption (SIGUSR1/SIGTERM), saves interrupted state, updates the progress tracker, and raises GracefulExit.

Parameters:

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

  • num_samples (int) – Number of trajectory frames to save. Ignored when report_interval is provided.

  • timestep_fs (float) – Time step in femtoseconds.

  • report_interval (int or None) – Fixed reporter interval in steps. When provided, this overrides the per-segment total_steps // num_samples calculation to keep frame spacing uniform across segments.

  • checkpoint_interval_s (float) – Wall-time interval in seconds between portable restart checkpoints. Also controls how frequently the loop checks for SLURM preemption signals. Set to 0 to disable wall-time checkpoints (reverts to legacy behaviour).

Returns:

Dictionary with segment results.

Return type:

dict