Core Module

Restraints

Restraint definitions and application for OpenMM systems.

This module provides classes for defining and applying various types of restraints (flat-bottom, harmonic, etc.) to OpenMM simulations.

class polyzymd.core.restraints.RestraintType(value)[source]

Bases: str, Enum

Types of restraints that can be applied.

FLAT_BOTTOM = 'flat_bottom'

HARMONIC = 'harmonic'

UPPER_WALL = 'upper_wall'

LOWER_WALL = 'lower_wall'

class polyzymd.core.restraints.AtomSelection(selection, description=None)[source]

Bases: object

Represents a selection of atoms using MDAnalysis-style syntax.

This class provides a flexible way to specify atoms for restraints using selection strings that are compatible with MDAnalysis.

Variables:

selection (str) – MDAnalysis-compatible selection string
description (str | None) – Human-readable description of what this selects

Example

>>> sel = AtomSelection("resid 77 and name OG", "Catalytic serine oxygen")
>>> indices = sel.resolve(topology)

selection: str

description: str | None = None

resolve(topology)[source]

Resolve the selection to atom indices.

For OpenMM topologies, we parse the selection string and find matching atoms. Supports basic MDAnalysis-style selections.

Parameters:: topology (openmm.app.Topology) – OpenMM Topology object
Returns:: List of atom indices matching the selection
Raises:: ValueError – If selection syntax is invalid or no atoms match
Return type:: List[int]

__init__(selection, description=None)

class polyzymd.core.restraints.RestraintDefinition(restraint_type, name, atom1, atom2, distance=<factory>, force_constant=<factory>, enabled=True)[source]

Bases: object

Definition of a single restraint to be applied to a system.

Variables:

restraint_type (polyzymd.core.restraints.RestraintType) – Type of restraint (flat_bottom, harmonic, etc.)
name (str) – Human-readable identifier
atom1 (polyzymd.core.restraints.AtomSelection) – First atom selection
atom2 (polyzymd.core.restraints.AtomSelection) – Second atom selection
distance (openmm.unit.Quantity) – Target or threshold distance
force_constant (openmm.unit.Quantity) – Force constant for the restraint
enabled (bool) – Whether this restraint should be applied

Example

>>> restraint = RestraintDefinition(
...     restraint_type=RestraintType.FLAT_BOTTOM,
...     name="catalytic_serine",
...     atom1=AtomSelection("resid 77 and name OG"),
...     atom2=AtomSelection("resname LIG and name C12"),
...     distance=3.3 * angstrom,
...     force_constant=10000 * kilojoule_per_mole / nanometer**2
... )

restraint_type: RestraintType

name: str

atom1: AtomSelection

atom2: AtomSelection

distance: openmm.unit.Quantity

force_constant: openmm.unit.Quantity

enabled: bool = True

apply(topology, system)[source]

Apply this restraint to an OpenMM system.

Parameters:

topology (openmm.app.Topology) – OpenMM Topology for resolving atom selections
system (openmm.System) – OpenMM System to add the force to

Returns:

Index of the added force, or None if restraint is disabled

Return type:

int | None

__init__(restraint_type, name, atom1, atom2, distance=<factory>, force_constant=<factory>, enabled=True)

class polyzymd.core.restraints.RestraintFactory[source]

Bases: object

Factory for creating restraints from configuration.

This class bridges the configuration schema with the restraint implementation, creating RestraintDefinition objects from config.

static from_config(config)[source]

Create a RestraintDefinition from a configuration dictionary.

Parameters:: config (Dict[str, Any]) – Dictionary with restraint configuration
Returns:: RestraintDefinition instance
Return type:: RestraintDefinition

static create_flat_bottom(name, atom1_selection, atom2_selection, distance, force_constant=10000.0)[source]

Convenience method to create a flat-bottom restraint.

Parameters:

name (str) – Restraint identifier
atom1_selection (str) – Selection string for first atom
atom2_selection (str) – Selection string for second atom
distance (float) – Threshold distance in Angstroms
force_constant (float) – Force constant in kJ/mol/nm^2

Returns:

RestraintDefinition for flat-bottom potential

Return type:

RestraintDefinition

static create_harmonic(name, atom1_selection, atom2_selection, distance, force_constant=10000.0)[source]

Convenience method to create a harmonic restraint.

Parameters:

name (str) – Restraint identifier
atom1_selection (str) – Selection string for first atom
atom2_selection (str) – Selection string for second atom
distance (float) – Equilibrium distance in Angstroms
force_constant (float) – Force constant in kJ/mol/nm^2

Returns:

RestraintDefinition for harmonic potential

Return type:

RestraintDefinition

polyzymd.core.restraints.apply_restraints(restraints, topology, system)[source]

Apply multiple restraints to a system.

Parameters:

restraints (List[RestraintDefinition]) – List of restraint definitions
topology (openmm.app.Topology) – OpenMM Topology for resolving selections
system (openmm.System) – OpenMM System to modify

Returns:

List of force indices for the added restraints

Return type:

List[int]

Simulation Parameters

Simulation parameter dataclasses for OpenMM integration.

This module provides dataclasses that bridge the configuration system with OpenMM’s simulation parameters, supporting serialization to JSON for checkpoint/restart capabilities.

polyzymd.core.parameters.quantity_to_dict(q)[source]

Convert an OpenMM Quantity to a serializable dictionary.

Parameters:: q (openmm.unit.Quantity) – OpenMM Quantity with units
Returns:: Dictionary with value and unit information
Return type:: Dict[str, Any]

polyzymd.core.parameters.quantity_from_dict(d)[source]

Restore an OpenMM Quantity from a serialized dictionary.

Parameters:: d (Dict[str, Any]) – Dictionary with value and unit information
Returns:: OpenMM Quantity
Return type:: openmm.unit.Quantity

class polyzymd.core.parameters.ThermostatParameters(temperature=<factory>, timescale=<factory>, thermostat='LangevinMiddle')[source]

Bases: object

Parameters for temperature control.

Variables:

temperature (openmm.unit.Quantity) – Target temperature
timescale (openmm.unit.Quantity) – Coupling timescale (friction coefficient)
thermostat (str) – Type of thermostat to use

temperature: openmm.unit.Quantity

timescale: openmm.unit.Quantity

thermostat: str = 'LangevinMiddle'

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

__init__(temperature=<factory>, timescale=<factory>, thermostat='LangevinMiddle')

class polyzymd.core.parameters.BarostatParameters(pressure=<factory>, temperature=<factory>, update_frequency=25, barostat='MC')[source]

Bases: object

Parameters for pressure control.

Variables:

pressure (openmm.unit.Quantity) – Target pressure
temperature (openmm.unit.Quantity) – Temperature for barostat (should match thermostat)
update_frequency (int) – Steps between barostat updates
barostat (str) – Type of barostat to use

pressure: openmm.unit.Quantity

temperature: openmm.unit.Quantity

update_frequency: int = 25

barostat: str = 'MC'

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

__init__(pressure=<factory>, temperature=<factory>, update_frequency=25, barostat='MC')

class polyzymd.core.parameters.ThermoParameters(thermostat_params=<factory>, barostat_params=None)[source]

Bases: object

Combined thermodynamic parameters.

Variables:

thermostat_params (polyzymd.core.parameters.ThermostatParameters) – Temperature control parameters
barostat_params (polyzymd.core.parameters.BarostatParameters | None) – Pressure control parameters (None for NVT/NVE)

thermostat_params: ThermostatParameters

barostat_params: BarostatParameters | None = None

property ensemble: str: Infer ensemble from parameters.

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

__init__(thermostat_params=<factory>, barostat_params=None)

class polyzymd.core.parameters.IntegratorParameters(time_step=<factory>, total_time=<factory>, num_samples=100)[source]

Bases: object

Parameters for the MD integrator.

Variables:

time_step (openmm.unit.Quantity) – Integration time step
total_time (openmm.unit.Quantity) – Total simulation time
num_samples (int) – Number of frames to save

time_step: openmm.unit.Quantity

total_time: openmm.unit.Quantity

num_samples: int = 100

property total_steps: int: Calculate total number of integration steps.

property reporting_interval: int: Calculate steps between trajectory frames.

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

__init__(time_step=<factory>, total_time=<factory>, num_samples=100)

class polyzymd.core.parameters.ReporterParameters(traj_ext='dcd', state_data=<factory>, report_checkpoint=True, report_state=True, report_trajectory=True, report_state_data=True)[source]

Bases: object

Parameters for simulation output/reporting.

Variables:

traj_ext (str) – Trajectory file extension
state_data (Dict[str, bool]) – Dictionary of state data properties to report
report_checkpoint (bool) – Whether to save checkpoint files
report_state (bool) – Whether to save state XML files
report_trajectory (bool) – Whether to save trajectory
report_state_data (bool) – Whether to save CSV state data

traj_ext: str = 'dcd'

state_data: Dict[str, bool]

report_checkpoint: bool = True

report_state: bool = True

report_trajectory: bool = True

report_state_data: bool = True

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

__init__(traj_ext='dcd', state_data=<factory>, report_checkpoint=True, report_state=True, report_trajectory=True, report_state_data=True)

class polyzymd.core.parameters.SimulationParameters(thermo_params=<factory>, integ_params=<factory>, reporter_params=<factory>)[source]

Bases: object

Complete simulation parameters combining all components.

Variables:

thermo_params (polyzymd.core.parameters.ThermoParameters) – Thermodynamic parameters
integ_params (polyzymd.core.parameters.IntegratorParameters) – Integrator parameters
reporter_params (polyzymd.core.parameters.ReporterParameters) – Reporter parameters

thermo_params: ThermoParameters

integ_params: IntegratorParameters

reporter_params: ReporterParameters

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

to_json(path)[source]

Save parameters to JSON file.

classmethod from_json(path)[source]

Load parameters from JSON file.

__init__(thermo_params=<factory>, integ_params=<factory>, reporter_params=<factory>)

class polyzymd.core.parameters.SimulationPhase(name, parameters)[source]

Bases: object

Represents a single simulation phase (e.g., equilibration, production).

Variables:

name (str) – Phase identifier
parameters (polyzymd.core.parameters.SimulationParameters) – Simulation parameters for this phase

name: str

parameters: SimulationParameters

to_dict()[source]

Serialize to dictionary.

classmethod from_dict(d)[source]

Deserialize from dictionary.

__init__(name, parameters)

polyzymd.core.parameters.create_nvt_parameters(temperature, duration, samples, time_step=2.0, thermostat_timescale=1.0)[source]

Factory function to create NVT simulation parameters.

Parameters:

temperature (float) – Temperature in Kelvin
duration (float) – Simulation duration in nanoseconds
samples (int) – Number of trajectory frames to save
time_step (float) – Integration time step in femtoseconds
thermostat_timescale (float) – Thermostat timescale in picoseconds

Returns:

Configured SimulationParameters for NVT

Return type:

SimulationParameters

polyzymd.core.parameters.create_npt_parameters(temperature, pressure, duration, samples, time_step=2.0, thermostat_timescale=1.0, barostat_frequency=25)[source]

Factory function to create NPT simulation parameters.

Parameters:

temperature (float) – Temperature in Kelvin
pressure (float) – Pressure in atmospheres
duration (float) – Simulation duration in nanoseconds
samples (int) – Number of trajectory frames to save
time_step (float) – Integration time step in femtoseconds
thermostat_timescale (float) – Thermostat timescale in picoseconds
barostat_frequency (int) – Steps between barostat updates

Returns:

Configured SimulationParameters for NPT

Return type:

SimulationParameters