From 1610300f7ae1279d4ffb0ab857f41d61fb54d99c Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 31 Jul 2023 13:16:08 +0200 Subject: [PATCH 01/66] start plain MD scripts --- .../protocols/openmm_md/plain_md_methods.py | 515 ++++++++++++++++++ .../protocols/openmm_md/plain_md_settings.py | 35 ++ openfe/protocols/openmm_utils/omm_settings.py | 458 ++++++++++++++++ 3 files changed, 1008 insertions(+) create mode 100644 openfe/protocols/openmm_md/plain_md_methods.py create mode 100644 openfe/protocols/openmm_md/plain_md_settings.py create mode 100644 openfe/protocols/openmm_utils/omm_settings.py diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py new file mode 100644 index 000000000..12e96991f --- /dev/null +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -0,0 +1,515 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +"""OpenMM MD Protocol --- :mod:`openfe.protocols.openmm_md.plain_md_methods` +=========================================================================================== + +This module implements the necessary methodology tools to run an MD +simulation using OpenMM tools. + +""" +from __future__ import annotations + +import os +import logging + +from collections import defaultdict +import gufe +import numpy.typing as npt +import openmm +from openff.units import unit +from openff.units.openmm import from_openmm, to_openmm +from typing import Optional +from openmm import app +from openmm import unit as omm_unit +import pathlib +from typing import Any, Iterable +import openmmtools +import uuid +import time + +from gufe import ( + settings, ChemicalSystem, SmallMoleculeComponent, + ProteinComponent +) +from openfe.protocols.openmm_md.plain_md_settings import ( + PlainMDProtocolSettings, SystemSettings, + SolvationSettings, OpenMMEngineSettings, + IntegratorSettings, SimulationSettings, + RepeatSettings +) +from openfe.protocols.openmm_rfe._rfe_utils import compute +from openfe.protocols.openmm_utils import ( + system_validation, settings_validation, system_creation +) + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +class PlainMDProtocolResult(gufe.ProtocolResult): + """EMPTY, Dict-like container for the output of a PlainMDProtocol""" + def __init__(self, **data): + super().__init__(**data) + # data is mapping of str(repeat_id): list[protocolunitresults] + if any(len(pur_list) > 2 for pur_list in self.data.values()): + raise NotImplementedError("Can't stitch together results yet") + + def get_estimate(self): + """Since no results as output --> returns None + + Returns + ------- + None + """ + + return None + + def get_uncertainty(self): + """Since no results as output --> returns None""" + + return None + + +class PlainMDProtocol(gufe.Protocol): + result_cls = PlainMDProtocolResult + _settings: PlainMDProtocolSettings + + @classmethod + def _default_settings(cls): + """A dictionary of initial settings for this creating this Protocol + + These settings are intended as a suitable starting point for creating + an instance of this protocol. It is recommended, however that care is + taken to inspect and customize these before performing a Protocol. + + Returns + ------- + Settings + a set of default settings + """ + return PlainMDProtocolSettings( + forcefield_settings=settings.OpenMMSystemGeneratorFFSettings(), + thermo_settings=settings.ThermoSettings( + temperature=298.15 * unit.kelvin, + pressure=1 * unit.bar, + ), + system_settings=SystemSettings(), + solvation_settings=SolvationSettings(), + engine_settings=OpenMMEngineSettings(), + integrator_settings=IntegratorSettings(), + simulation_settings=SimulationSettings( + equilibration_length=1.0 * unit.nanosecond, + production_length=5.0 * unit.nanosecond, + ), + repeat_settings=RepeatSettings(), + ) + + def _create( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: Optional[dict[str, gufe.ComponentMapping]] = None, + extends: Optional[gufe.ProtocolDAGResult] = None, + ) -> list[gufe.ProtocolUnit]: + # TODO: Extensions? + if extends: + raise NotImplementedError("Can't extend simulations yet") + + # Validate solvent component + nonbond = self.settings.system_settings.nonbonded_method + system_validation.validate_solvent(stateA, nonbond) + + # Validate protein component + system_validation.validate_protein(stateA) + + # actually create and return Units + lig_name = stateA.components['ligand'].name + # our DAG has no dependencies, so just list units + n_repeats = self.settings.repeat_settings.n_repeats + units = [PlainMDProtocolUnit( + stateA=stateA, stateB=stateB, ligandmapping=None, + settings=self.settings, + generation=0, repeat_id=int(uuid.uuid4()), + name=f'{lig_name} repeat {i} generation 0') + for i in range(n_repeats)] + + return units + + def _gather( + self, protocol_dag_results: Iterable[gufe.ProtocolDAGResult] + ) -> dict[str, Any]: + # result units will have a repeat_id and generations within this + # repeat_id + # first group according to repeat_id + unsorted_repeats = defaultdict(list) + for d in protocol_dag_results: + pu: gufe.ProtocolUnitResult + for pu in d.protocol_unit_results: + if not pu.ok(): + continue + + unsorted_repeats[pu.outputs['repeat_id']].append(pu) + + # then sort by generation within each repeat_id list + repeats: dict[str, list[gufe.ProtocolUnitResult]] = {} + for k, v in unsorted_repeats.items(): + repeats[str(k)] = sorted(v, key=lambda x: x.outputs['generation']) + + # returns a dict of repeat_id: sorted list of ProtocolUnitResult + return repeats + + +class PlainMDProtocolUnit(gufe.ProtocolUnit): + """ + Base class for plain MD simulations (NonTransformation. + """ + + def __init__(self, *, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + settings: PlainMDProtocolSettings, + generation: int, + repeat_id: int, + name: Optional[str] = None, + ): + """ + Parameters + ---------- + stateA, stateB : ChemicalSystem + the two ligand SmallMoleculeComponents to transform between. The + transformation will go from ligandA to ligandB. + ligandmapping : LigandAtomMapping + the mapping of atoms between the two ligand components + settings : settings.Settings + the settings for the Method. This can be constructed using the + get_default_settings classmethod to give a starting point that + can be updated to suit. + repeat_id : int + identifier for which repeat (aka replica/clone) this Unit is + generation : int + counter for how many times this repeat has been extended + name : str, optional + human-readable identifier for this Unit + + Notes + ----- + The mapping used must not involve any elemental changes. A check for + this is done on class creation. + """ + super().__init__( + name=name, + stateA=stateA, + stateB=stateB, + ligandmapping=None, + settings=settings, + repeat_id=repeat_id, + generation=generation + ) + + @staticmethod + def _pre_minimize(system: openmm.System, + positions: omm_unit.Quantity) -> npt.NDArray: + """ + Short CPU minization of System to avoid GPU NaNs + + Parameters + ---------- + system : openmm.System + An OpenMM System to minimize. + positions : openmm.unit.Quantity + Initial positions for the system. + + Returns + ------- + minimized_positions : npt.NDArray + Minimized positions + """ + integrator = openmm.VerletIntegrator(0.001) + context = openmm.Context( + system, integrator, + openmm.Platform.getPlatformByName('CPU'), + ) + context.setPositions(positions) + # Do a quick 100 steps minimization, usually avoids NaNs + openmm.LocalEnergyMinimizer.minimize( + context, maxIterations=100 + ) + state = context.getState(getPositions=True) + minimized_positions = state.getPositions(asNumpy=True) + return minimized_positions + + def run(self, *, dry=False, verbose=True, + scratch_basepath=None, + shared_basepath=None) -> dict[str, Any]: + """Run the MD simulation. + + Parameters + ---------- + dry : bool + Do a dry run of the calculation, creating all necessary hybrid + system components (topology, system, sampler, etc...) but without + running the simulation. + verbose : bool + Verbose output of the simulation progress. Output is provided via + INFO level logging. + scratch_basepath: Pathlike, optional + Where to store temporary files, defaults to current working directory + shared_basepath : Pathlike, optional + Where to run the calculation, defaults to current working directory + + Returns + ------- + dict + Outputs created in the basepath directory or the debug objects + (i.e. sampler) if ``dry==True``. + + Raises + ------ + error + Exception if anything failed + """ + if verbose: + logger.info("creating system") + if scratch_basepath is None: + scratch_basepath = pathlib.Path('.') + if shared_basepath is None: + # use cwd + shared_basepath = pathlib.Path('.') + + # 0. General setup and settings dependency resolution step + + # Extract relevant settings + protocol_settings: PlainMDProtocolSettings = self._inputs[ + 'settings'] + stateA = self._inputs['stateA'] + + forcefield_settings: settings.OpenMMSystemGeneratorFFSettings = \ + protocol_settings.forcefield_settings + thermo_settings: settings.ThermoSettings = \ + protocol_settings.thermo_settings + system_settings: SystemSettings = protocol_settings.system_settings + solvation_settings: SolvationSettings = \ + protocol_settings.solvation_settings + sim_settings: SimulationSettings = \ + protocol_settings.simulation_settings + timestep = protocol_settings.integrator_settings.timestep + mc_steps = protocol_settings.integrator_settings.n_steps.m + + # is the timestep good for the mass? + settings_validation.validate_timestep( + forcefield_settings.hydrogen_mass, timestep + ) + equil_steps, prod_steps = settings_validation.get_simsteps( + equil_length=sim_settings.equilibration_length, + prod_length=sim_settings.production_length, + timestep=timestep, mc_steps=mc_steps + ) + + solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) + + # 1. Create stateA system + # a. get a system generator + if sim_settings.forcefield_cache is not None: + ffcache = shared_basepath / sim_settings.forcefield_cache + else: + ffcache = None + + system_generator = system_creation.get_system_generator( + forcefield_settings=forcefield_settings, + thermo_settings=thermo_settings, + system_settings=system_settings, + cache=ffcache, + has_solvent=solvent_comp is not None, + ) + + # b. force the creation of parameters + # This is necessary because we need to have the FF generated ahead of + # solvating the system. + # Note: by default this is cached to ctx.shared/db.json so shouldn't + # incur too large a cost + for comp in small_mols: + offmol = comp.to_openff() + system_generator.create_system(offmol.to_topology().to_openmm(), + molecules=[offmol]) + + # c. get OpenMM Modeller + a dictionary of resids for each component + stateA_modeller, comp_resids = system_creation.get_omm_modeller( + protein_comp=protein_comp, + solvent_comp=solvent_comp, + small_mols=small_mols, + omm_forcefield=system_generator.forcefield, + solvent_settings=solvation_settings, + ) + + # d. get topology & positions + # Note: roundtrip positions to remove vec3 issues + stateA_topology = stateA_modeller.getTopology() + stateA_positions = to_openmm( + from_openmm(stateA_modeller.getPositions()) + ) + + # e. create the stateA System + stateA_system = system_generator.create_system( + stateA_modeller.topology, + molecules=[s.to_openff() for s in small_mols], + ) + + # 9. Create the reporter + # Get the sub selection of the system to print coords for + # selection_indices = stateA_topology.select( + # sim_settings.output_indices + # ) + + # a. Create the reporter + reporter_dcd = openmm.app.DCDReporter( + shared_basepath / sim_settings.output_filename, + sim_settings.checkpoint_interval.m, + enforcePeriodicBox=False) + # reporter = multistate.MultiStateReporter( + # storage=shared_basepath / sim_settings.output_filename, + # analysis_particle_indices=selection_indices, + # checkpoint_interval=sim_settings.checkpoint_interval.m, + # checkpoint_storage=sim_settings.checkpoint_storage, + # ) + + # 10. Get platform + platform = compute.get_openmm_platform( + protocol_settings.engine_settings.compute_platform + ) + + # 11. Set the integrator + # a. get integrator settings + integrator_settings = protocol_settings.integrator_settings + + # Validate settings + # Virtual sites sanity check - ensure we restart velocities when + # there are virtual sites in the system + if stateA_system.has_virtual_sites: + if not integrator_settings.reassign_velocities: + errmsg = ("Simulations with virtual sites without velocity " + "reassignments are unstable in openmmtools") + raise ValueError(errmsg) + + # b. create langevin integrator + integrator = openmmtools.mcmc.LangevinDynamicsMove( + timestep=to_openmm(integrator_settings.timestep), + collision_rate=to_openmm(integrator_settings.collision_rate), + n_steps=integrator_settings.n_steps.m, + reassign_velocities=integrator_settings.reassign_velocities, + n_restart_attempts=integrator_settings.n_restart_attempts, + constraint_tolerance=integrator_settings.constraint_tolerance, + ) + + try: + # # Create context caches (energy + sampler) + # energy_context_cache = openmmtools.cache.ContextCache( + # capacity=None, time_to_live=None, platform=platform, + # ) + # + # sampler_context_cache = openmmtools.cache.ContextCache( + # capacity=None, time_to_live=None, platform=platform, + # ) + # + # sampler.energy_context_cache = energy_context_cache + # sampler.sampler_context_cache = sampler_context_cache + + if not dry: # pragma: no-cover + + simulation = openmm.app.Simulation( + stateA_topology, + stateA_system, + integrator, + platform=platform + ) + simulation.context.setPositions(stateA_positions) + + # minimize + if verbose: + logger.info("minimizing systems") + + simulation.minimizeEnergy() + # sampler.minimize( + # max_iterations=sim_settings.minimization_steps) + + # equilibrate + if verbose: + logger.info("equilibrating systems") + + simulation.context.setVelocitiesToTemperature( + thermo_settings.temperature * openmm.unit.kelvin) + simulation.step(equil_steps) + + # sampler.equilibrate( + # int(equil_steps / mc_steps)) # type: ignore + + # production + if verbose: + logger.info("running production phase") + + simulation.reporters.append(reporter_dcd) + + t0 = time.time() + simulation.step(prod_steps) + t1 = time.time() + logger.info(f"Completed simulation in {t1 - t0} seconds") + # sampler.extend(int(prod_steps / mc_steps)) # type: ignore + + # # calculate estimate of results from this individual unit + # ana = multistate.MultiStateSamplerAnalyzer(reporter) + # est, _ = ana.get_free_energy() + # est = (est[0, -1] * ana.kT).in_units_of( + # omm_unit.kilocalories_per_mole) + # est = ensure_quantity(est, 'openff') + # ana.clear() # clean up cached values + + nc = shared_basepath / sim_settings.output_filename + chk = shared_basepath / sim_settings.checkpoint_storage + else: + # clean up the reporter file + fns = [shared_basepath / sim_settings.output_filename, + shared_basepath / sim_settings.checkpoint_storage] + for fn in fns: + os.remove(fn) + finally: + # close reporter when you're done, prevent file handle clashes + reporter_dcd.close() + del reporter_dcd + + # # clear GPU contexts + # # replace with above + # for context in list(energy_context_cache._lru._data.keys()): + # del energy_context_cache._lru._data[context] + # for context in list(sampler_context_cache._lru._data.keys()): + # del sampler_context_cache._lru._data[context] + # # cautiously clear out the global context cache too + # for context in list( + # openmmtools.cache.global_context_cache._lru._data + # .keys()): + # del openmmtools.cache.global_context_cache._lru._data[ + # context] + # + # del sampler_context_cache, energy_context_cache + + if not dry: + del integrator, simulation + + if not dry: # pragma: no-cover + return { + 'nc': nc, + 'last_checkpoint': chk, + } + else: + return {'debug': {'simulation': simulation}} + + def _execute( + self, ctx: gufe.Context, **kwargs, + ) -> dict[str, Any]: + outputs = self.run(scratch_basepath=ctx.scratch, + shared_basepath=ctx.shared) + + return { + 'repeat_id': self._inputs['repeat_id'], + 'generation': self._inputs['generation'], + **outputs + } diff --git a/openfe/protocols/openmm_md/plain_md_settings.py b/openfe/protocols/openmm_md/plain_md_settings.py new file mode 100644 index 000000000..edf022949 --- /dev/null +++ b/openfe/protocols/openmm_md/plain_md_settings.py @@ -0,0 +1,35 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +"""Settings class for plain MD Protocols using OpenMM + OpenMMTools + +This module implements the settings necessary to run MD simulations using +:class:`openfe.protocols.openmm_md.plain_md_methods.py` + +""" +from openfe.protocols.openmm_utils.omm_settings import ( + Settings, SystemSettings, + SolvationSettings, OpenMMEngineSettings, + IntegratorSettings, SimulationSettings, RepeatSettings +) + + +class PlainMDProtocolSettings(Settings): + class Config: + arbitrary_types_allowed = True + + # Things for creating the systems + system_settings: SystemSettings + solvation_settings: SolvationSettings + + # MD Engine things + engine_settings: OpenMMEngineSettings + + # Sampling State defining things + integrator_settings: IntegratorSettings + + # Simulation run settings + simulation_settings: SimulationSettings + + # Setting number of repeats of md simulation + repeat_settings : RepeatSettings \ No newline at end of file diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py new file mode 100644 index 000000000..d4e503a64 --- /dev/null +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -0,0 +1,458 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Standardised set of settings for OpenMM based Protocols. +""" +from __future__ import annotations + +import abc +from typing import Optional, Union +from pydantic import Extra, validator, BaseModel, PositiveFloat, Field +from openff.units import unit +import os + + +if os.environ.get('SPHINX', False): #pragma: no cover + class SettingsBaseModel(BaseModel): + class Config: + extra = Extra.forbid + validate_assignment = True + arbitrary_types_allowed = True + smart_union = True + + + class ThermoSettings(SettingsBaseModel): + """Settings for thermodynamic parameters. + + .. note:: + No checking is done to ensure a valid thermodynamic ensemble is + possible. + """ + + temperature = 298.15 * unit.kelvin + """Simulation temperature, default units kelvin""" + + pressure = 1.0 * unit.bar + """Simulation pressure, default units standard atmosphere (atm)""" + + ph: Union[PositiveFloat, None] = Field(None, description="Simulation pH") + redox_potential: Optional[float] = Field( + None, description="Simulation redox potential" + ) + + + class BaseForceFieldSettings(SettingsBaseModel, abc.ABC): + """Base class for ForceFieldSettings objects""" + ... + + + class OpenMMSystemGeneratorFFSettings(SettingsBaseModel): + """Parameters to set up the force field with OpenMM ForceFields + + .. note:: + Right now we just basically just grab what we need for the + :class:`openmmforcefields.system_generators.SystemGenerator` + signature. See the `OpenMMForceField SystemGenerator documentation`_ + for more details. + + + .. _`OpenMMForceField SystemGenerator documentation`: + https://github.com/openmm/openmmforcefields#automating-force-field-management-with-systemgenerator + """ + constraints: Optional[str] = 'hbonds' + """Constraints to be applied to system. + One of 'hbonds', 'allbonds', 'hangles' or None, default 'hbonds'""" + + rigid_water: bool = True + remove_com: bool = False + hydrogen_mass: float = 3.0 + """Mass to be repartitioned to hydrogens from neighbouring + heavy atoms (in amu), default 3.0""" + + forcefields: list[str] = [ + "amber/ff14SB.xml", # ff14SB protein force field + "amber/tip3p_standard.xml", # TIP3P and recommended monovalent ion parameters + "amber/tip3p_HFE_multivalent.xml", # for divalent ions + "amber/phosaa10.xml", # Handles THE TPO + ] + """List of force field paths for all components except :class:`SmallMoleculeComponent` """ + + small_molecule_forcefield: str = "openff-2.0.0" # other default ideas 'openff-2.0.0', 'gaff-2.11', 'espaloma-0.2.0' + """Name of the force field to be used for :class:`SmallMoleculeComponent` """ + + @validator('constraints') + def constraint_check(cls, v): + allowed = {'hbonds', 'hangles', 'allbonds'} + + if not (v is None or v.lower() in allowed): + raise ValueError(f"Bad constraints value, use one of {allowed}") + + return v + + + class Settings(SettingsBaseModel): + """ + Container for all settings needed by a protocol + + This represents the minimal surface that all settings objects will have. + + Protocols can subclass this to extend this to cater for their additional settings. + """ + forcefield_settings: BaseForceFieldSettings + thermo_settings: ThermoSettings + + @classmethod + def get_defaults(cls): + return cls( + forcefield_settings=OpenMMSystemGeneratorFFSettings(), + thermo_settings=ThermoSettings(temperature=300 * unit.kelvin), + ) +else: + from gufe.settings import Settings, SettingsBaseModel, ThermoSettings, OpenMMSystemGeneratorFFSettings # type: ignore + + +class SystemSettings(SettingsBaseModel): + """Settings describing the simulation system settings.""" + + class Config: + arbitrary_types_allowed = True + + nonbonded_method = 'PME' + """ + Method for treating nonbonded interactions, currently only PME and + NoCutoff are allowed. Default PME. + """ + nonbonded_cutoff = 1.0 * unit.nanometer + """ + Cutoff value for short range nonbonded interactions. + Default 1.0 * unit.nanometer. + """ + + @validator('nonbonded_method') + def allowed_nonbonded(cls, v): + if v.lower() not in ['pme', 'nocutoff']: + errmsg = ("Only PME and NoCutoff are allowed nonbonded_methods") + raise ValueError(errmsg) + return v + + @validator('nonbonded_cutoff') + def is_positive_distance(cls, v): + # these are time units, not simulation steps + if not v.is_compatible_with(unit.nanometer): + raise ValueError("nonbonded_cutoff must be in distance units " + "(i.e. nanometers)") + if v < 0: + errmsg = "nonbonded_cutoff must be a positive value" + raise ValueError(errmsg) + return v + + +class SolvationSettings(SettingsBaseModel): + """Settings for solvating the system + + Note + ---- + No solvation will happen if a SolventComponent is not passed. + + """ + class Config: + arbitrary_types_allowed = True + + solvent_model = 'tip3p' + """ + Force field water model to use. + Allowed values are; `tip3p`, `spce`, `tip4pew`, and `tip5p`. + """ + + solvent_padding = 1.2 * unit.nanometer + """Minimum distance from any solute atoms to the solvent box edge.""" + + @validator('solvent_model') + def allowed_solvent(cls, v): + allowed_models = ['tip3p', 'spce', 'tip4pew', 'tip5p'] + if v.lower() not in allowed_models: + errmsg = ( + f"Only {allowed_models} are allowed solvent_model values" + ) + raise ValueError(errmsg) + return v + + @validator('solvent_padding') + def is_positive_distance(cls, v): + # these are time units, not simulation steps + if not v.is_compatible_with(unit.nanometer): + raise ValueError("solvent_padding must be in distance units " + "(i.e. nanometers)") + if v < 0: + errmsg = "solvent_padding must be a positive value" + raise ValueError(errmsg) + return v + + +class AlchemicalSamplerSettings(SettingsBaseModel): + """Settings for the Equilibrium Alchemical sampler, currently supporting + either MultistateSampler, SAMSSampler or ReplicaExchangeSampler. + + """ + + """ + TODO + ---- + * It'd be great if we could pass in the sampler object rather than using + strings to define which one we want. + * Make n_replicas optional such that: If `None` or greater than the number + of lambda windows set in :class:`AlchemicalSettings`, this will default + to the number of lambda windows. If less than the number of lambda + windows, the replica lambda states will be picked at equidistant + intervals along the lambda schedule. + """ + class Config: + arbitrary_types_allowed = True + + sampler_method = "repex" + """ + Alchemical sampling method, must be one of; + `repex` (Hamiltonian Replica Exchange), + `sams` (Self-Adjusted Mixture Sampling), + or `independent` (independently sampled lambda windows). + Default `repex`. + """ + online_analysis_interval: Optional[int] = 250 + """ + MCMC steps (i.e. ``IntegratorSettings.n_steps``) interval at which + to perform an analysis of the free energies. + + At each interval, real time analysis data (e.g. current free energy + estimate and timing data) will be written to a yaml file named + ``_real_time_analysis.yaml``. The + current error in the estimate will also be assed and if it drops + below ``AlchemicalSamplerSettings.online_analysis_target_error`` + the simulation will be terminated. + + If ``None``, no real time analysis will be performed and the yaml + file will not be written. + + Must be a multiple of ``SimulationSettings.checkpoint_interval`` + + Default `250`. + """ + online_analysis_target_error = 0.0 * unit.boltzmann_constant * unit.kelvin + """ + Target error for the online analysis measured in kT. Once the free energy + is at or below this value, the simulation will be considered complete. A + suggested value of 0.2 * `unit.boltzmann_constant` * `unit.kelvin` has + shown to be effective in both hydration and binding free energy benchmarks. + Default 0.0 * `unit.boltzmann_constant` * `unit.kelvin`, i.e. no early + termination will occur. + """ + online_analysis_minimum_iterations = 500 + """ + Number of iterations which must pass before online analysis is + carried out. Default 500. + """ + n_repeats: int = 3 + """ + Number of independent repeats to run. Default 3 + """ + flatness_criteria = 'logZ-flatness' + """ + SAMS only. Method for assessing when to switch to asymptomatically + optimal scheme. + One of ['logZ-flatness', 'minimum-visits', 'histogram-flatness']. + Default 'logZ-flatness'. + """ + gamma0 = 1.0 + """SAMS only. Initial weight adaptation rate. Default 1.0.""" + n_replicas = 11 + """Number of replicas to use. Default 11.""" + + @validator('flatness_criteria') + def supported_flatness(cls, v): + supported = [ + 'logz-flatness', 'minimum-visits', 'histogram-flatness' + ] + if v.lower() not in supported: + errmsg = ("Only the following flatness_criteria are " + f"supported: {supported}") + raise ValueError(errmsg) + return v + + @validator('sampler_method') + def supported_sampler(cls, v): + supported = ['repex', 'sams', 'independent'] + if v.lower() not in supported: + errmsg = ("Only the following sampler_method values are " + f"supported: {supported}") + raise ValueError(errmsg) + return v + + @validator('n_repeats', 'n_replicas') + def must_be_positive(cls, v): + if v <= 0: + errmsg = "n_repeats and n_replicas must be positive values" + raise ValueError(errmsg) + return v + + @validator('online_analysis_target_error', 'n_repeats', + 'online_analysis_minimum_iterations', 'gamma0', 'n_replicas') + def must_be_zero_or_positive(cls, v): + if v < 0: + errmsg = ("Online analysis target error, minimum iteration " + "and SAMS gamm0 must be 0 or positive values.") + raise ValueError(errmsg) + return v + +class RepeatSettings(SettingsBaseModel): + """Settings for how many independent MD runs to perform.""" + + n_repeats: int = 1 + """ + Number of independent repeats to run. Default 3 + """ + +class OpenMMEngineSettings(SettingsBaseModel): + """OpenMM MD engine settings""" + + + """ + TODO + ---- + * In the future make precision and deterministic forces user defined too. + """ + + compute_platform: Optional[str] = None + """ + OpenMM compute platform to perform MD integration with. If None, will + choose fastest available platform. Default None. + """ + + +class IntegratorSettings(SettingsBaseModel): + """Settings for the LangevinSplittingDynamicsMove integrator""" + + class Config: + arbitrary_types_allowed = True + + timestep = 4 * unit.femtosecond + """Size of the simulation timestep. Default 4 * unit.femtosecond.""" + collision_rate = 1.0 / unit.picosecond + """Collision frequency. Default 1.0 / unit.pisecond.""" + n_steps = 250 * unit.timestep + """ + Number of integration timesteps between each time the MCMC move + is applied. Default 250 * unit.timestep. + """ + reassign_velocities = False + """ + If True, velocities are reassigned from the Maxwell-Boltzmann + distribution at the beginning of move. Default False. + """ + n_restart_attempts = 20 + """ + Number of attempts to restart from Context if there are NaNs in the + energies after integration. Default 20. + """ + constraint_tolerance = 1e-06 + """Tolerance for the constraint solver. Default 1e-6.""" + barostat_frequency = 25 * unit.timestep + """ + Frequency at which volume scaling changes should be attempted. + Default 25 * unit.timestep. + """ + + @validator('collision_rate', 'n_restart_attempts') + def must_be_positive_or_zero(cls, v): + if v < 0: + errmsg = ("collision_rate, and n_restart_attempts must be " + "zero or positive values") + raise ValueError(errmsg) + return v + + @validator('timestep', 'n_steps', 'constraint_tolerance') + def must_be_positive(cls, v): + if v <= 0: + errmsg = ("timestep, n_steps, constraint_tolerance " + "must be positive values") + raise ValueError(errmsg) + return v + + @validator('timestep') + def is_time(cls, v): + # these are time units, not simulation steps + if not v.is_compatible_with(unit.picosecond): + raise ValueError("timestep must be in time units " + "(i.e. picoseconds)") + return v + + @validator('collision_rate') + def must_be_inverse_time(cls, v): + if not v.is_compatible_with(1 / unit.picosecond): + raise ValueError("collision_rate must be in inverse time " + "(i.e. 1/picoseconds)") + return v + + +class SimulationSettings(SettingsBaseModel): + """ + Settings for simulation control, including lengths, + writing to disk, etc... + """ + class Config: + arbitrary_types_allowed = True + + minimization_steps = 5000 + """Number of minimization steps to perform. Default 5000.""" + equilibration_length: unit.Quantity + """ + Length of the equilibration phase in units of time. The total number of + steps from this equilibration length + (i.e. ``equilibration_length`` / :class:`IntegratorSettings.timestep`) + must be a multiple of the value defined for + :class:`IntegratorSettings.n_steps`. + """ + production_length: unit.Quantity + """ + Length of the production phase in units of time. The total number of + steps from this production length (i.e. + ``production_length`` / :class:`IntegratorSettings.timestep`) must be + a multiple of the value defined for :class:`IntegratorSettings.nsteps`. + """ + + # reporter settings + output_filename = 'simulation.nc' + """Path to the storage file for analysis. Default 'simulation.nc'.""" + output_indices = 'not water' + """ + Selection string for which part of the system to write coordinates for. + Default 'not water'. + """ + checkpoint_interval = 250 * unit.timestep + """ + Frequency to write the checkpoint file. Default 250 * unit.timestep. + """ + checkpoint_storage = 'checkpoint.nc' + """ + Separate filename for the checkpoint file. Note, this should + not be a full path, just a filename. Default 'checkpoint.nc'. + """ + forcefield_cache: Optional[str] = 'db.json' + """ + Filename for caching small molecule residue templates so they can be + later reused. + """ + + @validator('equilibration_length', 'production_length') + def is_time(cls, v): + # these are time units, not simulation steps + if not v.is_compatible_with(unit.picosecond): + raise ValueError("Durations must be in time units") + return v + + @validator('minimization_steps', 'equilibration_length', + 'production_length', 'checkpoint_interval') + def must_be_positive(cls, v): + if v <= 0: + errmsg = ("Minimization steps, MD lengths, and checkpoint " + "intervals must be positive") + raise ValueError(errmsg) + return v \ No newline at end of file From e64fb75c3dc52f847a04d9f82b637eef67d68d50 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 1 Aug 2023 13:56:02 +0200 Subject: [PATCH 02/66] Add tests, fixing some errors in methods --- .../protocols/openmm_md/plain_md_methods.py | 115 ++---- .../test_openmm_plain_md_protocols.py | 384 ++++++++++++++++++ 2 files changed, 422 insertions(+), 77 deletions(-) create mode 100644 openfe/tests/protocols/test_openmm_plain_md_protocols.py diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 12e96991f..469ab7383 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -10,7 +10,6 @@ """ from __future__ import annotations -import os import logging from collections import defaultdict @@ -27,6 +26,7 @@ import openmmtools import uuid import time +from mdtraj.reporters import NetCDFReporter from gufe import ( settings, ChemicalSystem, SmallMoleculeComponent, @@ -124,11 +124,13 @@ def _create( system_validation.validate_protein(stateA) # actually create and return Units - lig_name = stateA.components['ligand'].name + solvent_comp, protein_comp, small_mols = \ + system_validation.get_components(stateA) + lig_name = small_mols[0].name # our DAG has no dependencies, so just list units n_repeats = self.settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( - stateA=stateA, stateB=stateB, ligandmapping=None, + stateA=stateA, stateB=stateB, settings=self.settings, generation=0, repeat_id=int(uuid.uuid4()), name=f'{lig_name} repeat {i} generation 0') @@ -179,8 +181,6 @@ def __init__(self, *, stateA, stateB : ChemicalSystem the two ligand SmallMoleculeComponents to transform between. The transformation will go from ligandA to ligandB. - ligandmapping : LigandAtomMapping - the mapping of atoms between the two ligand components settings : settings.Settings the settings for the Method. This can be constructed using the get_default_settings classmethod to give a starting point that @@ -201,7 +201,6 @@ def __init__(self, *, name=name, stateA=stateA, stateB=stateB, - ligandmapping=None, settings=settings, repeat_id=repeat_id, generation=generation @@ -295,6 +294,7 @@ def run(self, *, dry=False, verbose=True, protocol_settings.simulation_settings timestep = protocol_settings.integrator_settings.timestep mc_steps = protocol_settings.integrator_settings.n_steps.m + integrator_settings = protocol_settings.integrator_settings # is the timestep good for the mass? settings_validation.validate_timestep( @@ -351,54 +351,39 @@ def run(self, *, dry=False, verbose=True, # e. create the stateA System stateA_system = system_generator.create_system( - stateA_modeller.topology, + stateA_topology, molecules=[s.to_openff() for s in small_mols], ) - # 9. Create the reporter - # Get the sub selection of the system to print coords for - # selection_indices = stateA_topology.select( - # sim_settings.output_indices - # ) - - # a. Create the reporter - reporter_dcd = openmm.app.DCDReporter( - shared_basepath / sim_settings.output_filename, - sim_settings.checkpoint_interval.m, - enforcePeriodicBox=False) - # reporter = multistate.MultiStateReporter( - # storage=shared_basepath / sim_settings.output_filename, - # analysis_particle_indices=selection_indices, - # checkpoint_interval=sim_settings.checkpoint_interval.m, - # checkpoint_storage=sim_settings.checkpoint_storage, - # ) - # 10. Get platform platform = compute.get_openmm_platform( protocol_settings.engine_settings.compute_platform ) # 11. Set the integrator - # a. get integrator settings - integrator_settings = protocol_settings.integrator_settings + # # a. get integrator settings + # integrator_settings = protocol_settings.integrator_settings + # + # # b. create langevin integrator + # integrator = openmmtools.mcmc.LangevinDynamicsMove( + # timestep=to_openmm(integrator_settings.timestep), + # collision_rate=to_openmm(integrator_settings.collision_rate), + # n_steps=integrator_settings.n_steps.m, + # reassign_velocities=integrator_settings.reassign_velocities, + # n_restart_attempts=integrator_settings.n_restart_attempts, + # constraint_tolerance=integrator_settings.constraint_tolerance, + # ) + integrator = openmm.LangevinIntegrator( + to_openmm(thermo_settings.temperature), + integrator_settings.n_steps.m, + to_openmm(integrator_settings.timestep), + ) - # Validate settings - # Virtual sites sanity check - ensure we restart velocities when - # there are virtual sites in the system - if stateA_system.has_virtual_sites: - if not integrator_settings.reassign_velocities: - errmsg = ("Simulations with virtual sites without velocity " - "reassignments are unstable in openmmtools") - raise ValueError(errmsg) - - # b. create langevin integrator - integrator = openmmtools.mcmc.LangevinDynamicsMove( - timestep=to_openmm(integrator_settings.timestep), - collision_rate=to_openmm(integrator_settings.collision_rate), - n_steps=integrator_settings.n_steps.m, - reassign_velocities=integrator_settings.reassign_velocities, - n_restart_attempts=integrator_settings.n_restart_attempts, - constraint_tolerance=integrator_settings.constraint_tolerance, + simulation = openmm.app.Simulation( + stateA_modeller.topology, + stateA_system, + integrator, + platform=platform ) try: @@ -416,12 +401,6 @@ def run(self, *, dry=False, verbose=True, if not dry: # pragma: no-cover - simulation = openmm.app.Simulation( - stateA_topology, - stateA_system, - integrator, - platform=platform - ) simulation.context.setPositions(stateA_positions) # minimize @@ -447,7 +426,8 @@ def run(self, *, dry=False, verbose=True, if verbose: logger.info("running production phase") - simulation.reporters.append(reporter_dcd) + simulation.reporters.append(NetCDFReporter( + shared_basepath / sim_settings.output_filename, 1000)) t0 = time.time() simulation.step(prod_steps) @@ -462,34 +442,15 @@ def run(self, *, dry=False, verbose=True, # omm_unit.kilocalories_per_mole) # est = ensure_quantity(est, 'openff') # ana.clear() # clean up cached values - nc = shared_basepath / sim_settings.output_filename chk = shared_basepath / sim_settings.checkpoint_storage - else: - # clean up the reporter file - fns = [shared_basepath / sim_settings.output_filename, - shared_basepath / sim_settings.checkpoint_storage] - for fn in fns: - os.remove(fn) + # else: + # # clean up the reporter file + # fns = [shared_basepath / sim_settings.output_filename, + # shared_basepath / sim_settings.checkpoint_storage] + # for fn in fns: + # os.remove(fn) finally: - # close reporter when you're done, prevent file handle clashes - reporter_dcd.close() - del reporter_dcd - - # # clear GPU contexts - # # replace with above - # for context in list(energy_context_cache._lru._data.keys()): - # del energy_context_cache._lru._data[context] - # for context in list(sampler_context_cache._lru._data.keys()): - # del sampler_context_cache._lru._data[context] - # # cautiously clear out the global context cache too - # for context in list( - # openmmtools.cache.global_context_cache._lru._data - # .keys()): - # del openmmtools.cache.global_context_cache._lru._data[ - # context] - # - # del sampler_context_cache, energy_context_cache if not dry: del integrator, simulation @@ -500,7 +461,7 @@ def run(self, *, dry=False, verbose=True, 'last_checkpoint': chk, } else: - return {'debug': {'simulation': simulation}} + return {'debug': {'sampler': stateA_system}} def _execute( self, ctx: gufe.Context, **kwargs, diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py new file mode 100644 index 000000000..c79ac1934 --- /dev/null +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -0,0 +1,384 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import os +from io import StringIO +import numpy as np +import gufe +from gufe.tests.test_tokenization import GufeTokenizableTestsMixin +import pytest +from unittest import mock +from openff.units import unit +from importlib import resources +import xml.etree.ElementTree as ET +from openff.units.openmm import from_openmm, to_openmm +from openmm import app, XmlSerializer, MonteCarloBarostat +from openmm import unit as omm_unit +from openmmtools.states import ThermodynamicState +import pathlib +from rdkit import Chem +from rdkit.Geometry import Point3D +import mdtraj as mdt + +import openfe +from openfe import setup +from openfe.protocols import openmm_md +from openfe.protocols.openmm_utils import system_creation +from openmmforcefields.generators import SMIRNOFFTemplateGenerator +from openff.units.openmm import ensure_quantity +from openfe.protocols.openmm_md.plain_md_methods import ( + PlainMDProtocol, PlainMDProtocolUnit, PlainMDProtocolResult, +) +from openfe.protocols.openmm_md.plain_md_settings import PlainMDProtocolSettings + + +def test_create_default_settings(): + settings = PlainMDProtocol.default_settings() + + assert settings + + +def test_create_default_protocol(): + # this is roughly how it should be created + protocol = PlainMDProtocol( + settings=PlainMDProtocol.default_settings(), + ) + + assert protocol + + +def test_serialize_protocol(): + protocol = PlainMDProtocol( + settings=PlainMDProtocol.default_settings(), + ) + + ser = protocol.to_dict() + + ret = PlainMDProtocol.from_dict(ser) + + assert protocol == ret + + +def test_create_independent_repeat_ids(benzene_system): + # if we create two dags each with 3 repeats, they should give 6 repeat_ids + # this allows multiple DAGs in flight for one Transformation that don't clash on gather + settings = PlainMDProtocol.default_settings() + # Default protocol is 1 repeat, change to 3 repeats + settings.repeat_settings.n_repeats = 3 + protocol = PlainMDProtocol( + settings=settings, + ) + dag1 = protocol.create( + stateA=benzene_system, + stateB=benzene_system, + mapping=None, + ) + dag2 = protocol.create( + stateA=benzene_system, + stateB=benzene_system, + mapping=None, + ) + + repeat_ids = set() + u: PlainMDProtocolUnit + for u in dag1.protocol_units: + repeat_ids.add(u.inputs['repeat_id']) + for u in dag2.protocol_units: + repeat_ids.add(u.inputs['repeat_id']) + + assert len(repeat_ids) == 6 + + +def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): + + vac_settings = PlainMDProtocol.default_settings() + vac_settings.system_settings.nonbonded_method = 'nocutoff' + + protocol = PlainMDProtocol( + settings=vac_settings, + ) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sim = dag_unit.run(dry=True)['debug']['sampler'] + print(sim) + assert not ThermodynamicState(sim, temperature= + to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic + assert ThermodynamicState(sim, temperature= + to_openmm(protocol.settings.thermo_settings.temperature)).barostat is None + + # Check MDTtraj Topologies + # 16 atoms: + # 11 common atoms, 1 extra hydrogen in benzene, 4 extra in toluene + # 12 bonds in benzene + 4 extra tolunee bonds + # assert len(list(.atoms)) == 16 + # assert len(list(htf.omm_hybrid_topology.atoms())) == 16 + # assert len(list(htf.hybrid_topology.bonds)) == 16 + # assert len(list(htf.omm_hybrid_topology.bonds())) == 16 + + # # smoke test - can convert back the mdtraj topology + # ret_top = mdt.Topology.to_openmm(htf.hybrid_topology) + # assert len(list(ret_top.atoms())) == 16 + # assert len(list(ret_top.bonds())) == 16 + + +def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir): + vac_settings = PlainMDProtocol.default_settings() + vac_settings.system_settings.nonbonded_method = 'nocutoff' + vac_settings.forcefield_settings.small_molecule_forcefield = 'gaff-2.11' + + protocol = PlainMDProtocol( + settings=vac_settings, + ) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = unit.run(dry=True)['debug']['sampler'] + + +def test_dry_many_molecules_solvent( + benzene_many_solv_system, tmpdir +): + """ + A basic test flushing "will it work if you pass multiple molecules" + """ + settings = PlainMDProtocol.default_settings() + + protocol = PlainMDProtocol( + settings=settings, + ) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=benzene_many_solv_system, + stateB=benzene_many_solv_system, + mapping=None, + ) + unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = unit.run(dry=True)['debug']['sampler'] + + +BENZ = """\ +benzene + PyMOL2.5 3D 0 + + 12 12 0 0 0 0 0 0 0 0999 V2000 + 1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7022 1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 2.5079 -0.0000 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1.2540 2.1720 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2540 2.1720 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -2.5079 -0.0000 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2540 -2.1719 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1.2540 -2.1720 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 2 0 0 0 0 + 1 6 1 0 0 0 0 + 1 7 1 0 0 0 0 + 2 3 1 0 0 0 0 + 2 8 1 0 0 0 0 + 3 4 2 0 0 0 0 + 3 9 1 0 0 0 0 + 4 5 1 0 0 0 0 + 4 10 1 0 0 0 0 + 5 6 2 0 0 0 0 + 5 11 1 0 0 0 0 + 6 12 1 0 0 0 0 +M END +$$$$ +""" + + +PYRIDINE = """\ +pyridine + PyMOL2.5 3D 0 + + 11 11 0 0 0 0 0 0 0 0999 V2000 + 1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 2.4940 -0.0325 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1.2473 -2.1604 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2473 -2.1604 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -2.4945 -0.0000 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2753 2.1437 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7525 1.3034 0.0000 N 0 0 0 0 0 0 0 0 0 0 0 0 + 1 5 1 0 0 0 0 + 1 6 1 0 0 0 0 + 1 11 2 0 0 0 0 + 2 3 2 0 0 0 0 + 2 10 1 0 0 0 0 + 3 4 1 0 0 0 0 + 3 9 1 0 0 0 0 + 4 5 2 0 0 0 0 + 4 8 1 0 0 0 0 + 5 7 1 0 0 0 0 + 2 11 1 0 0 0 0 +M END +$$$$ +""" + +def test_dry_run_ligand_tip4p(benzene_system, tmpdir): + """ + Test that we can create a system with virtual sites in the + environment (waters) + """ + settings = PlainMDProtocol.default_settings() + settings.forcefield_settings.forcefields = [ + "amber/ff14SB.xml", # ff14SB protein force field + "amber/tip4pew_standard.xml", # FF we are testsing with the fun VS + "amber/phosaa10.xml", # Handles THE TPO + ] + settings.solvation_settings.solvent_padding = 1.0 * unit.nanometer + settings.system_settings.nonbonded_cutoff = 0.9 * unit.nanometer + settings.solvation_settings.solvent_model = 'tip4pew' + settings.integrator_settings.reassign_velocities = True + + protocol = PlainMDProtocol( + settings=settings, + ) + dag = protocol.create( + stateA=benzene_system, + stateB=benzene_system, + mapping=None, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = dag_unit.run(dry=True)['debug']['sampler'] + assert sampler + + +# @pytest.mark.slow +# def test_dry_run_complex(benzene_complex_system, tmpdir): +# # this will be very time consuming +# settings = PlainMDProtocol.default_settings() +# +# protocol = PlainMDProtocol( +# settings=settings, +# ) +# dag = protocol.create( +# stateA=benzene_complex_system, +# stateB=benzene_complex_system, +# mapping=None, +# ) +# dag_unit = list(dag.protocol_units)[0] +# +# with tmpdir.as_cwd(): +# sim = dag_unit.run(dry=True)['debug']['sampler'] +# assert not ThermodynamicState(sim, temperature= +# to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic +# assert isinstance(ThermodynamicState(sim, temperature= +# to_openmm(protocol.settings.thermo_settings.temperature)).barostat, +# MonteCarloBarostat) +# assert ThermodynamicState(sim, temperature= +# to_openmm(protocol.settings.thermo_settings.temperature)).pressure == 1 * omm_unit.bar + + +def test_hightimestep(benzene_vacuum_system, tmpdir): + settings = PlainMDProtocol.default_settings() + settings.forcefield_settings.hydrogen_mass = 1.0 + settings.system_settings.nonbonded_method = 'nocutoff' + + p = PlainMDProtocol( + settings=settings, + ) + + dag = p.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + dag_unit = list(dag.protocol_units)[0] + + errmsg = "too large for hydrogen mass" + with tmpdir.as_cwd(): + with pytest.raises(ValueError, match=errmsg): + dag_unit.run(dry=True) + + +def test_vaccuum_PME_error(benzene_vacuum_system): + + p = PlainMDProtocol( + settings=PlainMDProtocol.default_settings(), + ) + errmsg = "PME cannot be used for vacuum transform" + with pytest.raises(ValueError, match=errmsg): + _ = p.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + + +@pytest.fixture +def solvent_protocol_dag(benzene_system): + settings = PlainMDProtocol.default_settings() + settings.repeat_settings.n_repeats = 3 + protocol = PlainMDProtocol( + settings=settings, + ) + + return protocol.create( + stateA=benzene_system, stateB=benzene_system, + mapping=None, + ) + + +def test_unit_tagging(solvent_protocol_dag, tmpdir): + # test that executing the Units includes correct generation and repeat info + dag_units = solvent_protocol_dag.protocol_units + with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', + return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}): + results = [] + for u in dag_units: + ret = u.execute(context=gufe.Context(tmpdir, tmpdir)) + results.append(ret) + + repeats = set() + for ret in results: + assert isinstance(ret, gufe.ProtocolUnitResult) + assert ret.outputs['generation'] == 0 + repeats.add(ret.outputs['repeat_id']) + # repeats are random ints, so check we got 3 individual numbers + assert len(repeats) == 3 + + +# def test_gather(solvent_protocol_dag, tmpdir): +# # check .gather behaves as expected +# with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', +# return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}): +# dagres = gufe.protocols.execute_DAG(solvent_protocol_dag, +# shared_basedir=tmpdir, +# scratch_basedir=tmpdir, +# keep_shared=True) +# +# settings = PlainMDProtocol.default_settings() +# settings.repeat_settings.n_repeats = 3 +# prot = PlainMDProtocol( +# settings=settings +# ) +# +# res = prot.gather([dagres]) +# +# assert isinstance(res, PlainMDProtocolResult) From 86a60c9480ff66042aebba128c00a6d3473e2193 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 2 Aug 2023 13:05:46 +0200 Subject: [PATCH 03/66] Add gathertest --- .../test_openmm_plain_md_protocols.py | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index c79ac1934..048f6f7e2 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -364,21 +364,21 @@ def test_unit_tagging(solvent_protocol_dag, tmpdir): assert len(repeats) == 3 -# def test_gather(solvent_protocol_dag, tmpdir): -# # check .gather behaves as expected -# with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', -# return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}): -# dagres = gufe.protocols.execute_DAG(solvent_protocol_dag, -# shared_basedir=tmpdir, -# scratch_basedir=tmpdir, -# keep_shared=True) -# -# settings = PlainMDProtocol.default_settings() -# settings.repeat_settings.n_repeats = 3 -# prot = PlainMDProtocol( -# settings=settings -# ) -# -# res = prot.gather([dagres]) -# -# assert isinstance(res, PlainMDProtocolResult) +def test_gather(solvent_protocol_dag, tmpdir): + # check .gather behaves as expected + with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', + return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}): + dagres = gufe.protocols.execute_DAG(solvent_protocol_dag, + shared_basedir=tmpdir, + scratch_basedir=tmpdir, + keep_shared=True) + + settings = PlainMDProtocol.default_settings() + settings.repeat_settings.n_repeats = 3 + prot = PlainMDProtocol( + settings=settings + ) + + res = prot.gather([dagres]) + + assert isinstance(res, PlainMDProtocolResult) From b9c841323fac72d0259141d08f06307f419b96d3 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 2 Aug 2023 15:25:06 +0200 Subject: [PATCH 04/66] Add StateDataReporter --- openfe/protocols/openmm_md/plain_md_methods.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 469ab7383..e349bdcbc 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import sys from collections import defaultdict import gufe @@ -427,8 +428,15 @@ def run(self, *, dry=False, verbose=True, logger.info("running production phase") simulation.reporters.append(NetCDFReporter( - shared_basepath / sim_settings.output_filename, 1000)) - + shared_basepath / sim_settings.output_filename, + sim_settings.checkpoint_interval.m)) + simulation.reporters.append(openmm.app.StateDataReporter( + sys.stdout, + sim_settings.checkpoint_interval.m, + step=True, + potentialEnergy=True, + temperature=True + )) t0 = time.time() simulation.step(prod_steps) t1 = time.time() From f33e21d9464d97a42eb1d462705128947d86a350 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 3 Aug 2023 10:19:52 +0200 Subject: [PATCH 05/66] Cleanup functions --- .../protocols/openmm_md/plain_md_methods.py | 53 +++---------------- .../test_openmm_plain_md_protocols.py | 12 ++--- 2 files changed, 13 insertions(+), 52 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index e349bdcbc..5b5781b5b 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -362,19 +362,7 @@ def run(self, *, dry=False, verbose=True, ) # 11. Set the integrator - # # a. get integrator settings - # integrator_settings = protocol_settings.integrator_settings - # - # # b. create langevin integrator - # integrator = openmmtools.mcmc.LangevinDynamicsMove( - # timestep=to_openmm(integrator_settings.timestep), - # collision_rate=to_openmm(integrator_settings.collision_rate), - # n_steps=integrator_settings.n_steps.m, - # reassign_velocities=integrator_settings.reassign_velocities, - # n_restart_attempts=integrator_settings.n_restart_attempts, - # constraint_tolerance=integrator_settings.constraint_tolerance, - # ) - integrator = openmm.LangevinIntegrator( + integrator = openmm.LangevinMiddleIntegrator( to_openmm(thermo_settings.temperature), integrator_settings.n_steps.m, to_openmm(integrator_settings.timestep), @@ -388,17 +376,6 @@ def run(self, *, dry=False, verbose=True, ) try: - # # Create context caches (energy + sampler) - # energy_context_cache = openmmtools.cache.ContextCache( - # capacity=None, time_to_live=None, platform=platform, - # ) - # - # sampler_context_cache = openmmtools.cache.ContextCache( - # capacity=None, time_to_live=None, platform=platform, - # ) - # - # sampler.energy_context_cache = energy_context_cache - # sampler.sampler_context_cache = sampler_context_cache if not dry: # pragma: no-cover @@ -408,9 +385,9 @@ def run(self, *, dry=False, verbose=True, if verbose: logger.info("minimizing systems") - simulation.minimizeEnergy() - # sampler.minimize( - # max_iterations=sim_settings.minimization_steps) + simulation.minimizeEnergy( + maxIterations=sim_settings.minimization_steps + ) # equilibrate if verbose: @@ -420,9 +397,6 @@ def run(self, *, dry=False, verbose=True, thermo_settings.temperature * openmm.unit.kelvin) simulation.step(equil_steps) - # sampler.equilibrate( - # int(equil_steps / mc_steps)) # type: ignore - # production if verbose: logger.info("running production phase") @@ -441,23 +415,10 @@ def run(self, *, dry=False, verbose=True, simulation.step(prod_steps) t1 = time.time() logger.info(f"Completed simulation in {t1 - t0} seconds") - # sampler.extend(int(prod_steps / mc_steps)) # type: ignore - - # # calculate estimate of results from this individual unit - # ana = multistate.MultiStateSamplerAnalyzer(reporter) - # est, _ = ana.get_free_energy() - # est = (est[0, -1] * ana.kT).in_units_of( - # omm_unit.kilocalories_per_mole) - # est = ensure_quantity(est, 'openff') - # ana.clear() # clean up cached values + nc = shared_basepath / sim_settings.output_filename chk = shared_basepath / sim_settings.checkpoint_storage - # else: - # # clean up the reporter file - # fns = [shared_basepath / sim_settings.output_filename, - # shared_basepath / sim_settings.checkpoint_storage] - # for fn in fns: - # os.remove(fn) + finally: if not dry: @@ -469,7 +430,7 @@ def run(self, *, dry=False, verbose=True, 'last_checkpoint': chk, } else: - return {'debug': {'sampler': stateA_system}} + return {'debug': {'system': stateA_system}} def _execute( self, ctx: gufe.Context, **kwargs, diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index 048f6f7e2..fc6aea4b2 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -106,7 +106,7 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): dag_unit = list(dag.protocol_units)[0] with tmpdir.as_cwd(): - sim = dag_unit.run(dry=True)['debug']['sampler'] + sim = dag_unit.run(dry=True)['debug']['system'] print(sim) assert not ThermodynamicState(sim, temperature= to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic @@ -146,7 +146,7 @@ def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir): unit = list(dag.protocol_units)[0] with tmpdir.as_cwd(): - sampler = unit.run(dry=True)['debug']['sampler'] + system = unit.run(dry=True)['debug']['system'] def test_dry_many_molecules_solvent( @@ -170,7 +170,7 @@ def test_dry_many_molecules_solvent( unit = list(dag.protocol_units)[0] with tmpdir.as_cwd(): - sampler = unit.run(dry=True)['debug']['sampler'] + system = unit.run(dry=True)['debug']['system'] BENZ = """\ @@ -265,8 +265,8 @@ def test_dry_run_ligand_tip4p(benzene_system, tmpdir): dag_unit = list(dag.protocol_units)[0] with tmpdir.as_cwd(): - sampler = dag_unit.run(dry=True)['debug']['sampler'] - assert sampler + system = dag_unit.run(dry=True)['debug']['system'] + assert system # @pytest.mark.slow @@ -285,7 +285,7 @@ def test_dry_run_ligand_tip4p(benzene_system, tmpdir): # dag_unit = list(dag.protocol_units)[0] # # with tmpdir.as_cwd(): -# sim = dag_unit.run(dry=True)['debug']['sampler'] +# sim = dag_unit.run(dry=True)['debug']['system'] # assert not ThermodynamicState(sim, temperature= # to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic # assert isinstance(ThermodynamicState(sim, temperature= From 72ee3648c3ce82ef596257fc30f5084843a17f00 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 3 Aug 2023 14:52:36 +0200 Subject: [PATCH 06/66] fix integrator units --- openfe/protocols/openmm_md/plain_md_methods.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 5b5781b5b..3ed79f2f8 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -39,6 +39,7 @@ IntegratorSettings, SimulationSettings, RepeatSettings ) + from openfe.protocols.openmm_rfe._rfe_utils import compute from openfe.protocols.openmm_utils import ( system_validation, settings_validation, system_creation @@ -364,7 +365,7 @@ def run(self, *, dry=False, verbose=True, # 11. Set the integrator integrator = openmm.LangevinMiddleIntegrator( to_openmm(thermo_settings.temperature), - integrator_settings.n_steps.m, + to_openmm(integrator_settings.collision_rate), to_openmm(integrator_settings.timestep), ) @@ -388,13 +389,20 @@ def run(self, *, dry=False, verbose=True, simulation.minimizeEnergy( maxIterations=sim_settings.minimization_steps ) - + min_pdb_file = "minimized.pdb" + positions = simulation.context.getState( + getPositions=True, enforcePeriodicBox=False + ).getPositions() + with open(shared_basepath / min_pdb_file, "w") as f: + openmm.app.PDBFile.writeFile( + simulation.topology, positions, file=f, keepIds=True + ) # equilibrate if verbose: logger.info("equilibrating systems") - - simulation.context.setVelocitiesToTemperature( - thermo_settings.temperature * openmm.unit.kelvin) + # simulation.context.setVelocitiesToTemperature( + # to_openmm(thermo_settings.temperature), 8) + simulation.context.setVelocitiesToTemperature(298.15 * openmm.unit.kelvin) simulation.step(equil_steps) # production From 1516300649894a4b0cf8dcb28171f53054b59f27 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 3 Aug 2023 14:59:31 +0200 Subject: [PATCH 07/66] fix setvelocity unit --- openfe/protocols/openmm_md/plain_md_methods.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 3ed79f2f8..0a4126eb4 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -400,9 +400,8 @@ def run(self, *, dry=False, verbose=True, # equilibrate if verbose: logger.info("equilibrating systems") - # simulation.context.setVelocitiesToTemperature( - # to_openmm(thermo_settings.temperature), 8) - simulation.context.setVelocitiesToTemperature(298.15 * openmm.unit.kelvin) + simulation.context.setVelocitiesToTemperature( + to_openmm(thermo_settings.temperature)) simulation.step(equil_steps) # production From b9f9b6d5d682b493533801cdb657673fc126ad58 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 3 Aug 2023 16:03:44 +0200 Subject: [PATCH 08/66] test reporter --- openfe/protocols/openmm_md/plain_md_methods.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 0a4126eb4..0a4bc5931 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -407,9 +407,12 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - + traj = "traj.nc" + # simulation.reporters.append(NetCDFReporter( + # shared_basepath / sim_settings.output_filename, + # sim_settings.checkpoint_interval.m)) simulation.reporters.append(NetCDFReporter( - shared_basepath / sim_settings.output_filename, + shared_basepath / traj, sim_settings.checkpoint_interval.m)) simulation.reporters.append(openmm.app.StateDataReporter( sys.stdout, From 033a9acef0020c5b2a740cdd72c3e69eef980297 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 3 Aug 2023 16:10:59 +0200 Subject: [PATCH 09/66] change reporter to xtc --- openfe/protocols/openmm_md/plain_md_methods.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 0a4bc5931..4f98fd319 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -27,7 +27,7 @@ import openmmtools import uuid import time -from mdtraj.reporters import NetCDFReporter +from mdtraj.reporters import NetCDFReporter, XTCReporter from gufe import ( settings, ChemicalSystem, SmallMoleculeComponent, @@ -407,13 +407,15 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - traj = "traj.nc" + traj = "traj.xtc" # simulation.reporters.append(NetCDFReporter( # shared_basepath / sim_settings.output_filename, # sim_settings.checkpoint_interval.m)) - simulation.reporters.append(NetCDFReporter( - shared_basepath / traj, - sim_settings.checkpoint_interval.m)) + # simulation.reporters.append(NetCDFReporter( + # shared_basepath / traj, + # sim_settings.checkpoint_interval.m)) + simulation.reporters.append(XTCReporter( + shared_basepath / traj, sim_settings.checkpoint_interval.m)) simulation.reporters.append(openmm.app.StateDataReporter( sys.stdout, sim_settings.checkpoint_interval.m, From 3901dfae17c326e4d6e2431bbb8e230a741d511a Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 3 Aug 2023 16:17:25 +0200 Subject: [PATCH 10/66] change reporter to dcd --- openfe/protocols/openmm_md/plain_md_methods.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 4f98fd319..984f7180c 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -407,15 +407,18 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - traj = "traj.xtc" + traj = "traj.dcd" # simulation.reporters.append(NetCDFReporter( # shared_basepath / sim_settings.output_filename, # sim_settings.checkpoint_interval.m)) # simulation.reporters.append(NetCDFReporter( # shared_basepath / traj, # sim_settings.checkpoint_interval.m)) - simulation.reporters.append(XTCReporter( - shared_basepath / traj, sim_settings.checkpoint_interval.m)) + # simulation.reporters.append(XTCReporter( + # shared_basepath / traj, sim_settings.checkpoint_interval.m)) + simulation.reporters.append(openmm.app.DCDReporter( + shared_basepath / traj, sim_settings.checkpoint_interval.m + )) simulation.reporters.append(openmm.app.StateDataReporter( sys.stdout, sim_settings.checkpoint_interval.m, From dbcacf3e9b5908da0ae7ca501ca479bb13ed8328 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 4 Aug 2023 09:15:01 +0200 Subject: [PATCH 11/66] clean up reporters --- openfe/protocols/openmm_md/plain_md_methods.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 984f7180c..82fc6d6d6 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -411,9 +411,6 @@ def run(self, *, dry=False, verbose=True, # simulation.reporters.append(NetCDFReporter( # shared_basepath / sim_settings.output_filename, # sim_settings.checkpoint_interval.m)) - # simulation.reporters.append(NetCDFReporter( - # shared_basepath / traj, - # sim_settings.checkpoint_interval.m)) # simulation.reporters.append(XTCReporter( # shared_basepath / traj, sim_settings.checkpoint_interval.m)) simulation.reporters.append(openmm.app.DCDReporter( From 7b4c04deed44e46b650e9a7f9bfd1dea1f9752b8 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 7 Aug 2023 09:41:21 +0200 Subject: [PATCH 12/66] Add checkpoint reporter --- openfe/protocols/openmm_md/plain_md_methods.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 82fc6d6d6..5b7608b76 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -407,14 +407,18 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - traj = "traj.dcd" - # simulation.reporters.append(NetCDFReporter( - # shared_basepath / sim_settings.output_filename, - # sim_settings.checkpoint_interval.m)) + traj = "prod.dcd" + simulation.reporters.append(NetCDFReporter( + shared_basepath / sim_settings.output_filename, + sim_settings.checkpoint_interval.m)) # simulation.reporters.append(XTCReporter( # shared_basepath / traj, sim_settings.checkpoint_interval.m)) - simulation.reporters.append(openmm.app.DCDReporter( - shared_basepath / traj, sim_settings.checkpoint_interval.m + # simulation.reporters.append(openmm.app.DCDReporter( + # shared_basepath / traj, sim_settings.checkpoint_interval.m + # )) + simulation.reporters.append(openmm.app.CheckpointReporter( + shared_basepath / sim_settings.checkpoint_storage, + sim_settings.checkpoint_interval.m )) simulation.reporters.append(openmm.app.StateDataReporter( sys.stdout, From 95b9c16b5c23a7cd5316402dfd55195fa3074de2 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 7 Aug 2023 10:02:13 +0200 Subject: [PATCH 13/66] change reporter to xtc --- openfe/protocols/openmm_md/plain_md_methods.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 5b7608b76..6572589be 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -28,6 +28,7 @@ import uuid import time from mdtraj.reporters import NetCDFReporter, XTCReporter +import netCDF4 as netcdf from gufe import ( settings, ChemicalSystem, SmallMoleculeComponent, @@ -407,12 +408,12 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - traj = "prod.dcd" - simulation.reporters.append(NetCDFReporter( - shared_basepath / sim_settings.output_filename, - sim_settings.checkpoint_interval.m)) - # simulation.reporters.append(XTCReporter( - # shared_basepath / traj, sim_settings.checkpoint_interval.m)) + traj = "prod.xtc" + # simulation.reporters.append(NetCDFReporter( + # shared_basepath / sim_settings.output_filename, + # sim_settings.checkpoint_interval.m)) + simulation.reporters.append(XTCReporter(file = str( + shared_basepath / traj), reportInterval=sim_settings.checkpoint_interval.m)) # simulation.reporters.append(openmm.app.DCDReporter( # shared_basepath / traj, sim_settings.checkpoint_interval.m # )) From 96b69f07795a39904d8fbe4f7cf2c192d16ba937 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 7 Aug 2023 10:54:42 +0200 Subject: [PATCH 14/66] Fix reporters --- .../protocols/openmm_md/plain_md_methods.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 6572589be..83ce5a42f 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -408,25 +408,28 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - traj = "prod.xtc" - # simulation.reporters.append(NetCDFReporter( - # shared_basepath / sim_settings.output_filename, - # sim_settings.checkpoint_interval.m)) - simulation.reporters.append(XTCReporter(file = str( - shared_basepath / traj), reportInterval=sim_settings.checkpoint_interval.m)) - # simulation.reporters.append(openmm.app.DCDReporter( - # shared_basepath / traj, sim_settings.checkpoint_interval.m - # )) - simulation.reporters.append(openmm.app.CheckpointReporter( - shared_basepath / sim_settings.checkpoint_storage, + chk_out = "checkpoint.chk" + log_out = "output.log" + simulation.reporters.append(NetCDFReporter(file= + str(shared_basepath / sim_settings.output_filename), reportInterval= + sim_settings.checkpoint_interval.m)) + # simulation.reporters.append(XTCReporter(file = str( + # shared_basepath / traj), reportInterval=sim_settings.checkpoint_interval.m)) + simulation.reporters.append(openmm.app.CheckpointReporter(file=str( + shared_basepath / chk_out), reportInterval= sim_settings.checkpoint_interval.m )) simulation.reporters.append(openmm.app.StateDataReporter( - sys.stdout, + shared_basepath / log_out, sim_settings.checkpoint_interval.m, step=True, + time=True, potentialEnergy=True, - temperature=True + kineticEnergy=True, + totalEnergy=True, + temperature=True, + volume=True, + density=True, )) t0 = time.time() simulation.step(prod_steps) @@ -434,7 +437,7 @@ def run(self, *, dry=False, verbose=True, logger.info(f"Completed simulation in {t1 - t0} seconds") nc = shared_basepath / sim_settings.output_filename - chk = shared_basepath / sim_settings.checkpoint_storage + chk = shared_basepath / chk_out finally: From fca9c238294dd51a7077f7fa05d5494514a4556b Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 7 Aug 2023 11:03:34 +0200 Subject: [PATCH 15/66] Fix StateReporter --- openfe/protocols/openmm_md/plain_md_methods.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 83ce5a42f..d291a5031 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -420,7 +420,7 @@ def run(self, *, dry=False, verbose=True, sim_settings.checkpoint_interval.m )) simulation.reporters.append(openmm.app.StateDataReporter( - shared_basepath / log_out, + str(shared_basepath / log_out), sim_settings.checkpoint_interval.m, step=True, time=True, From 48b626c2271f0a13e670d04539a46e8e7b7f79d0 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 7 Aug 2023 14:26:07 +0200 Subject: [PATCH 16/66] Add output indices traj and other fixes --- .../protocols/openmm_md/plain_md_methods.py | 43 +++++++++---------- .../protocols/openmm_md/plain_md_settings.py | 6 +-- openfe/protocols/openmm_utils/omm_settings.py | 24 ++++++++++- 3 files changed, 47 insertions(+), 26 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index d291a5031..3241b78a0 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -11,7 +11,6 @@ from __future__ import annotations import logging -import sys from collections import defaultdict import gufe @@ -24,20 +23,18 @@ from openmm import unit as omm_unit import pathlib from typing import Any, Iterable -import openmmtools import uuid import time -from mdtraj.reporters import NetCDFReporter, XTCReporter -import netCDF4 as netcdf +import mdtraj +from mdtraj.reporters import XTCReporter from gufe import ( - settings, ChemicalSystem, SmallMoleculeComponent, - ProteinComponent + settings, ChemicalSystem, ) from openfe.protocols.openmm_md.plain_md_settings import ( PlainMDProtocolSettings, SystemSettings, SolvationSettings, OpenMMEngineSettings, - IntegratorSettings, SimulationSettings, + IntegratorSettings, SimulationSettingsMD, RepeatSettings ) @@ -101,7 +98,7 @@ def _default_settings(cls): solvation_settings=SolvationSettings(), engine_settings=OpenMMEngineSettings(), integrator_settings=IntegratorSettings(), - simulation_settings=SimulationSettings( + simulation_settings=SimulationSettingsMD( equilibration_length=1.0 * unit.nanosecond, production_length=5.0 * unit.nanosecond, ), @@ -293,7 +290,7 @@ def run(self, *, dry=False, verbose=True, system_settings: SystemSettings = protocol_settings.system_settings solvation_settings: SolvationSettings = \ protocol_settings.solvation_settings - sim_settings: SimulationSettings = \ + sim_settings: SimulationSettingsMD = \ protocol_settings.simulation_settings timestep = protocol_settings.integrator_settings.timestep mc_steps = protocol_settings.integrator_settings.n_steps.m @@ -408,19 +405,21 @@ def run(self, *, dry=False, verbose=True, # production if verbose: logger.info("running production phase") - chk_out = "checkpoint.chk" - log_out = "output.log" - simulation.reporters.append(NetCDFReporter(file= - str(shared_basepath / sim_settings.output_filename), reportInterval= - sim_settings.checkpoint_interval.m)) - # simulation.reporters.append(XTCReporter(file = str( - # shared_basepath / traj), reportInterval=sim_settings.checkpoint_interval.m)) - simulation.reporters.append(openmm.app.CheckpointReporter(file=str( - shared_basepath / chk_out), reportInterval= - sim_settings.checkpoint_interval.m - )) + + # Get the sub selection of the system to print coords for + selection_indices = mdtraj.Topology.from_openmm( + stateA_topology).select(sim_settings.output_indices) + + # Setup the reporters + simulation.reporters.append(XTCReporter( + file=str(shared_basepath / sim_settings.output_filename), + reportInterval=sim_settings.checkpoint_interval.m, + atomSubset=selection_indices)) + simulation.reporters.append(openmm.app.CheckpointReporter( + file=str(shared_basepath / sim_settings.checkpoint_storage), + reportInterval=sim_settings.checkpoint_interval.m)) simulation.reporters.append(openmm.app.StateDataReporter( - str(shared_basepath / log_out), + str(shared_basepath / sim_settings.log_output), sim_settings.checkpoint_interval.m, step=True, time=True, @@ -437,7 +436,7 @@ def run(self, *, dry=False, verbose=True, logger.info(f"Completed simulation in {t1 - t0} seconds") nc = shared_basepath / sim_settings.output_filename - chk = shared_basepath / chk_out + chk = shared_basepath / sim_settings.checkpoint_storage finally: diff --git a/openfe/protocols/openmm_md/plain_md_settings.py b/openfe/protocols/openmm_md/plain_md_settings.py index edf022949..81c7321d1 100644 --- a/openfe/protocols/openmm_md/plain_md_settings.py +++ b/openfe/protocols/openmm_md/plain_md_settings.py @@ -9,8 +9,8 @@ """ from openfe.protocols.openmm_utils.omm_settings import ( Settings, SystemSettings, - SolvationSettings, OpenMMEngineSettings, - IntegratorSettings, SimulationSettings, RepeatSettings + SolvationSettings, OpenMMEngineSettings, SimulationSettingsMD, + IntegratorSettings, RepeatSettings ) @@ -29,7 +29,7 @@ class Config: integrator_settings: IntegratorSettings # Simulation run settings - simulation_settings: SimulationSettings + simulation_settings: SimulationSettingsMD # Setting number of repeats of md simulation repeat_settings : RepeatSettings \ No newline at end of file diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index d4e503a64..daae8c317 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -455,4 +455,26 @@ def must_be_positive(cls, v): errmsg = ("Minimization steps, MD lengths, and checkpoint " "intervals must be positive") raise ValueError(errmsg) - return v \ No newline at end of file + return v + +class SimulationSettingsMD(SimulationSettings): + """ + Settings for simulation control for plain MD simulations, including + writing outputs + """ + class Config: + arbitrary_types_allowed = True + + # reporter settings + output_filename = 'simulation.xtc' + """Path to the storage file for analysis. Default 'simulation.xtc'.""" + checkpoint_storage = 'checkpoint.chk' + """ + Separate filename for the checkpoint file. Note, this should + not be a full path, just a filename. Default 'checkpoint.chk'. + """ + log_output = 'simulation.log' + """ + Filename for writing the log of the MD simulation, including timesteps, + energies, density, etc. + """ \ No newline at end of file From 0b61cbf13dcb792232153e59cd5a33bba3077c20 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 8 Aug 2023 10:24:22 +0200 Subject: [PATCH 17/66] changes test --- .../test_openmm_plain_md_protocols.py | 36 ++++++------------- 1 file changed, 10 insertions(+), 26 deletions(-) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index fc6aea4b2..b8108b8f7 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -1,34 +1,17 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/openfe -import os -from io import StringIO -import numpy as np + import gufe -from gufe.tests.test_tokenization import GufeTokenizableTestsMixin import pytest from unittest import mock from openff.units import unit -from importlib import resources -import xml.etree.ElementTree as ET -from openff.units.openmm import from_openmm, to_openmm -from openmm import app, XmlSerializer, MonteCarloBarostat -from openmm import unit as omm_unit + +from openff.units.openmm import to_openmm from openmmtools.states import ThermodynamicState -import pathlib -from rdkit import Chem -from rdkit.Geometry import Point3D -import mdtraj as mdt - -import openfe -from openfe import setup -from openfe.protocols import openmm_md -from openfe.protocols.openmm_utils import system_creation -from openmmforcefields.generators import SMIRNOFFTemplateGenerator -from openff.units.openmm import ensure_quantity + from openfe.protocols.openmm_md.plain_md_methods import ( PlainMDProtocol, PlainMDProtocolUnit, PlainMDProtocolResult, ) -from openfe.protocols.openmm_md.plain_md_settings import PlainMDProtocolSettings def test_create_default_settings(): @@ -108,10 +91,10 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): with tmpdir.as_cwd(): sim = dag_unit.run(dry=True)['debug']['system'] print(sim) - assert not ThermodynamicState(sim, temperature= - to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic - assert ThermodynamicState(sim, temperature= - to_openmm(protocol.settings.thermo_settings.temperature)).barostat is None + assert not ThermodynamicState(sim, temperature=to_openmm( + protocol.settings.thermo_settings.temperature)).is_periodic + assert ThermodynamicState(sim, temperature=to_openmm( + protocol.settings.thermo_settings.temperature)).barostat is None # Check MDTtraj Topologies # 16 atoms: @@ -238,6 +221,7 @@ def test_dry_many_molecules_solvent( $$$$ """ + def test_dry_run_ligand_tip4p(benzene_system, tmpdir): """ Test that we can create a system with virtual sites in the @@ -246,7 +230,7 @@ def test_dry_run_ligand_tip4p(benzene_system, tmpdir): settings = PlainMDProtocol.default_settings() settings.forcefield_settings.forcefields = [ "amber/ff14SB.xml", # ff14SB protein force field - "amber/tip4pew_standard.xml", # FF we are testsing with the fun VS + "amber/tip4pew_standard.xml", # FF we are testsing with the fun VS "amber/phosaa10.xml", # Handles THE TPO ] settings.solvation_settings.solvent_padding = 1.0 * unit.nanometer From e77cad6230f178fd174437c57b0736ec7535b3c3 Mon Sep 17 00:00:00 2001 From: Hannah Baumann <43765638+hannahbaumann@users.noreply.github.com> Date: Fri, 3 Nov 2023 10:23:03 +0100 Subject: [PATCH 18/66] style_fix Co-authored-by: Irfan Alibay --- openfe/protocols/openmm_md/plain_md_methods.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 3241b78a0..df2efc214 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -124,8 +124,7 @@ def _create( system_validation.validate_protein(stateA) # actually create and return Units - solvent_comp, protein_comp, small_mols = \ - system_validation.get_components(stateA) + solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) lig_name = small_mols[0].name # our DAG has no dependencies, so just list units n_repeats = self.settings.repeat_settings.n_repeats From e0ea862dac424ae4affd62a9b9233defcf7b2d7c Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 6 Nov 2023 08:58:05 +0100 Subject: [PATCH 19/66] remove preequil --- .../protocols/openmm_md/plain_md_methods.py | 63 +++++++------------ .../test_openmm_plain_md_protocols.py | 14 ----- 2 files changed, 21 insertions(+), 56 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index df2efc214..234333273 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -125,14 +125,17 @@ def _create( # actually create and return Units solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) - lig_name = small_mols[0].name + if len(small_mols) == 0: + system_name = f'{protein_comp.name} {solvent_comp.name}' + else: + system_name = f'{protein_comp.name} {small_mols[0].name} {solvent_comp.name}' # our DAG has no dependencies, so just list units n_repeats = self.settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( stateA=stateA, stateB=stateB, settings=self.settings, generation=0, repeat_id=int(uuid.uuid4()), - name=f'{lig_name} repeat {i} generation 0') + name=f'{system_name} repeat {i} generation 0') for i in range(n_repeats)] return units @@ -205,37 +208,6 @@ def __init__(self, *, generation=generation ) - @staticmethod - def _pre_minimize(system: openmm.System, - positions: omm_unit.Quantity) -> npt.NDArray: - """ - Short CPU minization of System to avoid GPU NaNs - - Parameters - ---------- - system : openmm.System - An OpenMM System to minimize. - positions : openmm.unit.Quantity - Initial positions for the system. - - Returns - ------- - minimized_positions : npt.NDArray - Minimized positions - """ - integrator = openmm.VerletIntegrator(0.001) - context = openmm.Context( - system, integrator, - openmm.Platform.getPlatformByName('CPU'), - ) - context.setPositions(positions) - # Do a quick 100 steps minimization, usually avoids NaNs - openmm.LocalEnergyMinimizer.minimize( - context, maxIterations=100 - ) - state = context.getState(getPositions=True) - minimized_positions = state.getPositions(asNumpy=True) - return minimized_positions def run(self, *, dry=False, verbose=True, scratch_basepath=None, @@ -322,15 +294,22 @@ def run(self, *, dry=False, verbose=True, has_solvent=solvent_comp is not None, ) - # b. force the creation of parameters - # This is necessary because we need to have the FF generated ahead of - # solvating the system. - # Note: by default this is cached to ctx.shared/db.json so shouldn't - # incur too large a cost - for comp in small_mols: - offmol = comp.to_openff() - system_generator.create_system(offmol.to_topology().to_openmm(), - molecules=[offmol]) + smc_components: dict[SmallMoleculeComponent, OFFMolecule] + + smc_components = {small_mols[0]: small_mols[0].to_openff()} + for mol in smc_components.values(): + # don't do this if we have user charges + if not (mol.partial_charges is not None and np.any( + mol.partial_charges)): + # due to issues with partial charge generation in ambertools + # we default to using the input conformer for charge generation + mol.assign_partial_charges( + 'am1bcc', use_conformers=mol.conformers + ) + + system_generator.create_system( + mol.to_topology().to_openmm(), molecules=[mol] + ) # c. get OpenMM Modeller + a dictionary of resids for each component stateA_modeller, comp_resids = system_creation.get_omm_modeller( diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index b8108b8f7..4e4d15500 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -96,20 +96,6 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): assert ThermodynamicState(sim, temperature=to_openmm( protocol.settings.thermo_settings.temperature)).barostat is None - # Check MDTtraj Topologies - # 16 atoms: - # 11 common atoms, 1 extra hydrogen in benzene, 4 extra in toluene - # 12 bonds in benzene + 4 extra tolunee bonds - # assert len(list(.atoms)) == 16 - # assert len(list(htf.omm_hybrid_topology.atoms())) == 16 - # assert len(list(htf.hybrid_topology.bonds)) == 16 - # assert len(list(htf.omm_hybrid_topology.bonds())) == 16 - - # # smoke test - can convert back the mdtraj topology - # ret_top = mdt.Topology.to_openmm(htf.hybrid_topology) - # assert len(list(ret_top.atoms())) == 16 - # assert len(list(ret_top.bonds())) == 16 - def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir): vac_settings = PlainMDProtocol.default_settings() From 4b52c2537908ae4068b50d4ee948536c661ad70d Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 8 Nov 2023 14:13:40 +0100 Subject: [PATCH 20/66] Change system name --- openfe/protocols/openmm_md/plain_md_methods.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 234333273..ec8367aa8 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -126,9 +126,9 @@ def _create( # actually create and return Units solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) if len(small_mols) == 0: - system_name = f'{protein_comp.name} {solvent_comp.name}' + system_name = 'MD' else: - system_name = f'{protein_comp.name} {small_mols[0].name} {solvent_comp.name}' + system_name = f'MD {small_mols[0].name}' # our DAG has no dependencies, so just list units n_repeats = self.settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( @@ -166,7 +166,7 @@ def _gather( class PlainMDProtocolUnit(gufe.ProtocolUnit): """ - Base class for plain MD simulations (NonTransformation. + Base class for plain MD simulations (NonTransformation). """ def __init__(self, *, From 8338cf391ee13a8cc7e03f04b5b03e31f0f1b7a1 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 9 Nov 2023 11:21:26 +0100 Subject: [PATCH 21/66] Save pdb of the system and of minimized structure --- .../protocols/openmm_md/plain_md_methods.py | 44 ++++++++++++++----- openfe/protocols/openmm_utils/omm_settings.py | 9 ++++ 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index ec8367aa8..cd6f428f3 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -44,11 +44,11 @@ ) logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) class PlainMDProtocolResult(gufe.ProtocolResult): - """EMPTY, Dict-like container for the output of a PlainMDProtocol""" + """Dict-like container for the output of a PlainMDProtocol + outputs filenames for the pdb file and trajectory""" def __init__(self, **data): super().__init__(**data) # data is mapping of str(repeat_id): list[protocolunitresults] @@ -70,6 +70,20 @@ def get_uncertainty(self): return None + def get_traj_filename(self): + """String of trajectory file name""" + protocol_settings: PlainMDProtocolSettings = self._inputs[ + 'settings'] + + return protocol_settings.simulation_settings.output_filename + + def get_pdb_filename(self): + """String of pdb file name""" + protocol_settings: PlainMDProtocolSettings = self._inputs[ + 'settings'] + + return protocol_settings.simulation_settings.output_structure + class PlainMDProtocol(gufe.Protocol): result_cls = PlainMDProtocolResult @@ -124,6 +138,7 @@ def _create( system_validation.validate_protein(stateA) # actually create and return Units + # TODO: Deal with multiple ProteinComponents solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) if len(small_mols) == 0: system_name = 'MD' @@ -296,7 +311,7 @@ def run(self, *, dry=False, verbose=True, smc_components: dict[SmallMoleculeComponent, OFFMolecule] - smc_components = {small_mols[0]: small_mols[0].to_openff()} + smc_components = {i: i.to_openff() for i in small_mols} for mol in smc_components.values(): # don't do this if we have user charges if not (mol.partial_charges is not None and np.any( @@ -333,6 +348,12 @@ def run(self, *, dry=False, verbose=True, molecules=[s.to_openff() for s in small_mols], ) + # f. Save pdb of entire system + with open(shared_basepath / sim_settings.output_structure, "w") as f: + openmm.app.PDBFile.writeFile( + stateA_topology, stateA_positions, file=f, keepIds=True + ) + # 10. Get platform platform = compute.get_openmm_platform( protocol_settings.engine_settings.compute_platform @@ -365,13 +386,20 @@ def run(self, *, dry=False, verbose=True, simulation.minimizeEnergy( maxIterations=sim_settings.minimization_steps ) - min_pdb_file = "minimized.pdb" + + # Get the sub selection of the system to print coords for + selection_indices = mdtraj.Topology.from_openmm( + stateA_topology).select(sim_settings.output_indices) + positions = simulation.context.getState( getPositions=True, enforcePeriodicBox=False ).getPositions() - with open(shared_basepath / min_pdb_file, "w") as f: + # Store subset of atoms, specified in input, as PDB file + with open(shared_basepath / sim_settings.minimized_structure, "w") as f: openmm.app.PDBFile.writeFile( - simulation.topology, positions, file=f, keepIds=True + simulation.topology.subset(selection_indices), + positions[selection_indices, :], + file=f, keepIds=True, ) # equilibrate if verbose: @@ -384,10 +412,6 @@ def run(self, *, dry=False, verbose=True, if verbose: logger.info("running production phase") - # Get the sub selection of the system to print coords for - selection_indices = mdtraj.Topology.from_openmm( - stateA_topology).select(sim_settings.output_indices) - # Setup the reporters simulation.reporters.append(XTCReporter( file=str(shared_basepath / sim_settings.output_filename), diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 6e4b91f3f..273312300 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -385,6 +385,15 @@ class Config: # reporter settings output_filename = 'simulation.xtc' """Path to the storage file for analysis. Default 'simulation.xtc'.""" + trajectory_interval = 250 * unit.timestep + """ + Frequency to write the xtc file. Default 5000 * unit.timestep. + """ + output_structure = 'system.pdb' + """Path to the pdb file of the full system. Default 'system.pdb'.""" + minimized_structure = 'minimized.pdb' + """Path to the pdb file of the system after minimization. + Only the specified atom subset is saved. Default 'minimized.pdb'.""" checkpoint_storage = 'checkpoint.chk' """ Separate filename for the checkpoint file. Note, this should From 9c74b84f8638b9d1904bd4e56a27881144a3068f Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 14 Nov 2023 15:04:04 +0100 Subject: [PATCH 22/66] Fixes plain MD method --- openfe/protocols/openmm_md/plain_md_methods.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index cd6f428f3..000545756 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -395,17 +395,21 @@ def run(self, *, dry=False, verbose=True, getPositions=True, enforcePeriodicBox=False ).getPositions() # Store subset of atoms, specified in input, as PDB file - with open(shared_basepath / sim_settings.minimized_structure, "w") as f: - openmm.app.PDBFile.writeFile( - simulation.topology.subset(selection_indices), - positions[selection_indices, :], - file=f, keepIds=True, - ) + mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) + traj = mdtraj.Trajectory( + positions[selection_indices, :], + mdtraj_top.subset(selection_indices), + ) + traj.save_pdb( + shared_basepath / sim_settings.minimized_structure + ) # equilibrate + # NVT equilibration if verbose: logger.info("equilibrating systems") simulation.context.setVelocitiesToTemperature( to_openmm(thermo_settings.temperature)) + simulation.context.setParameter(MonteCarloBarostat.setFrequency(), 0) simulation.step(equil_steps) # production From 3ff34e2006ffc32804cba5d4e97f1cd70eaf21b1 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 17 Nov 2023 09:49:49 +0100 Subject: [PATCH 23/66] Fix docstring --- openfe/protocols/openmm_utils/omm_settings.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 273312300..29b4a6b94 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -371,7 +371,7 @@ class RepeatSettings(SettingsBaseModel): n_repeats: int = 1 """ - Number of independent repeats to run. Default 3 + Number of independent repeats to run. Default 1 """ class SimulationSettingsMD(SimulationSettings): @@ -387,10 +387,10 @@ class Config: """Path to the storage file for analysis. Default 'simulation.xtc'.""" trajectory_interval = 250 * unit.timestep """ - Frequency to write the xtc file. Default 5000 * unit.timestep. + Frequency to write the xtc file. Default 250 * unit.timestep. """ output_structure = 'system.pdb' - """Path to the pdb file of the full system. Default 'system.pdb'.""" + """Path to the pdb file of the full pre-minimized system. Default 'system.pdb'.""" minimized_structure = 'minimized.pdb' """Path to the pdb file of the system after minimization. Only the specified atom subset is saved. Default 'minimized.pdb'.""" From 22eec317b7dc8e6618c75ea963fd7a3939c2c4eb Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 17 Nov 2023 10:11:44 +0100 Subject: [PATCH 24/66] Add ffcache test --- .../test_openmm_plain_md_protocols.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index 4e4d15500..1c24ebb7c 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -96,6 +96,32 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): assert ThermodynamicState(sim, temperature=to_openmm( protocol.settings.thermo_settings.temperature)).barostat is None +def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir): + + vac_settings = PlainMDProtocol.default_settings() + vac_settings.system_settings.nonbonded_method = 'nocutoff' + vac_settings.simulation_settings.forcefield_cache = None + + protocol = PlainMDProtocol( + settings=vac_settings, + ) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sim = dag_unit.run(dry=True)['debug']['system'] + print(sim) + assert not ThermodynamicState(sim, temperature=to_openmm( + protocol.settings.thermo_settings.temperature)).is_periodic + assert ThermodynamicState(sim, temperature=to_openmm( + protocol.settings.thermo_settings.temperature)).barostat is None + def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir): vac_settings = PlainMDProtocol.default_settings() @@ -352,3 +378,6 @@ def test_gather(solvent_protocol_dag, tmpdir): res = prot.gather([dagres]) assert isinstance(res, PlainMDProtocolResult) + +def test_ffcache_none(): + dryrun = pu.run(dry=True, shared_basepath=tmp) \ No newline at end of file From e160600314a967f5d96bc43945927a0a195746b2 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 20 Nov 2023 11:24:02 +0100 Subject: [PATCH 25/66] Add NVT and NPT equilibration --- .../protocols/openmm_md/plain_md_methods.py | 31 ++++++++++++++++--- openfe/protocols/openmm_utils/omm_settings.py | 20 ++++++++++-- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 000545756..bd0ab9eb5 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -286,8 +286,9 @@ def run(self, *, dry=False, verbose=True, settings_validation.validate_timestep( forcefield_settings.hydrogen_mass, timestep ) - equil_steps, prod_steps = settings_validation.get_simsteps( - equil_length=sim_settings.equilibration_length, + equil_steps_nvt, equil_steps_npt, prod_steps = settings_validation.get_simsteps( + equil_length_nvt=sim_settings.equilibration_length_nvt, + equil_length_npt=sim_settings.equilibration_length_npt, prod_length=sim_settings.production_length, timestep=timestep, mc_steps=mc_steps ) @@ -406,11 +407,31 @@ def run(self, *, dry=False, verbose=True, # equilibrate # NVT equilibration if verbose: - logger.info("equilibrating systems") + logger.info("Running NVT equilibration") simulation.context.setVelocitiesToTemperature( to_openmm(thermo_settings.temperature)) simulation.context.setParameter(MonteCarloBarostat.setFrequency(), 0) - simulation.step(equil_steps) + + t0 = time.time() + simulation.step(equil_steps_nvt) + t1 = time.time() + logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") + + # NPT equilibration + if verbose: + logger.info("Running NPT equilibration") + simulation.context.setVelocitiesToTemperature( + to_openmm(thermo_settings.temperature)) + simulation.context.setParameter( + MonteCarloBarostat.setFrequency(), + integrator_settings.barostat_frequency, + ) + + t0 = time.time() + simulation.step(equil_steps_npt) + t1 = time.time() + logger.info( + f"Completed NPT equilibration in {t1 - t0} seconds") # production if verbose: @@ -419,7 +440,7 @@ def run(self, *, dry=False, verbose=True, # Setup the reporters simulation.reporters.append(XTCReporter( file=str(shared_basepath / sim_settings.output_filename), - reportInterval=sim_settings.checkpoint_interval.m, + reportInterval=sim_settings.trajectory_interval.m, atomSubset=selection_indices)) simulation.reporters.append(openmm.app.CheckpointReporter( file=str(shared_basepath / sim_settings.checkpoint_storage), diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 29b4a6b94..e0895399f 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -382,12 +382,28 @@ class SimulationSettingsMD(SimulationSettings): class Config: arbitrary_types_allowed = True + equilibration_length_nvt: unit.Quantity + """ + Length of the equilibration phase in the NVT ensemble in units of time. + The total number of steps from this equilibration length + (i.e. ``equilibration_length`` / :class:`IntegratorSettings.timestep`) + must be a multiple of the value defined for + :class:`IntegratorSettings.n_steps`. + """ + equilibration_length_npt: unit.Quantity + """ + Length of the equilibration phase in the NPT ensemble in units of time. + The total number of steps from this equilibration length + (i.e. ``equilibration_length`` / :class:`IntegratorSettings.timestep`) + must be a multiple of the value defined for + :class:`IntegratorSettings.n_steps`. + """ # reporter settings output_filename = 'simulation.xtc' """Path to the storage file for analysis. Default 'simulation.xtc'.""" - trajectory_interval = 250 * unit.timestep + trajectory_interval = 5000 * unit.timestep """ - Frequency to write the xtc file. Default 250 * unit.timestep. + Frequency to write the xtc file. Default 5000 * unit.timestep. """ output_structure = 'system.pdb' """Path to the pdb file of the full pre-minimized system. Default 'system.pdb'.""" From 24d5a2f1ef6d8b33d7ef32078a566cda5c50817e Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 20 Nov 2023 11:38:03 +0100 Subject: [PATCH 26/66] Fix equilbration settings --- openfe/protocols/openmm_md/plain_md_methods.py | 3 ++- openfe/protocols/openmm_utils/omm_settings.py | 10 +--------- .../tests/protocols/test_openmm_plain_md_protocols.py | 5 +---- 3 files changed, 4 insertions(+), 14 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index bd0ab9eb5..5dd69a954 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -113,6 +113,7 @@ def _default_settings(cls): engine_settings=OpenMMEngineSettings(), integrator_settings=IntegratorSettings(), simulation_settings=SimulationSettingsMD( + equilibration_length_nvt=0.1 * unit.nanosecond, equilibration_length=1.0 * unit.nanosecond, production_length=5.0 * unit.nanosecond, ), @@ -288,7 +289,7 @@ def run(self, *, dry=False, verbose=True, ) equil_steps_nvt, equil_steps_npt, prod_steps = settings_validation.get_simsteps( equil_length_nvt=sim_settings.equilibration_length_nvt, - equil_length_npt=sim_settings.equilibration_length_npt, + equil_length=sim_settings.equilibration_length, prod_length=sim_settings.production_length, timestep=timestep, mc_steps=mc_steps ) diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index e0895399f..cb6f1140a 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -386,15 +386,7 @@ class Config: """ Length of the equilibration phase in the NVT ensemble in units of time. The total number of steps from this equilibration length - (i.e. ``equilibration_length`` / :class:`IntegratorSettings.timestep`) - must be a multiple of the value defined for - :class:`IntegratorSettings.n_steps`. - """ - equilibration_length_npt: unit.Quantity - """ - Length of the equilibration phase in the NPT ensemble in units of time. - The total number of steps from this equilibration length - (i.e. ``equilibration_length`` / :class:`IntegratorSettings.timestep`) + (i.e. ``equilibration_length_nvt`` / :class:`IntegratorSettings.timestep`) must be a multiple of the value defined for :class:`IntegratorSettings.n_steps`. """ diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index 1c24ebb7c..ba3df33dc 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -377,7 +377,4 @@ def test_gather(solvent_protocol_dag, tmpdir): res = prot.gather([dagres]) - assert isinstance(res, PlainMDProtocolResult) - -def test_ffcache_none(): - dryrun = pu.run(dry=True, shared_basepath=tmp) \ No newline at end of file + assert isinstance(res, PlainMDProtocolResult) \ No newline at end of file From 0c66327126e4a75aa365b132817038a0c8d3aa9e Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 20 Nov 2023 15:09:53 +0100 Subject: [PATCH 27/66] Change get_simsteps to handle NVT equilibration --- .../protocols/openmm_md/plain_md_methods.py | 15 +++++-- .../openmm_utils/settings_validation.py | 43 ++++++++----------- 2 files changed, 28 insertions(+), 30 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 5dd69a954..59b625363 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -287,10 +287,17 @@ def run(self, *, dry=False, verbose=True, settings_validation.validate_timestep( forcefield_settings.hydrogen_mass, timestep ) - equil_steps_nvt, equil_steps_npt, prod_steps = settings_validation.get_simsteps( - equil_length_nvt=sim_settings.equilibration_length_nvt, - equil_length=sim_settings.equilibration_length, - prod_length=sim_settings.production_length, + + equil_steps_nvt = settings_validation.get_simsteps( + sim_length=sim_settings.equilibration_length_nvt, + timestep=timestep, mc_steps=mc_steps + ) + equil_steps_npt = settings_validation.get_simsteps( + sim_length=sim_settings.equilibration_length, + timestep=timestep, mc_steps=mc_steps + ) + prod_steps = settings_validation.get_simsteps( + sim_length=sim_settings.production_length, timestep=timestep, mc_steps=mc_steps ) diff --git a/openfe/protocols/openmm_utils/settings_validation.py b/openfe/protocols/openmm_utils/settings_validation.py index 5745a178d..2deb2f90d 100644 --- a/openfe/protocols/openmm_utils/settings_validation.py +++ b/openfe/protocols/openmm_utils/settings_validation.py @@ -37,17 +37,15 @@ def validate_timestep(hmass: float, timestep: unit.Quantity): raise ValueError(errmsg) -def get_simsteps(equil_length: unit.Quantity, prod_length: unit.Quantity, - timestep: unit.Quantity, mc_steps: int) -> Tuple[int, int]: +def get_simsteps(sim_length: unit.Quantity, + timestep: unit.Quantity, mc_steps: int) -> int: """ - Gets and validates the number of equilibration and production steps. + Gets and validates the number of simulation steps. Parameters ---------- - equil_length : unit.Quantity - Simulation equilibration length. - prod_length : unit.Quantity - Simulation production length. + sim_length : unit.Quantity + Simulation length. timestep : unit.Quantity Integration timestep. mc_steps : int @@ -55,29 +53,22 @@ def get_simsteps(equil_length: unit.Quantity, prod_length: unit.Quantity, Returns ------- - equil_steps : int - The number of equilibration timesteps. - prod_steps : int - The number of production timesteps. + sim_steps : int + The number of simulation timesteps. """ - equil_time = round(equil_length.to('attosecond').m) - prod_time = round(prod_length.to('attosecond').m) + sim_time = round(sim_length.to('attosecond').m) ts = round(timestep.to('attosecond').m) - equil_steps, mod = divmod(equil_time, ts) + sim_steps, mod = divmod(sim_time, ts) if mod != 0: - raise ValueError("Equilibration time not divisible by timestep") - prod_steps, mod = divmod(prod_time, ts) - if mod != 0: - raise ValueError("Production time not divisible by timestep") + raise ValueError("Simulation time not divisible by timestep") - for var in [("Equilibration", equil_steps, equil_time), - ("Production", prod_steps, prod_time)]: - if (var[1] % mc_steps) != 0: - errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " - "number of steps divisible by the number of integrator " - f"timesteps between MC moves {mc_steps}") - raise ValueError(errmsg) + var = ["Simulation", sim_steps, sim_time] + if (var[1] % mc_steps) != 0: + errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " + "number of steps divisible by the number of integrator " + f"timesteps between MC moves {mc_steps}") + raise ValueError(errmsg) - return equil_steps, prod_steps + return sim_steps From df7f92102a356e5a2122f318ff8cacb8dea02e76 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 20 Nov 2023 18:08:14 +0100 Subject: [PATCH 28/66] Fix positions for pdb storage --- openfe/protocols/openmm_md/plain_md_methods.py | 10 +++++++--- .../openmm_utils/settings_validation.py | 2 +- openfe/tests/protocols/test_openmmutils.py | 16 +++++++--------- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 59b625363..8f9df47e2 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -398,20 +398,24 @@ def run(self, *, dry=False, verbose=True, # Get the sub selection of the system to print coords for selection_indices = mdtraj.Topology.from_openmm( - stateA_topology).select(sim_settings.output_indices) + simulation.topology).select(sim_settings.output_indices) - positions = simulation.context.getState( + positions = to_openmm( + from_openmm(simulation.context.getState( getPositions=True, enforcePeriodicBox=False - ).getPositions() + ).getPositions())) # Store subset of atoms, specified in input, as PDB file mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) + traj = mdtraj.Trajectory( positions[selection_indices, :], mdtraj_top.subset(selection_indices), ) + print('here1') traj.save_pdb( shared_basepath / sim_settings.minimized_structure ) + print('here2') # equilibrate # NVT equilibration if verbose: diff --git a/openfe/protocols/openmm_utils/settings_validation.py b/openfe/protocols/openmm_utils/settings_validation.py index 2deb2f90d..3806e1618 100644 --- a/openfe/protocols/openmm_utils/settings_validation.py +++ b/openfe/protocols/openmm_utils/settings_validation.py @@ -66,7 +66,7 @@ def get_simsteps(sim_length: unit.Quantity, var = ["Simulation", sim_steps, sim_time] if (var[1] % mc_steps) != 0: - errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " + errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " "number of steps divisible by the number of integrator " f"timesteps between MC moves {mc_steps}") raise ValueError(errmsg) diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index 596601f92..857fea65f 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -22,17 +22,15 @@ def test_validate_timestep(): settings_validation.validate_timestep(2.0, 4.0 * unit.femtoseconds) -@pytest.mark.parametrize('e,p,ts,mc,es,ps', [ - [1 * unit.nanoseconds, 5 * unit.nanoseconds, 4 * unit.femtoseconds, - 250, 250000, 1250000], - [1 * unit.picoseconds, 1 * unit.picoseconds, 2 * unit.femtoseconds, - 250, 500, 500], +@pytest.mark.parametrize('s,ts,mc,es', [ + [5 * unit.nanoseconds, 4 * unit.femtoseconds, 250, 1250000], + [1 * unit.nanoseconds, 4 * unit.femtoseconds, 250, 250000], + [1 * unit.picoseconds, 2 * unit.femtoseconds, 250, 500], ]) -def test_get_simsteps(e, p, ts, mc, es, ps): - equil_steps, prod_steps = settings_validation.get_simsteps(e, p, ts, mc) +def test_get_simsteps(s, ts, mc, es): + sim_steps = settings_validation.get_simsteps(s, ts, mc) - assert equil_steps == es - assert prod_steps == ps + assert sim_steps == es @pytest.mark.parametrize('nametype, timelengths', [ From a51def26c57e1ca1f856269965e28650e3540f26 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 09:54:09 +0100 Subject: [PATCH 29/66] Fix get_simsteps tests --- openfe/tests/protocols/test_openmmutils.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index 857fea65f..eff8c15d0 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -34,29 +34,25 @@ def test_get_simsteps(s, ts, mc, es): @pytest.mark.parametrize('nametype, timelengths', [ - ['Equilibration', [1.003 * unit.picoseconds, 1 * unit.picoseconds]], - ['Production', [1 * unit.picoseconds, 1.003 * unit.picoseconds]], + ['Simulation', [1.003 * unit.picoseconds, 1 * unit.picoseconds]], ]) def test_get_simsteps_indivisible_simtime(nametype, timelengths): errmsg = f"{nametype} time not divisible by timestep" with pytest.raises(ValueError, match=errmsg): settings_validation.get_simsteps( timelengths[0], - timelengths[1], 2 * unit.femtoseconds, 100) @pytest.mark.parametrize('nametype, timelengths', [ - ['Equilibration', [1 * unit.picoseconds, 10 * unit.picoseconds]], - ['Production', [10 * unit.picoseconds, 1 * unit.picoseconds]], + ['Simulation', [1 * unit.picoseconds, 10 * unit.picoseconds]], ]) def test_mc_indivisible(nametype, timelengths): errmsg = f"{nametype} time 1.0 ps should contain" with pytest.raises(ValueError, match=errmsg): settings_validation.get_simsteps( - timelengths[0], timelengths[1], - 2 * unit.femtoseconds, 1000) + timelengths[0], 2 * unit.femtoseconds, 1000) def test_get_alchemical_components(benzene_modifications, From c5f566707d9a18d0d67976e1e3c958aac55bd8ed Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 13:54:55 +0100 Subject: [PATCH 30/66] Disable barostat for NVT equilibration, bring it back for NPT --- .../protocols/openmm_md/plain_md_methods.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 8f9df47e2..1a00ccf61 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -375,6 +375,14 @@ def run(self, *, dry=False, verbose=True, to_openmm(integrator_settings.timestep), ) + # forces = stateA_system.getForces() + # for x in forces: + # if x.getName() == 'MonteCarloBarostat': + # stateA_system.removeForce(forces.index(x)) + # barostat = openmm.MonteCarloBarostat(to_openmm(thermo_settings.pressure), + # to_openmm(thermo_settings.temperature), 0) + # stateA_system.addForce(barostat) + simulation = openmm.app.Simulation( stateA_modeller.topology, stateA_system, @@ -411,18 +419,21 @@ def run(self, *, dry=False, verbose=True, positions[selection_indices, :], mdtraj_top.subset(selection_indices), ) - print('here1') traj.save_pdb( shared_basepath / sim_settings.minimized_structure ) - print('here2') # equilibrate # NVT equilibration if verbose: logger.info("Running NVT equilibration") + + # Set barostat frequency to zero for NVT + for x in simulation.context.getSystem().getForces(): + if x.getName() == 'MonteCarloBarostat': + x.setFrequency(0) + print(x.getFrequency()) simulation.context.setVelocitiesToTemperature( to_openmm(thermo_settings.temperature)) - simulation.context.setParameter(MonteCarloBarostat.setFrequency(), 0) t0 = time.time() simulation.step(equil_steps_nvt) @@ -434,10 +445,11 @@ def run(self, *, dry=False, verbose=True, logger.info("Running NPT equilibration") simulation.context.setVelocitiesToTemperature( to_openmm(thermo_settings.temperature)) - simulation.context.setParameter( - MonteCarloBarostat.setFrequency(), - integrator_settings.barostat_frequency, - ) + + # Enable the barostat for NPT + for x in simulation.context.getSystem().getForces(): + if x.getName() == 'MonteCarloBarostat': + x.setFrequency(integrator_settings.barostat_frequency.m) t0 = time.time() simulation.step(equil_steps_npt) From 2f586a9e135fdd29cb695519fd763b6a635e0645 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 14:11:05 +0100 Subject: [PATCH 31/66] Output pdb file after NVT equilibration and after NPT equilibration --- .../protocols/openmm_md/plain_md_methods.py | 43 +++++++++++++------ openfe/protocols/openmm_utils/omm_settings.py | 6 +++ 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 1a00ccf61..e25ff7c13 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -375,14 +375,6 @@ def run(self, *, dry=False, verbose=True, to_openmm(integrator_settings.timestep), ) - # forces = stateA_system.getForces() - # for x in forces: - # if x.getName() == 'MonteCarloBarostat': - # stateA_system.removeForce(forces.index(x)) - # barostat = openmm.MonteCarloBarostat(to_openmm(thermo_settings.pressure), - # to_openmm(thermo_settings.temperature), 0) - # stateA_system.addForce(barostat) - simulation = openmm.app.Simulation( stateA_modeller.topology, stateA_system, @@ -404,7 +396,7 @@ def run(self, *, dry=False, verbose=True, maxIterations=sim_settings.minimization_steps ) - # Get the sub selection of the system to print coords for + # Get the sub selection of the system to save coords for selection_indices = mdtraj.Topology.from_openmm( simulation.topology).select(sim_settings.output_indices) @@ -431,7 +423,7 @@ def run(self, *, dry=False, verbose=True, for x in simulation.context.getSystem().getForces(): if x.getName() == 'MonteCarloBarostat': x.setFrequency(0) - print(x.getFrequency()) + simulation.context.setVelocitiesToTemperature( to_openmm(thermo_settings.temperature)) @@ -440,6 +432,20 @@ def run(self, *, dry=False, verbose=True, t1 = time.time() logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") + # Save last frame NVT equilibration + positions = to_openmm( + from_openmm(simulation.context.getState( + getPositions=True, enforcePeriodicBox=False + ).getPositions())) + + traj = mdtraj.Trajectory( + positions[selection_indices, :], + mdtraj_top.subset(selection_indices), + ) + traj.save_pdb( + shared_basepath / sim_settings.equ_NVT_structure + ) + # NPT equilibration if verbose: logger.info("Running NPT equilibration") @@ -454,8 +460,21 @@ def run(self, *, dry=False, verbose=True, t0 = time.time() simulation.step(equil_steps_npt) t1 = time.time() - logger.info( - f"Completed NPT equilibration in {t1 - t0} seconds") + logger.info(f"Completed NPT equilibration in {t1 - t0} seconds") + + # Save last frame NPT equilibration + positions = to_openmm( + from_openmm(simulation.context.getState( + getPositions=True, enforcePeriodicBox=False + ).getPositions())) + + traj = mdtraj.Trajectory( + positions[selection_indices, :], + mdtraj_top.subset(selection_indices), + ) + traj.save_pdb( + shared_basepath / sim_settings.equ_NPT_structure + ) # production if verbose: diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index cb6f1140a..57059e692 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -402,6 +402,12 @@ class Config: minimized_structure = 'minimized.pdb' """Path to the pdb file of the system after minimization. Only the specified atom subset is saved. Default 'minimized.pdb'.""" + equ_NVT_structure = 'equ_NVT.pdb' + """Path to the pdb file of the system after NVT equilibration. + Only the specified atom subset is saved. Default 'equ_NVT.pdb'.""" + equ_NPT_structure = 'equ_NPT.pdb' + """Path to the pdb file of the system after NPT equilibration. + Only the specified atom subset is saved. Default 'equ_NPT.pdb'.""" checkpoint_storage = 'checkpoint.chk' """ Separate filename for the checkpoint file. Note, this should From 917041385a4ce8981895b0a2a60a829623d7b37f Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 14:21:32 +0100 Subject: [PATCH 32/66] Add without_oechem_backend wrapper --- openfe/protocols/openmm_md/plain_md_methods.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index e25ff7c13..3c0493aef 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -27,6 +27,7 @@ import time import mdtraj from mdtraj.reporters import XTCReporter +from openfe.utils import without_oechem_backend, log_system_probe from gufe import ( settings, ChemicalSystem, @@ -524,8 +525,11 @@ def run(self, *, dry=False, verbose=True, def _execute( self, ctx: gufe.Context, **kwargs, ) -> dict[str, Any]: - outputs = self.run(scratch_basepath=ctx.scratch, - shared_basepath=ctx.shared) + log_system_probe(logging.INFO, paths=[ctx.scratch]) + + with without_oechem_backend(): + outputs = self.run(scratch_basepath=ctx.scratch, + shared_basepath=ctx.shared) return { 'repeat_id': self._inputs['repeat_id'], From b33bc7a914352d58db309bb19f5d00eb8aec4f72 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 14:33:40 +0100 Subject: [PATCH 33/66] Add log_system_probe --- openfe/utils/__init__.py | 1 + openfe/utils/system_probe.py | 41 ++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/openfe/utils/__init__.py b/openfe/utils/__init__.py index a16b6fa61..0d276a97e 100644 --- a/openfe/utils/__init__.py +++ b/openfe/utils/__init__.py @@ -4,3 +4,4 @@ from . import custom_typing from .optional_imports import requires_package from .remove_oechem import without_oechem_backend +from .system_probe import log_system_probe diff --git a/openfe/utils/system_probe.py b/openfe/utils/system_probe.py index eba22b485..535335fcc 100644 --- a/openfe/utils/system_probe.py +++ b/openfe/utils/system_probe.py @@ -3,6 +3,7 @@ import socket import sys import subprocess +from typing import Iterable, Optional import psutil @@ -453,6 +454,46 @@ def _probe_system() -> dict: } +def log_system_probe(level=logging.DEBUG, + paths: Optional[Iterable[os.PathLike]] = None): + """Print the system information via configurable logging. + + This creates a logger tree under "{__name__}.log", allowing one to turn + on/off logging of GPU info or hostname info the with + "{__name__}.log.gpu" and "{__name__}.log.hostname" loggers. + """ + if paths is None: + paths = [] + pl_paths = (pathlib.Path(p) for p in paths) + basename = __name__ + ".log" + base = logging.getLogger(basename) + gpu = logging.getLogger(basename + ".gpu") + hostname = logging.getLogger(basename + ".hostname") + loggers = [base, gpu, hostname] + if any(l.isEnabledFor(level) for l in loggers): + sysinfo = _probe_system(pl_paths)['system information'] + base.log(level, "SYSTEM CONFIG DETAILS:") + hostname.log(level, f"hostname: '{sysinfo['hostname']}'") + if gpuinfo := sysinfo['gpu information']: + for uuid, gpu_card in gpuinfo.items(): + gpu.log(level, f"GPU: {uuid=} {gpu_card['name']} " + f"mode={gpu_card['compute_mode']}") + # gpu.log(level, f"CUDA driver: {...}") + # gpu.log(level, f"CUDA toolkit: {...}") + else: # -no-cov- + gpu.log(level, f"CUDA-based GPU not found") + + psutilinfo = sysinfo["psutil information"] + memused = psutilinfo['virtual_memory']['used'] + mempct = psutilinfo['virtual_memory']['percent'] + base.log(level, f"Memory used: {bytes2human(memused)} ({mempct}%)") + for diskdev, disk in sysinfo['disk usage information'].items(): + base.log(level, f"{diskdev}: {disk['percent_used']} full " + f"({disk['available']} free)") + + if __name__ == "__main__": from pprint import pprint + pprint(_probe_system()) + From 8049534ea4e65d3e3004b67c18773b02fcee96c3 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 14:36:39 +0100 Subject: [PATCH 34/66] Add verbosity flag --- .../protocols/openmm_md/plain_md_methods.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 3c0493aef..243570aba 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -257,7 +257,7 @@ def run(self, *, dry=False, verbose=True, Exception if anything failed """ if verbose: - logger.info("creating system") + self.logger.info("Creating system") if scratch_basepath is None: scratch_basepath = pathlib.Path('.') if shared_basepath is None: @@ -391,7 +391,7 @@ def run(self, *, dry=False, verbose=True, # minimize if verbose: - logger.info("minimizing systems") + self.logger.info("minimizing systems") simulation.minimizeEnergy( maxIterations=sim_settings.minimization_steps @@ -418,7 +418,7 @@ def run(self, *, dry=False, verbose=True, # equilibrate # NVT equilibration if verbose: - logger.info("Running NVT equilibration") + self.logger.info("Running NVT equilibration") # Set barostat frequency to zero for NVT for x in simulation.context.getSystem().getForces(): @@ -431,7 +431,8 @@ def run(self, *, dry=False, verbose=True, t0 = time.time() simulation.step(equil_steps_nvt) t1 = time.time() - logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") + if verbose: + self.logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") # Save last frame NVT equilibration positions = to_openmm( @@ -449,7 +450,7 @@ def run(self, *, dry=False, verbose=True, # NPT equilibration if verbose: - logger.info("Running NPT equilibration") + self.logger.info("Running NPT equilibration") simulation.context.setVelocitiesToTemperature( to_openmm(thermo_settings.temperature)) @@ -461,7 +462,8 @@ def run(self, *, dry=False, verbose=True, t0 = time.time() simulation.step(equil_steps_npt) t1 = time.time() - logger.info(f"Completed NPT equilibration in {t1 - t0} seconds") + if verbose: + self.logger.info(f"Completed NPT equilibration in {t1 - t0} seconds") # Save last frame NPT equilibration positions = to_openmm( @@ -479,7 +481,7 @@ def run(self, *, dry=False, verbose=True, # production if verbose: - logger.info("running production phase") + self.logger.info("running production phase") # Setup the reporters simulation.reporters.append(XTCReporter( @@ -504,7 +506,8 @@ def run(self, *, dry=False, verbose=True, t0 = time.time() simulation.step(prod_steps) t1 = time.time() - logger.info(f"Completed simulation in {t1 - t0} seconds") + if verbose: + self.logger.info(f"Completed simulation in {t1 - t0} seconds") nc = shared_basepath / sim_settings.output_filename chk = shared_basepath / sim_settings.checkpoint_storage From 97171dccaec75a0755cc6a17cbf58c684995c975 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 14:47:30 +0100 Subject: [PATCH 35/66] Small fixes --- openfe/protocols/openmm_md/plain_md_methods.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 243570aba..1522ba2b9 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -14,23 +14,22 @@ from collections import defaultdict import gufe -import numpy.typing as npt import openmm from openff.units import unit from openff.units.openmm import from_openmm, to_openmm from typing import Optional from openmm import app -from openmm import unit as omm_unit import pathlib from typing import Any, Iterable import uuid import time +import numpy as np import mdtraj from mdtraj.reporters import XTCReporter from openfe.utils import without_oechem_backend, log_system_probe - from gufe import ( - settings, ChemicalSystem, + settings, ChemicalSystem, SmallMoleculeComponent, + ProteinComponent, SolventComponent ) from openfe.protocols.openmm_md.plain_md_settings import ( PlainMDProtocolSettings, SystemSettings, @@ -38,6 +37,7 @@ IntegratorSettings, SimulationSettingsMD, RepeatSettings ) +from openff.toolkit.topology import Molecule as OFFMolecule from openfe.protocols.openmm_rfe._rfe_utils import compute from openfe.protocols.openmm_utils import ( @@ -225,7 +225,6 @@ def __init__(self, *, generation=generation ) - def run(self, *, dry=False, verbose=True, scratch_basepath=None, shared_basepath=None) -> dict[str, Any]: @@ -401,10 +400,10 @@ def run(self, *, dry=False, verbose=True, selection_indices = mdtraj.Topology.from_openmm( simulation.topology).select(sim_settings.output_indices) - positions = to_openmm( - from_openmm(simulation.context.getState( - getPositions=True, enforcePeriodicBox=False - ).getPositions())) + positions = to_openmm(from_openmm( + simulation.context.getState(getPositions=True, + enforcePeriodicBox=False + ).getPositions())) # Store subset of atoms, specified in input, as PDB file mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) From b3bb0eb3df05392d242df7f8864b74d016f1269a Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 14:49:18 +0100 Subject: [PATCH 36/66] Pep8 fixes --- openfe/protocols/openmm_md/plain_md_settings.py | 3 ++- openfe/protocols/openmm_utils/omm_settings.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_settings.py b/openfe/protocols/openmm_md/plain_md_settings.py index 81c7321d1..26ff9d4a1 100644 --- a/openfe/protocols/openmm_md/plain_md_settings.py +++ b/openfe/protocols/openmm_md/plain_md_settings.py @@ -32,4 +32,5 @@ class Config: simulation_settings: SimulationSettingsMD # Setting number of repeats of md simulation - repeat_settings : RepeatSettings \ No newline at end of file + repeat_settings: RepeatSettings + \ No newline at end of file diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 57059e692..f97cb316b 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -216,7 +216,6 @@ def must_be_zero_or_positive(cls, v): class OpenMMEngineSettings(SettingsBaseModel): """OpenMM MD engine settings""" - """ TODO ---- @@ -366,6 +365,7 @@ def must_be_positive(cls, v): raise ValueError(errmsg) return v + class RepeatSettings(SettingsBaseModel): """Settings for how many independent MD runs to perform.""" @@ -374,6 +374,7 @@ class RepeatSettings(SettingsBaseModel): Number of independent repeats to run. Default 1 """ + class SimulationSettingsMD(SimulationSettings): """ Settings for simulation control for plain MD simulations, including @@ -418,4 +419,3 @@ class Config: Filename for writing the log of the MD simulation, including timesteps, energies, density, etc. """ - From d0fbe972cb6f7970e45693feac00df7a45e24aa8 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 15:52:18 +0100 Subject: [PATCH 37/66] Change system naming --- .../protocols/openmm_md/plain_md_methods.py | 20 +++++++++++++++---- .../protocols/openmm_md/plain_md_settings.py | 1 - 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 1522ba2b9..1c5bb2787 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -142,10 +142,22 @@ def _create( # actually create and return Units # TODO: Deal with multiple ProteinComponents solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) - if len(small_mols) == 0: - system_name = 'MD' - else: - system_name = f'MD {small_mols[0].name}' + # if len(small_mols) == 0: + # system_name = 'MD' + # else: + # system_name = f'MD {small_mols[0].name}' + + system_name = "Solvent MD" if solvent_comp is not None else "Vacuum MD" + + for comp in [protein_comp] + small_mols: + if comp is not None: + comp_type = comp.key.split('-')[0] + if len(comp.name) == 0: + comp_name = 'NoName' + else: + comp_name = comp.name + system_name += f" {comp_type}:{comp_name}" + # our DAG has no dependencies, so just list units n_repeats = self.settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( diff --git a/openfe/protocols/openmm_md/plain_md_settings.py b/openfe/protocols/openmm_md/plain_md_settings.py index 26ff9d4a1..d17edbe2c 100644 --- a/openfe/protocols/openmm_md/plain_md_settings.py +++ b/openfe/protocols/openmm_md/plain_md_settings.py @@ -33,4 +33,3 @@ class Config: # Setting number of repeats of md simulation repeat_settings: RepeatSettings - \ No newline at end of file From 4fbf74f8f9c3ec81a46accb62ed227259b657519 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 12:37:57 +0100 Subject: [PATCH 38/66] Move actual MD stuff into separate function to be used in other protocols --- .../protocols/openmm_md/plain_md_methods.py | 428 ++++++++++++------ 1 file changed, 291 insertions(+), 137 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 1c5bb2787..19ae6db14 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -73,17 +73,16 @@ def get_uncertainty(self): def get_traj_filename(self): """String of trajectory file name""" - protocol_settings: PlainMDProtocolSettings = self._inputs[ - 'settings'] + print(self.data.values()) + traj = [pus[0].outputs['nc'] for pus in self.data.values()] - return protocol_settings.simulation_settings.output_filename + return traj def get_pdb_filename(self): """String of pdb file name""" - protocol_settings: PlainMDProtocolSettings = self._inputs[ - 'settings'] + pdbs = [pus[0].outputs['system_pdb'] for pus in self.data.values()] - return protocol_settings.simulation_settings.output_structure + return pdbs class PlainMDProtocol(gufe.Protocol): @@ -142,10 +141,6 @@ def _create( # actually create and return Units # TODO: Deal with multiple ProteinComponents solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) - # if len(small_mols) == 0: - # system_name = 'MD' - # else: - # system_name = f'MD {small_mols[0].name}' system_name = "Solvent MD" if solvent_comp is not None else "Vacuum MD" @@ -237,6 +232,158 @@ def __init__(self, *, generation=generation ) + @staticmethod + def _run_MD(simulation: openmm.app.Simulation, + positions: omm_unit.Quantity, + settings: PlainMDProtocolSettings, + equil_steps_nvt: int, + equil_steps_npt: int, + prod_steps: int, + verbose=True, + shared_basepath=None) -> npt.NDArray: + + """ + Energy minimization, Equilibration and Production MD to be reused + in multiple protocols + + Parameters + ---------- + simulation : openmm.app.Simulation + An OpenMM simulation to simulate. + positions : openmm.unit.Quantity + Initial positions for the system. + sim_settings : SimulationSettingsMD + Settings for MD simulation + Returns + ------- + + """ + if shared_basepath is None: + shared_basepath = pathlib.Path('.') + simulation.context.setPositions(positions) + # minimize + if verbose: + logger.info("minimizing systems") + + simulation.minimizeEnergy( + maxIterations=settings.simulation_settings.minimization_steps + ) + + # Get the sub selection of the system to save coords for + selection_indices = mdtraj.Topology.from_openmm( + simulation.topology).select(settings.simulation_settings.output_indices) + + positions = to_openmm(from_openmm( + simulation.context.getState(getPositions=True, + enforcePeriodicBox=False + ).getPositions())) + # Store subset of atoms, specified in input, as PDB file + mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) + traj = mdtraj.Trajectory( + positions[selection_indices, :], + mdtraj_top.subset(selection_indices), + ) + + traj.save_pdb( + shared_basepath / settings.simulation_settings.minimized_structure + ) + # equilibrate + # NVT equilibration + + if verbose: + logger.info("Running NVT equilibration") + + # Set barostat frequency to zero for NVT + for x in simulation.context.getSystem().getForces(): + if x.getName() == 'MonteCarloBarostat': + x.setFrequency(0) + + simulation.context.setVelocitiesToTemperature( + to_openmm(settings.thermo_settings.temperature)) + + t0 = time.time() + simulation.step(equil_steps_nvt) + t1 = time.time() + if verbose: + logger.info( + f"Completed NVT equilibration in {t1 - t0} seconds") + + # Save last frame NVT equilibration + positions = to_openmm( + from_openmm(simulation.context.getState( + getPositions=True, enforcePeriodicBox=False + ).getPositions())) + + traj = mdtraj.Trajectory( + positions[selection_indices, :], + mdtraj_top.subset(selection_indices), + ) + traj.save_pdb( + shared_basepath / settings.simulation_settings.equ_NVT_structure + ) + + # NPT equilibration + if verbose: + logger.info("Running NPT equilibration") + simulation.context.setVelocitiesToTemperature( + to_openmm(settings.thermo_settings.temperature)) + + # Enable the barostat for NPT + for x in simulation.context.getSystem().getForces(): + if x.getName() == 'MonteCarloBarostat': + x.setFrequency(settings.integrator_settings.barostat_frequency.m) + + t0 = time.time() + simulation.step(equil_steps_npt) + t1 = time.time() + if verbose: + logger.info( + f"Completed NPT equilibration in {t1 - t0} seconds") + + # Save last frame NPT equilibration + positions = to_openmm( + from_openmm(simulation.context.getState( + getPositions=True, enforcePeriodicBox=False + ).getPositions())) + + traj = mdtraj.Trajectory( + positions[selection_indices, :], + mdtraj_top.subset(selection_indices), + ) + traj.save_pdb( + shared_basepath / settings.simulation_settings.equ_NPT_structure + ) + + # production + if verbose: + logger.info("running production phase") + + # Setup the reporters + simulation.reporters.append(XTCReporter( + file=str(shared_basepath / settings.simulation_settings.output_filename), + reportInterval=settings.simulation_settings.trajectory_interval.m, + atomSubset=selection_indices)) + simulation.reporters.append(openmm.app.CheckpointReporter( + file=str(shared_basepath / settings.simulation_settings.checkpoint_storage), + reportInterval=settings.simulation_settings.checkpoint_interval.m)) + simulation.reporters.append(openmm.app.StateDataReporter( + str(shared_basepath / settings.simulation_settings.log_output), + settings.simulation_settings.checkpoint_interval.m, + step=True, + time=True, + potentialEnergy=True, + kineticEnergy=True, + totalEnergy=True, + temperature=True, + volume=True, + density=True, + )) + t0 = time.time() + simulation.step(prod_steps) + t1 = time.time() + if verbose: + logger.info(f"Completed simulation in {t1 - t0} seconds") + def run(self, *, dry=False, verbose=True, scratch_basepath=None, shared_basepath=None) -> dict[str, Any]: @@ -397,131 +544,134 @@ def run(self, *, dry=False, verbose=True, try: if not dry: # pragma: no-cover - - simulation.context.setPositions(stateA_positions) - - # minimize - if verbose: - self.logger.info("minimizing systems") - - simulation.minimizeEnergy( - maxIterations=sim_settings.minimization_steps - ) - - # Get the sub selection of the system to save coords for - selection_indices = mdtraj.Topology.from_openmm( - simulation.topology).select(sim_settings.output_indices) - - positions = to_openmm(from_openmm( - simulation.context.getState(getPositions=True, - enforcePeriodicBox=False - ).getPositions())) - # Store subset of atoms, specified in input, as PDB file - mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) - - traj = mdtraj.Trajectory( - positions[selection_indices, :], - mdtraj_top.subset(selection_indices), - ) - traj.save_pdb( - shared_basepath / sim_settings.minimized_structure - ) - # equilibrate - # NVT equilibration - if verbose: - self.logger.info("Running NVT equilibration") - - # Set barostat frequency to zero for NVT - for x in simulation.context.getSystem().getForces(): - if x.getName() == 'MonteCarloBarostat': - x.setFrequency(0) - - simulation.context.setVelocitiesToTemperature( - to_openmm(thermo_settings.temperature)) - - t0 = time.time() - simulation.step(equil_steps_nvt) - t1 = time.time() - if verbose: - self.logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") - - # Save last frame NVT equilibration - positions = to_openmm( - from_openmm(simulation.context.getState( - getPositions=True, enforcePeriodicBox=False - ).getPositions())) - - traj = mdtraj.Trajectory( - positions[selection_indices, :], - mdtraj_top.subset(selection_indices), - ) - traj.save_pdb( - shared_basepath / sim_settings.equ_NVT_structure - ) - - # NPT equilibration - if verbose: - self.logger.info("Running NPT equilibration") - simulation.context.setVelocitiesToTemperature( - to_openmm(thermo_settings.temperature)) - - # Enable the barostat for NPT - for x in simulation.context.getSystem().getForces(): - if x.getName() == 'MonteCarloBarostat': - x.setFrequency(integrator_settings.barostat_frequency.m) - - t0 = time.time() - simulation.step(equil_steps_npt) - t1 = time.time() - if verbose: - self.logger.info(f"Completed NPT equilibration in {t1 - t0} seconds") - - # Save last frame NPT equilibration - positions = to_openmm( - from_openmm(simulation.context.getState( - getPositions=True, enforcePeriodicBox=False - ).getPositions())) - - traj = mdtraj.Trajectory( - positions[selection_indices, :], - mdtraj_top.subset(selection_indices), - ) - traj.save_pdb( - shared_basepath / sim_settings.equ_NPT_structure - ) - - # production - if verbose: - self.logger.info("running production phase") - - # Setup the reporters - simulation.reporters.append(XTCReporter( - file=str(shared_basepath / sim_settings.output_filename), - reportInterval=sim_settings.trajectory_interval.m, - atomSubset=selection_indices)) - simulation.reporters.append(openmm.app.CheckpointReporter( - file=str(shared_basepath / sim_settings.checkpoint_storage), - reportInterval=sim_settings.checkpoint_interval.m)) - simulation.reporters.append(openmm.app.StateDataReporter( - str(shared_basepath / sim_settings.log_output), - sim_settings.checkpoint_interval.m, - step=True, - time=True, - potentialEnergy=True, - kineticEnergy=True, - totalEnergy=True, - temperature=True, - volume=True, - density=True, - )) - t0 = time.time() - simulation.step(prod_steps) - t1 = time.time() - if verbose: - self.logger.info(f"Completed simulation in {t1 - t0} seconds") - - nc = shared_basepath / sim_settings.output_filename - chk = shared_basepath / sim_settings.checkpoint_storage + self._run_MD(simulation, + stateA_positions, + protocol_settings, + equil_steps_nvt, + equil_steps_npt, + prod_steps, + ) + # simulation.context.setPositions(stateA_positions) + # + # # minimize + # if verbose: + # self.logger.info("minimizing systems") + # + # simulation.minimizeEnergy( + # maxIterations=sim_settings.minimization_steps + # ) + # + # # Get the sub selection of the system to save coords for + # selection_indices = mdtraj.Topology.from_openmm( + # simulation.topology).select(sim_settings.output_indices) + # + # positions = to_openmm(from_openmm( + # simulation.context.getState(getPositions=True, + # enforcePeriodicBox=False + # ).getPositions())) + # # Store subset of atoms, specified in input, as PDB file + # mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) + # + # traj = mdtraj.Trajectory( + # positions[selection_indices, :], + # mdtraj_top.subset(selection_indices), + # ) + # traj.save_pdb( + # shared_basepath / sim_settings.minimized_structure + # ) + # # equilibrate + # # NVT equilibration + # if verbose: + # self.logger.info("Running NVT equilibration") + # + # # Set barostat frequency to zero for NVT + # for x in simulation.context.getSystem().getForces(): + # if x.getName() == 'MonteCarloBarostat': + # x.setFrequency(0) + # + # simulation.context.setVelocitiesToTemperature( + # to_openmm(thermo_settings.temperature)) + # + # t0 = time.time() + # simulation.step(equil_steps_nvt) + # t1 = time.time() + # if verbose: + # self.logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") + # + # # Save last frame NVT equilibration + # positions = to_openmm( + # from_openmm(simulation.context.getState( + # getPositions=True, enforcePeriodicBox=False + # ).getPositions())) + # + # traj = mdtraj.Trajectory( + # positions[selection_indices, :], + # mdtraj_top.subset(selection_indices), + # ) + # traj.save_pdb( + # shared_basepath / sim_settings.equ_NVT_structure + # ) + # + # # NPT equilibration + # if verbose: + # self.logger.info("Running NPT equilibration") + # simulation.context.setVelocitiesToTemperature( + # to_openmm(thermo_settings.temperature)) + # + # # Enable the barostat for NPT + # for x in simulation.context.getSystem().getForces(): + # if x.getName() == 'MonteCarloBarostat': + # x.setFrequency(integrator_settings.barostat_frequency.m) + # + # t0 = time.time() + # simulation.step(equil_steps_npt) + # t1 = time.time() + # if verbose: + # self.logger.info(f"Completed NPT equilibration in {t1 - t0} seconds") + # + # # Save last frame NPT equilibration + # positions = to_openmm( + # from_openmm(simulation.context.getState( + # getPositions=True, enforcePeriodicBox=False + # ).getPositions())) + # + # traj = mdtraj.Trajectory( + # positions[selection_indices, :], + # mdtraj_top.subset(selection_indices), + # ) + # traj.save_pdb( + # shared_basepath / sim_settings.equ_NPT_structure + # ) + # + # # production + # if verbose: + # self.logger.info("running production phase") + # + # # Setup the reporters + # simulation.reporters.append(XTCReporter( + # file=str(shared_basepath / sim_settings.output_filename), + # reportInterval=sim_settings.trajectory_interval.m, + # atomSubset=selection_indices)) + # simulation.reporters.append(openmm.app.CheckpointReporter( + # file=str(shared_basepath / sim_settings.checkpoint_storage), + # reportInterval=sim_settings.checkpoint_interval.m)) + # simulation.reporters.append(openmm.app.StateDataReporter( + # str(shared_basepath / sim_settings.log_output), + # sim_settings.checkpoint_interval.m, + # step=True, + # time=True, + # potentialEnergy=True, + # kineticEnergy=True, + # totalEnergy=True, + # temperature=True, + # volume=True, + # density=True, + # )) + # t0 = time.time() + # simulation.step(prod_steps) + # t1 = time.time() + # if verbose: + # self.logger.info(f"Completed simulation in {t1 - t0} seconds") finally: @@ -530,8 +680,12 @@ def run(self, *, dry=False, verbose=True, if not dry: # pragma: no-cover return { - 'nc': nc, - 'last_checkpoint': chk, + 'system_pdb': shared_basepath / sim_settings.output_filename, + 'minimizes_pdb': shared_basepath / sim_settings.minimized_structure, + 'nvt_equil_pdb': shared_basepath / sim_settings.equ_NVT_structure, + 'npt_equil_pdb': shared_basepath / sim_settings.equ_NPT_structure, + 'nc': shared_basepath / sim_settings.output_filename, + 'last_checkpoint': shared_basepath / sim_settings.checkpoint_storage, } else: return {'debug': {'system': stateA_system}} From 323f1e8b17a8a9c6d9aff2dd63c0461a5889b0a6 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 12:43:41 +0100 Subject: [PATCH 39/66] Small fixes --- openfe/protocols/openmm_md/plain_md_methods.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 19ae6db14..2e0f18c50 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -17,6 +17,7 @@ import openmm from openff.units import unit from openff.units.openmm import from_openmm, to_openmm +import openmm.unit as omm_unit from typing import Optional from openmm import app import pathlib @@ -132,7 +133,7 @@ def _create( raise NotImplementedError("Can't extend simulations yet") # Validate solvent component - nonbond = self.settings.system_settings.nonbonded_method + nonbond = self._settings.system_settings.nonbonded_method system_validation.validate_solvent(stateA, nonbond) # Validate protein component @@ -154,10 +155,10 @@ def _create( system_name += f" {comp_type}:{comp_name}" # our DAG has no dependencies, so just list units - n_repeats = self.settings.repeat_settings.n_repeats + n_repeats = self._settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( stateA=stateA, stateB=stateB, - settings=self.settings, + settings=self._settings, generation=0, repeat_id=int(uuid.uuid4()), name=f'{system_name} repeat {i} generation 0') for i in range(n_repeats)] @@ -240,7 +241,7 @@ def _run_MD(simulation: openmm.app.Simulation, equil_steps_npt: int, prod_steps: int, verbose=True, - shared_basepath=None) -> npt.NDArray: + shared_basepath=None) -> np.NDArray: """ Energy minimization, Equilibration and Production MD to be reused @@ -252,8 +253,8 @@ def _run_MD(simulation: openmm.app.Simulation, An OpenMM simulation to simulate. positions : openmm.unit.Quantity Initial positions for the system. - sim_settings : SimulationSettingsMD - Settings for MD simulation + settings : PlainMDProtocolSettings + Settings for Plain MD protocol Returns ------- @@ -384,6 +385,8 @@ def _run_MD(simulation: openmm.app.Simulation, if verbose: logger.info(f"Completed simulation in {t1 - t0} seconds") + return + def run(self, *, dry=False, verbose=True, scratch_basepath=None, shared_basepath=None) -> dict[str, Any]: @@ -416,8 +419,6 @@ def run(self, *, dry=False, verbose=True, """ if verbose: self.logger.info("Creating system") - if scratch_basepath is None: - scratch_basepath = pathlib.Path('.') if shared_basepath is None: # use cwd shared_basepath = pathlib.Path('.') From 40978e49a023bca6205ff706dfc90168e3011216 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 12:59:40 +0100 Subject: [PATCH 40/66] Add benzene_modifications --- openfe/tests/protocols/conftest.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py index d340fe9ba..7ae6522ee 100644 --- a/openfe/tests/protocols/conftest.py +++ b/openfe/tests/protocols/conftest.py @@ -6,7 +6,17 @@ from rdkit.Geometry import Point3D import openfe from openff.units import unit +from gufe import SmallMoleculeComponent +@pytest.fixture +def benzene_modifications(): + files = {} + with importlib.resources.files('openfe.tests.data') as d: + fn = str(d / 'benzene_modifications.sdf') + supp = Chem.SDMolSupplier(str(fn), removeHs=False) + for rdmol in supp: + files[rdmol.GetProp('_Name')] = SmallMoleculeComponent(rdmol) + return files @pytest.fixture def benzene_vacuum_system(benzene_modifications): From 150730530bc4c2a848a6c87ac846849141a0df74 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 14:10:29 +0100 Subject: [PATCH 41/66] Remove benzene_modifications --- openfe/protocols/openmm_md/plain_md_methods.py | 2 +- openfe/tests/protocols/conftest.py | 10 ---------- 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 2e0f18c50..f59d19ec4 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -74,7 +74,6 @@ def get_uncertainty(self): def get_traj_filename(self): """String of trajectory file name""" - print(self.data.values()) traj = [pus[0].outputs['nc'] for pus in self.data.values()] return traj @@ -378,6 +377,7 @@ def _run_MD(simulation: openmm.app.Simulation, temperature=True, volume=True, density=True, + speed=True, )) t0 = time.time() simulation.step(prod_steps) diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py index 7ae6522ee..d340fe9ba 100644 --- a/openfe/tests/protocols/conftest.py +++ b/openfe/tests/protocols/conftest.py @@ -6,17 +6,7 @@ from rdkit.Geometry import Point3D import openfe from openff.units import unit -from gufe import SmallMoleculeComponent -@pytest.fixture -def benzene_modifications(): - files = {} - with importlib.resources.files('openfe.tests.data') as d: - fn = str(d / 'benzene_modifications.sdf') - supp = Chem.SDMolSupplier(str(fn), removeHs=False) - for rdmol in supp: - files[rdmol.GetProp('_Name')] = SmallMoleculeComponent(rdmol) - return files @pytest.fixture def benzene_vacuum_system(benzene_modifications): From 8cffa42ab2c3b96fb8dae6d2b86946162f3fef53 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 14:38:05 +0100 Subject: [PATCH 42/66] Change to smc_components --- openfe/protocols/openmm_md/plain_md_methods.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index f59d19ec4..1731dec41 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -499,7 +499,7 @@ def run(self, *, dry=False, verbose=True, stateA_modeller, comp_resids = system_creation.get_omm_modeller( protein_comp=protein_comp, solvent_comp=solvent_comp, - small_mols=small_mols, + small_mols=smc_components, omm_forcefield=system_generator.forcefield, solvent_settings=solvation_settings, ) From c119917b11a36efbae1e880eb437d17b54f84753 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 14:45:59 +0100 Subject: [PATCH 43/66] Move repeat_settings to protocol --- openfe/protocols/openmm_md/plain_md_settings.py | 12 +++++++++++- openfe/protocols/openmm_utils/omm_settings.py | 9 --------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_settings.py b/openfe/protocols/openmm_md/plain_md_settings.py index d17edbe2c..09b4b3b4e 100644 --- a/openfe/protocols/openmm_md/plain_md_settings.py +++ b/openfe/protocols/openmm_md/plain_md_settings.py @@ -10,8 +10,18 @@ from openfe.protocols.openmm_utils.omm_settings import ( Settings, SystemSettings, SolvationSettings, OpenMMEngineSettings, SimulationSettingsMD, - IntegratorSettings, RepeatSettings + IntegratorSettings ) +from gufe.settings import SettingsBaseModel + + +class RepeatSettings(SettingsBaseModel): + """Settings for how many independent MD runs to perform.""" + + n_repeats: int = 1 + """ + Number of independent repeats to run. Default 1 + """ class PlainMDProtocolSettings(Settings): diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 630e1e220..1acbd83b6 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -372,15 +372,6 @@ def must_be_positive(cls, v): return v -class RepeatSettings(SettingsBaseModel): - """Settings for how many independent MD runs to perform.""" - - n_repeats: int = 1 - """ - Number of independent repeats to run. Default 1 - """ - - class SimulationSettingsMD(SimulationSettings): """ Settings for simulation control for plain MD simulations, including From 11ff8cab5d0564502e454b3a4a355a4dc7b9ef0e Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 14:53:13 +0100 Subject: [PATCH 44/66] Change from equ_x to equil_x --- .../protocols/openmm_md/plain_md_methods.py | 132 +----------------- openfe/protocols/openmm_utils/omm_settings.py | 4 +- 2 files changed, 8 insertions(+), 128 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 1731dec41..8a821af85 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -319,7 +319,7 @@ def _run_MD(simulation: openmm.app.Simulation, mdtraj_top.subset(selection_indices), ) traj.save_pdb( - shared_basepath / settings.simulation_settings.equ_NVT_structure + shared_basepath / settings.simulation_settings.equil_NVT_structure ) # NPT equilibration @@ -351,7 +351,7 @@ def _run_MD(simulation: openmm.app.Simulation, mdtraj_top.subset(selection_indices), ) traj.save_pdb( - shared_basepath / settings.simulation_settings.equ_NPT_structure + shared_basepath / settings.simulation_settings.equil_NPT_structure ) # production @@ -551,128 +551,8 @@ def run(self, *, dry=False, verbose=True, equil_steps_nvt, equil_steps_npt, prod_steps, + shared_basepath=shared_basepath, ) - # simulation.context.setPositions(stateA_positions) - # - # # minimize - # if verbose: - # self.logger.info("minimizing systems") - # - # simulation.minimizeEnergy( - # maxIterations=sim_settings.minimization_steps - # ) - # - # # Get the sub selection of the system to save coords for - # selection_indices = mdtraj.Topology.from_openmm( - # simulation.topology).select(sim_settings.output_indices) - # - # positions = to_openmm(from_openmm( - # simulation.context.getState(getPositions=True, - # enforcePeriodicBox=False - # ).getPositions())) - # # Store subset of atoms, specified in input, as PDB file - # mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology) - # - # traj = mdtraj.Trajectory( - # positions[selection_indices, :], - # mdtraj_top.subset(selection_indices), - # ) - # traj.save_pdb( - # shared_basepath / sim_settings.minimized_structure - # ) - # # equilibrate - # # NVT equilibration - # if verbose: - # self.logger.info("Running NVT equilibration") - # - # # Set barostat frequency to zero for NVT - # for x in simulation.context.getSystem().getForces(): - # if x.getName() == 'MonteCarloBarostat': - # x.setFrequency(0) - # - # simulation.context.setVelocitiesToTemperature( - # to_openmm(thermo_settings.temperature)) - # - # t0 = time.time() - # simulation.step(equil_steps_nvt) - # t1 = time.time() - # if verbose: - # self.logger.info(f"Completed NVT equilibration in {t1 - t0} seconds") - # - # # Save last frame NVT equilibration - # positions = to_openmm( - # from_openmm(simulation.context.getState( - # getPositions=True, enforcePeriodicBox=False - # ).getPositions())) - # - # traj = mdtraj.Trajectory( - # positions[selection_indices, :], - # mdtraj_top.subset(selection_indices), - # ) - # traj.save_pdb( - # shared_basepath / sim_settings.equ_NVT_structure - # ) - # - # # NPT equilibration - # if verbose: - # self.logger.info("Running NPT equilibration") - # simulation.context.setVelocitiesToTemperature( - # to_openmm(thermo_settings.temperature)) - # - # # Enable the barostat for NPT - # for x in simulation.context.getSystem().getForces(): - # if x.getName() == 'MonteCarloBarostat': - # x.setFrequency(integrator_settings.barostat_frequency.m) - # - # t0 = time.time() - # simulation.step(equil_steps_npt) - # t1 = time.time() - # if verbose: - # self.logger.info(f"Completed NPT equilibration in {t1 - t0} seconds") - # - # # Save last frame NPT equilibration - # positions = to_openmm( - # from_openmm(simulation.context.getState( - # getPositions=True, enforcePeriodicBox=False - # ).getPositions())) - # - # traj = mdtraj.Trajectory( - # positions[selection_indices, :], - # mdtraj_top.subset(selection_indices), - # ) - # traj.save_pdb( - # shared_basepath / sim_settings.equ_NPT_structure - # ) - # - # # production - # if verbose: - # self.logger.info("running production phase") - # - # # Setup the reporters - # simulation.reporters.append(XTCReporter( - # file=str(shared_basepath / sim_settings.output_filename), - # reportInterval=sim_settings.trajectory_interval.m, - # atomSubset=selection_indices)) - # simulation.reporters.append(openmm.app.CheckpointReporter( - # file=str(shared_basepath / sim_settings.checkpoint_storage), - # reportInterval=sim_settings.checkpoint_interval.m)) - # simulation.reporters.append(openmm.app.StateDataReporter( - # str(shared_basepath / sim_settings.log_output), - # sim_settings.checkpoint_interval.m, - # step=True, - # time=True, - # potentialEnergy=True, - # kineticEnergy=True, - # totalEnergy=True, - # temperature=True, - # volume=True, - # density=True, - # )) - # t0 = time.time() - # simulation.step(prod_steps) - # t1 = time.time() - # if verbose: - # self.logger.info(f"Completed simulation in {t1 - t0} seconds") finally: @@ -682,9 +562,9 @@ def run(self, *, dry=False, verbose=True, if not dry: # pragma: no-cover return { 'system_pdb': shared_basepath / sim_settings.output_filename, - 'minimizes_pdb': shared_basepath / sim_settings.minimized_structure, - 'nvt_equil_pdb': shared_basepath / sim_settings.equ_NVT_structure, - 'npt_equil_pdb': shared_basepath / sim_settings.equ_NPT_structure, + 'minimized_pdb': shared_basepath / sim_settings.minimized_structure, + 'nvt_equil_pdb': shared_basepath / sim_settings.equil_NVT_structure, + 'npt_equil_pdb': shared_basepath / sim_settings.equil_NPT_structure, 'nc': shared_basepath / sim_settings.output_filename, 'last_checkpoint': shared_basepath / sim_settings.checkpoint_storage, } diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 1acbd83b6..76f90fbe8 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -400,10 +400,10 @@ class Config: minimized_structure = 'minimized.pdb' """Path to the pdb file of the system after minimization. Only the specified atom subset is saved. Default 'minimized.pdb'.""" - equ_NVT_structure = 'equ_NVT.pdb' + equil_NVT_structure = 'equil_NVT.pdb' """Path to the pdb file of the system after NVT equilibration. Only the specified atom subset is saved. Default 'equ_NVT.pdb'.""" - equ_NPT_structure = 'equ_NPT.pdb' + equil_NPT_structure = 'equil_NPT.pdb' """Path to the pdb file of the system after NPT equilibration. Only the specified atom subset is saved. Default 'equ_NPT.pdb'.""" checkpoint_storage = 'checkpoint.chk' From 9473a9673bda3f236ef7ea8fd03f5dff7a573108 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 14:57:12 +0100 Subject: [PATCH 45/66] Change comp_type --- openfe/protocols/openmm_md/plain_md_methods.py | 2 +- openfe/protocols/openmm_utils/omm_settings.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 8a821af85..20a101f2d 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -146,7 +146,7 @@ def _create( for comp in [protein_comp] + small_mols: if comp is not None: - comp_type = comp.key.split('-')[0] + comp_type = comp.__class__.__name__ if len(comp.name) == 0: comp_name = 'NoName' else: diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 76f90fbe8..152fca2c4 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -402,10 +402,10 @@ class Config: Only the specified atom subset is saved. Default 'minimized.pdb'.""" equil_NVT_structure = 'equil_NVT.pdb' """Path to the pdb file of the system after NVT equilibration. - Only the specified atom subset is saved. Default 'equ_NVT.pdb'.""" + Only the specified atom subset is saved. Default 'equil_NVT.pdb'.""" equil_NPT_structure = 'equil_NPT.pdb' """Path to the pdb file of the system after NPT equilibration. - Only the specified atom subset is saved. Default 'equ_NPT.pdb'.""" + Only the specified atom subset is saved. Default 'equil_NPT.pdb'.""" checkpoint_storage = 'checkpoint.chk' """ Separate filename for the checkpoint file. Note, this should From b63ae2a7e1a89d463b64b0d9fdb1e1b6888cdbf0 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 15:00:27 +0100 Subject: [PATCH 46/66] Remove stateB from ProtocolUnit --- openfe/protocols/openmm_md/plain_md_methods.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 20a101f2d..7dbfc3fb6 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -156,7 +156,7 @@ def _create( # our DAG has no dependencies, so just list units n_repeats = self._settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( - stateA=stateA, stateB=stateB, + stateA=stateA, settings=self._settings, generation=0, repeat_id=int(uuid.uuid4()), name=f'{system_name} repeat {i} generation 0') @@ -195,7 +195,6 @@ class PlainMDProtocolUnit(gufe.ProtocolUnit): def __init__(self, *, stateA: ChemicalSystem, - stateB: ChemicalSystem, settings: PlainMDProtocolSettings, generation: int, repeat_id: int, @@ -204,9 +203,8 @@ def __init__(self, *, """ Parameters ---------- - stateA, stateB : ChemicalSystem - the two ligand SmallMoleculeComponents to transform between. The - transformation will go from ligandA to ligandB. + stateA : ChemicalSystem + the chemical system for the MD simulation settings : settings.Settings the settings for the Method. This can be constructed using the get_default_settings classmethod to give a starting point that @@ -226,7 +224,6 @@ def __init__(self, *, super().__init__( name=name, stateA=stateA, - stateB=stateB, settings=settings, repeat_id=repeat_id, generation=generation From db4a847abf1ddec8f604dfd15e1ee4e19e72dc35 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 15:12:53 +0100 Subject: [PATCH 47/66] Use individual settings for _run_MD function --- .../protocols/openmm_md/plain_md_methods.py | 58 +++++++++++++------ 1 file changed, 39 insertions(+), 19 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 7dbfc3fb6..bfaefd9f4 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -36,7 +36,7 @@ PlainMDProtocolSettings, SystemSettings, SolvationSettings, OpenMMEngineSettings, IntegratorSettings, SimulationSettingsMD, - RepeatSettings + RepeatSettings, ) from openff.toolkit.topology import Molecule as OFFMolecule @@ -232,7 +232,9 @@ def __init__(self, *, @staticmethod def _run_MD(simulation: openmm.app.Simulation, positions: omm_unit.Quantity, - settings: PlainMDProtocolSettings, + simulation_settings: SimulationSettingsMD, + thermo_settings: settings.ThermoSettings, + integrator_settings: IntegratorSettings, equil_steps_nvt: int, equil_steps_npt: int, prod_steps: int, @@ -249,8 +251,24 @@ def _run_MD(simulation: openmm.app.Simulation, An OpenMM simulation to simulate. positions : openmm.unit.Quantity Initial positions for the system. - settings : PlainMDProtocolSettings - Settings for Plain MD protocol + simulation_settings : SimulationSettingsMD + Settings for MD simulation + thermo_settings: settings.ThermoSettings + temperature settings + integrator_settings: IntegratorSettings + Settings for the MD integrator + equil_steps_nvt: int + number of steps for NVT equilibration + equil_steps_npt: int + number of steps for NPT equilibration + prod_steps: int + number of steps for the production run + verbose: bool + Verbose output of the simulation progress. Output is provided via + INFO level logging. + shared_basepath : Pathlike, optional + Where to run the calculation, defaults to current working directory + Returns ------- @@ -263,12 +281,12 @@ def _run_MD(simulation: openmm.app.Simulation, logger.info("minimizing systems") simulation.minimizeEnergy( - maxIterations=settings.simulation_settings.minimization_steps + maxIterations=simulation_settings.minimization_steps ) # Get the sub selection of the system to save coords for selection_indices = mdtraj.Topology.from_openmm( - simulation.topology).select(settings.simulation_settings.output_indices) + simulation.topology).select(simulation_settings.output_indices) positions = to_openmm(from_openmm( simulation.context.getState(getPositions=True, @@ -282,7 +300,7 @@ def _run_MD(simulation: openmm.app.Simulation, ) traj.save_pdb( - shared_basepath / settings.simulation_settings.minimized_structure + shared_basepath / simulation_settings.minimized_structure ) # equilibrate # NVT equilibration @@ -296,7 +314,7 @@ def _run_MD(simulation: openmm.app.Simulation, x.setFrequency(0) simulation.context.setVelocitiesToTemperature( - to_openmm(settings.thermo_settings.temperature)) + to_openmm(thermo_settings.temperature)) t0 = time.time() simulation.step(equil_steps_nvt) @@ -316,19 +334,19 @@ def _run_MD(simulation: openmm.app.Simulation, mdtraj_top.subset(selection_indices), ) traj.save_pdb( - shared_basepath / settings.simulation_settings.equil_NVT_structure + shared_basepath / simulation_settings.equil_NVT_structure ) # NPT equilibration if verbose: logger.info("Running NPT equilibration") simulation.context.setVelocitiesToTemperature( - to_openmm(settings.thermo_settings.temperature)) + to_openmm(thermo_settings.temperature)) # Enable the barostat for NPT for x in simulation.context.getSystem().getForces(): if x.getName() == 'MonteCarloBarostat': - x.setFrequency(settings.integrator_settings.barostat_frequency.m) + x.setFrequency(integrator_settings.barostat_frequency.m) t0 = time.time() simulation.step(equil_steps_npt) @@ -348,7 +366,7 @@ def _run_MD(simulation: openmm.app.Simulation, mdtraj_top.subset(selection_indices), ) traj.save_pdb( - shared_basepath / settings.simulation_settings.equil_NPT_structure + shared_basepath / simulation_settings.equil_NPT_structure ) # production @@ -357,15 +375,15 @@ def _run_MD(simulation: openmm.app.Simulation, # Setup the reporters simulation.reporters.append(XTCReporter( - file=str(shared_basepath / settings.simulation_settings.output_filename), - reportInterval=settings.simulation_settings.trajectory_interval.m, + file=str(shared_basepath / simulation_settings.output_filename), + reportInterval=simulation_settings.trajectory_interval.m, atomSubset=selection_indices)) simulation.reporters.append(openmm.app.CheckpointReporter( - file=str(shared_basepath / settings.simulation_settings.checkpoint_storage), - reportInterval=settings.simulation_settings.checkpoint_interval.m)) + file=str(shared_basepath / simulation_settings.checkpoint_storage), + reportInterval=simulation_settings.checkpoint_interval.m)) simulation.reporters.append(openmm.app.StateDataReporter( - str(shared_basepath / settings.simulation_settings.log_output), - settings.simulation_settings.checkpoint_interval.m, + str(shared_basepath / simulation_settings.log_output), + simulation_settings.checkpoint_interval.m, step=True, time=True, potentialEnergy=True, @@ -544,7 +562,9 @@ def run(self, *, dry=False, verbose=True, if not dry: # pragma: no-cover self._run_MD(simulation, stateA_positions, - protocol_settings, + sim_settings, + thermo_settings, + integrator_settings, equil_steps_nvt, equil_steps_npt, prod_steps, From 0d527e004586f955f2121115929963314142aaf0 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 15:14:55 +0100 Subject: [PATCH 48/66] Small fix --- openfe/protocols/openmm_md/plain_md_methods.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index bfaefd9f4..b58a69405 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -132,7 +132,7 @@ def _create( raise NotImplementedError("Can't extend simulations yet") # Validate solvent component - nonbond = self._settings.system_settings.nonbonded_method + nonbond = self.settings.system_settings.nonbonded_method system_validation.validate_solvent(stateA, nonbond) # Validate protein component @@ -154,10 +154,10 @@ def _create( system_name += f" {comp_type}:{comp_name}" # our DAG has no dependencies, so just list units - n_repeats = self._settings.repeat_settings.n_repeats + n_repeats = self.settings.repeat_settings.n_repeats units = [PlainMDProtocolUnit( stateA=stateA, - settings=self._settings, + settings=self.settings, generation=0, repeat_id=int(uuid.uuid4()), name=f'{system_name} repeat {i} generation 0') for i in range(n_repeats)] From 98aa2ba54004440a29ec2c3b4ec2cda6dbdbaf0e Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 16:40:26 +0100 Subject: [PATCH 49/66] Fix npt issue --- openfe/protocols/openmm_md/plain_md_methods.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index b58a69405..a095765eb 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -25,6 +25,7 @@ import uuid import time import numpy as np +import numpy.typing as npt import mdtraj from mdtraj.reporters import XTCReporter from openfe.utils import without_oechem_backend, log_system_probe @@ -239,7 +240,7 @@ def _run_MD(simulation: openmm.app.Simulation, equil_steps_npt: int, prod_steps: int, verbose=True, - shared_basepath=None) -> np.NDArray: + shared_basepath=None) -> npt.NDArray: """ Energy minimization, Equilibration and Production MD to be reused From 8aa64a640e1f550b922383304f2218903100c9ee Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 16:44:45 +0100 Subject: [PATCH 50/66] small fix --- openfe/protocols/openmm_md/plain_md_methods.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index a095765eb..5b06ae021 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -401,7 +401,7 @@ def _run_MD(simulation: openmm.app.Simulation, if verbose: logger.info(f"Completed simulation in {t1 - t0} seconds") - return + return None def run(self, *, dry=False, verbose=True, scratch_basepath=None, From 5b9b57262d61bb0fcfaf605f8e1ca9578119d69f Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 16:48:31 +0100 Subject: [PATCH 51/66] small fix2 --- openfe/protocols/openmm_md/plain_md_methods.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 5b06ae021..3b80efd9c 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -25,7 +25,6 @@ import uuid import time import numpy as np -import numpy.typing as npt import mdtraj from mdtraj.reporters import XTCReporter from openfe.utils import without_oechem_backend, log_system_probe @@ -240,7 +239,7 @@ def _run_MD(simulation: openmm.app.Simulation, equil_steps_npt: int, prod_steps: int, verbose=True, - shared_basepath=None) -> npt.NDArray: + shared_basepath=None): """ Energy minimization, Equilibration and Production MD to be reused From 3996e3742bb7b83fb512b21ed99fd347aaad3edc Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 23 Nov 2023 11:31:05 +0100 Subject: [PATCH 52/66] Add test for result json file --- openfe/tests/protocols/conftest.py | 12 ++++ .../test_openmm_plain_md_protocols.py | 56 ++++++++++++++++--- 2 files changed, 60 insertions(+), 8 deletions(-) diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py index 8f0a79e6d..45b25723e 100644 --- a/openfe/tests/protocols/conftest.py +++ b/openfe/tests/protocols/conftest.py @@ -212,3 +212,15 @@ def afe_solv_transformation_json() -> str: with gzip.open((d / fname).as_posix(), 'r') as f: # type: ignore return f.read().decode() # type: ignore + + +@pytest.fixture +def md_json() -> str: + """ + string of a MD result (TYK ligand lig_ejm_31 in water) generated by quickrun + """ + d = resources.files('openfe.tests.data.openmm_md') + fname = "md_results.json.gz" + + with gzip.open((d / fname).as_posix(), 'r') as f: # type: ignore + return f.read().decode() # type: ignore diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index ba3df33dc..ff2e31a37 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -12,6 +12,10 @@ from openfe.protocols.openmm_md.plain_md_methods import ( PlainMDProtocol, PlainMDProtocolUnit, PlainMDProtocolResult, ) +import json +import openfe +from openfe.protocols import openmm_md +import pathlib def test_create_default_settings(): @@ -90,12 +94,12 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): with tmpdir.as_cwd(): sim = dag_unit.run(dry=True)['debug']['system'] - print(sim) assert not ThermodynamicState(sim, temperature=to_openmm( protocol.settings.thermo_settings.temperature)).is_periodic assert ThermodynamicState(sim, temperature=to_openmm( protocol.settings.thermo_settings.temperature)).barostat is None + def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir): vac_settings = PlainMDProtocol.default_settings() @@ -105,6 +109,7 @@ def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir): protocol = PlainMDProtocol( settings=vac_settings, ) + assert protocol.settings.simulation_settings.forcefield_cache is None # create DAG from protocol and take first (and only) work unit from within dag = protocol.create( @@ -115,12 +120,7 @@ def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir): dag_unit = list(dag.protocol_units)[0] with tmpdir.as_cwd(): - sim = dag_unit.run(dry=True)['debug']['system'] - print(sim) - assert not ThermodynamicState(sim, temperature=to_openmm( - protocol.settings.thermo_settings.temperature)).is_periodic - assert ThermodynamicState(sim, temperature=to_openmm( - protocol.settings.thermo_settings.temperature)).barostat is None + dag_unit.run(dry=True)['debug']['system'] def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir): @@ -377,4 +377,44 @@ def test_gather(solvent_protocol_dag, tmpdir): res = prot.gather([dagres]) - assert isinstance(res, PlainMDProtocolResult) \ No newline at end of file + assert isinstance(res, PlainMDProtocolResult) + + +class TestProtocolResult: + @pytest.fixture() + def protocolresult(self, md_json): + d = json.loads(md_json, cls=gufe.tokenization.JSON_HANDLER.decoder) + + pr = openfe.ProtocolResult.from_dict(d['protocol_result']) + + return pr + + def test_reload_protocol_result(self, md_json): + d = json.loads(md_json, cls=gufe.tokenization.JSON_HANDLER.decoder) + + pr = openmm_md.plain_md_methods.PlainMDProtocolResult.from_dict(d['protocol_result']) + + assert pr + + def test_get_estimate(self, protocolresult): + est = protocolresult.get_estimate() + + assert est is None + + + def test_get_uncertainty(self, protocolresult): + est = protocolresult.get_uncertainty() + + assert est is None + + def test_get_traj_filename(self, protocolresult): + traj = protocolresult.get_traj_filename() + + assert isinstance(traj, list) + assert isinstance(traj[0], pathlib.PurePath) + + def test_get_pdb_filename(self, protocolresult): + pdb = protocolresult.get_pdb_filename() + + assert isinstance(pdb, list) + assert isinstance(pdb[0], pathlib.PurePath) From 27b9b42a1a2565e225d8f568f9ef206ef7888a44 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 23 Nov 2023 11:38:03 +0100 Subject: [PATCH 53/66] Add results json files for MD run --- openfe/tests/data/openmm_md/md_results.json.gz | Bin 0 -> 3009 bytes 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 openfe/tests/data/openmm_md/md_results.json.gz diff --git a/openfe/tests/data/openmm_md/md_results.json.gz b/openfe/tests/data/openmm_md/md_results.json.gz new file mode 100644 index 0000000000000000000000000000000000000000..7ee1acdb3ded2b83763fd101d59873a3d8ad3ad8 GIT binary patch literal 3009 zcmV;y3qJH8iwFoK3|?gb18rnqa%FRMY;BNa>uN*{p;z>hNiB^NrQg#k=Rbgo2X2y*I17_72t&uQ zE$aEMZw9gBNi9h%r+4+dP1sUMA*?>WtIIl|#ATu9F~&UG5h6ko5*nv8B#y8zaix^J zYUaEHS9DzX@+-Ja7KY%6u!dGC+)`eblKAYIxj!QCm(ikQi-(KL=li2)-r|wB^Z(7E ztF$-$KVsU3vEYiat7LvqSNiW;;8ohcv(&r~t0Zec68H87(m# zCW>C8*#nw<4u;;5#G#o$7+NKX*1#@cT64G{S_9jF$saQMn-mg6Yv5yG8Z0e&EOP^G z;0gI`*{e&=D8ba-UgZ;Cke@Hom!r4+H>u>}nJupY zJLNa!jp@S`hza@EGx-^VvT-J_ECEXx=fx)|h=FMM$ka);I-TgcL zR2ExG`U_xZFe{)a10Ce`lQ-BLWe}{|0)$e1LkP<$?xfzIwtqXXt#+pYMx)Busx(5<7l!dF1rdK9_rBVw$ z0`o{8ER@P%a)$8&JY=St)hZ)f-IXxMLFIrC^f9R>5i=y#+y!z2>-$Nj*m&=rkp1!* zP~*d12Aaeis>7%c9L5}3(T6d{NF`@)Z~#gYM^n%GY4bEjlmbc};39pnV*1m>Kn`j> zC4cQLGfgD_h&X})bST6m4I&lB6~Ll0U{#86VZD}tefSV0c$n;gH?$k=Lk9<-VBjKx zZ4ERjjK{bn+?(jZ^ub8;Ndl(F@g8_9(;T2yoK1}4^&5ogZ4Hyi3T9Ho08moEa)?T2 z8o#rqE^25dN*yd_m~t7+3=t7bM`ITDin6eeB+c|x_+|s|Z@Nl9r+-5XKZcJ~m5^pb zTEDu1WAG7^b_S}M(6G?8G_F2d zdU!JbKof%{JBBh82O6&2o5=NGjJ^d+vh?gPNIr(g<(_hP0ip>|&|#2FnR^Q{rQ}U2 z00#`wPkT);ddo_N)_tkG@kr_a6DE+OO#H(A1^DQLhe}ee`7*%tcJ8Ts`y4zHKkbo7 zzSz9%FbIK#yhQM=c0wTSC{YG3{GD*I>g2^_^DR?h|)} z2x@2FYoJzR=;JlKWkkaxR{79GL}%>I8zuhhf539MSpps@9!Ka(a-+@um+%1M8D4$m zuSvWP$WWa$wwJBP`XXD4dShE6b@ zFMx#mD2R&Pea*g2X?FU__6$OydQfmLUuJ(h=ASVgo!0o9ikL^S`8AI+;GKI#r-%WkClzU`e;*bbgE1n z>P5*Y>!NAv2K#4q#ts&}!@e_UT*@Y3^lDBEIT%IITHM%Agmz(6)~O&O-*ld_ao?sJ zcXD2-<}K;4M2FozaQ#p!v+ zS-EWm#nd4a_=<6!S5pjI6eR`iI-D_Ru%0d$*F@J1tAy}H*SQ-@LRm&-I&P#4aT9gg z?4!Jkl6lAL@~HQjcqvX;ma{_RT~q0?5PW=$CVqp><+{Cpf4gd6^+MiQjId(63L?^- zh-HpVU}J`jCo$0G(>#rSqa0rN6?eW)&A097U> zBuTW$lBic?8Mo}f8j~>rEbtj?%lZocKLLjK0J#Z3zak!(34cB_=RB7WW%>JgS=E)y z1C~Z5voSQ$#@yw0*UubP1Vg&Qkc)j|eq-^qE4`VQszD%UJ{uGKrLPDcQB zDHD>^&_)`Xrb5I9-&oh;8&*y8iv4a4y$5X9>KnOv%{hz29>8U{R0FoKg-xM&U!mxO zCDXO4$@7wlENa#w795br)T~lAES>&ArZ;RjLUX?TE=y25J$$MH_maL@4uV6`ir7PQ;=j2Ye@+`A(L6MLmrWHUrR!0aYaBZ)L71C_ezhHNSkxVl0<| zjJeeG%H?_Q5B_3E%<6yeBhS>A%bY>G-iOToR4hHb{?001?_Zx#QIggi)d;Ep|8^u% zC}hUr(~h;>;fz)loyBxsikn09^5TOFOm>G$mmhXBDh~0{+x#MFvHL~ja8Qy~Vf)cp zyPA4%IyvJJCeTM`>}vf0W1ElN*b~jzAD6v@lLVfI;p{H==XG0C zaU_C~HBuhXe zum7fXOM9%}sTX%}*tvYZQR-^owh8_thX8Ef;BiPDn4~b=n&0LZxpLnB&F_PcWb6)v zXbF}qVKM)s}+G!AH4h@|+Nq3F*U88;1Xx}y3ca8R4qkY$CUyIPL z(Y|Z6?;7p9M*FVOzH7Aa8tuDA`>xTxYqakg?Yl<%uF<}0v|mG`{qpj^0rw|SNkaet D98}r6 literal 0 HcmV?d00001 From 2d6e38eb174a534a41f9a074a576d1b222ea6556 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 23 Nov 2023 13:49:32 +0100 Subject: [PATCH 54/66] Add tests for logging output --- .../test_openmm_plain_md_protocols.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index ff2e31a37..0cc724ebf 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -16,6 +16,7 @@ import openfe from openfe.protocols import openmm_md import pathlib +import logging def test_create_default_settings(): @@ -93,13 +94,44 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir): dag_unit = list(dag.protocol_units)[0] with tmpdir.as_cwd(): - sim = dag_unit.run(dry=True)['debug']['system'] + sim = dag_unit.run(dry=True, verbose=True)['debug']['system'] assert not ThermodynamicState(sim, temperature=to_openmm( protocol.settings.thermo_settings.temperature)).is_periodic assert ThermodynamicState(sim, temperature=to_openmm( protocol.settings.thermo_settings.temperature)).barostat is None +def test_dry_run_logger_output(benzene_vacuum_system, tmpdir, caplog): + + vac_settings = PlainMDProtocol.default_settings() + vac_settings.system_settings.nonbonded_method = 'nocutoff' + vac_settings.simulation_settings.equilibration_length_nvt = 1 * unit.picosecond + vac_settings.simulation_settings.equilibration_length = 1 * unit.picosecond + vac_settings.simulation_settings.production_length = 1 * unit.picosecond + + protocol = PlainMDProtocol( + settings=vac_settings, + ) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + caplog.set_level(logging.INFO) + dag_unit.run(dry=False, verbose=True) + + messages = [r.message for r in caplog.records] + assert "minimizing systems" in messages + assert "Running NVT equilibration" in messages + assert "Running NPT equilibration" in messages + assert "running production phase" in messages + + def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir): vac_settings = PlainMDProtocol.default_settings() From 45890d47e29e44a66f8ad764eaa3b39af2b30c40 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Wed, 22 Nov 2023 14:28:34 +0000 Subject: [PATCH 55/66] Add init files --- openfe/protocols/openmm_md/__init__.py | 1 + openfe/tests/data/openmm_md/__init__.py | 1 + 2 files changed, 2 insertions(+) create mode 100644 openfe/protocols/openmm_md/__init__.py create mode 100644 openfe/tests/data/openmm_md/__init__.py diff --git a/openfe/protocols/openmm_md/__init__.py b/openfe/protocols/openmm_md/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/openfe/protocols/openmm_md/__init__.py @@ -0,0 +1 @@ + diff --git a/openfe/tests/data/openmm_md/__init__.py b/openfe/tests/data/openmm_md/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/openfe/tests/data/openmm_md/__init__.py @@ -0,0 +1 @@ + From f5e6d999e535cac02d4ad3a53c109e55b79c0051 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 24 Nov 2023 09:39:47 +0100 Subject: [PATCH 56/66] Changes __init__.py --- openfe/protocols/openmm_md/__init__.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/openfe/protocols/openmm_md/__init__.py b/openfe/protocols/openmm_md/__init__.py index 8b1378917..c157ecbec 100644 --- a/openfe/protocols/openmm_md/__init__.py +++ b/openfe/protocols/openmm_md/__init__.py @@ -1 +1,20 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Run MD simulation using OpenMM and OpenMMTools. +""" + +from .plain_md_methods import ( + PlainMDProtocol, + PlainMDProtocolSettings, + PlainMDProtocolResult, + PlainMDProtocolUnit, +) + +__all__ = [ + "PlainMDProtocol", + "PlainMDProtocolSettings", + "PlainMDProtocolResult", + "PlainMDProtocolUnit", +] From 146b8d592cb0d26fd3408f1c3a5c1346c6f7654a Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 7 Dec 2023 16:06:00 +0100 Subject: [PATCH 57/66] add return annotation for results --- .../protocols/openmm_md/plain_md_methods.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 3b80efd9c..6b0bede3b 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -72,14 +72,28 @@ def get_uncertainty(self): return None - def get_traj_filename(self): - """String of trajectory file name""" + def get_traj_filename(self) -> list[pathlib.PurePath]: + """ + Get a list of trajectory paths + + Returns + ------- + traj : list + list of paths (pathlib.PurePath) to the simulation trajectory + """ traj = [pus[0].outputs['nc'] for pus in self.data.values()] return traj - def get_pdb_filename(self): - """String of pdb file name""" + def get_pdb_filename(self) -> list[pathlib.PurePath]: + """ + Get a list of paths to the pdb files of the pre-minimized system. + + Returns + ------- + pdbs : list + list of paths (pathlib.PurePath) to the pdb files + """ pdbs = [pus[0].outputs['system_pdb'] for pus in self.data.values()] return pdbs From 71f09d1f9465597384f57d020d4a12ca9f7ab5e0 Mon Sep 17 00:00:00 2001 From: Hannah Baumann <43765638+hannahbaumann@users.noreply.github.com> Date: Thu, 7 Dec 2023 16:07:34 +0100 Subject: [PATCH 58/66] Update openfe/protocols/openmm_md/plain_md_methods.py Co-authored-by: Irfan Alibay --- openfe/protocols/openmm_md/plain_md_methods.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 6b0bede3b..3976ae066 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -204,7 +204,7 @@ def _gather( class PlainMDProtocolUnit(gufe.ProtocolUnit): """ - Base class for plain MD simulations (NonTransformation). + Protocol unit for plain MD simulations (NonTransformation). """ def __init__(self, *, From 085919da9acca522bc84f6e2665f759e2fe3d205 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 7 Dec 2023 16:27:04 +0100 Subject: [PATCH 59/66] make settings names more specific --- openfe/protocols/openmm_md/plain_md_methods.py | 14 +++++++------- openfe/protocols/openmm_utils/omm_settings.py | 8 ++++---- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 3976ae066..1d0d78edc 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -389,11 +389,11 @@ def _run_MD(simulation: openmm.app.Simulation, # Setup the reporters simulation.reporters.append(XTCReporter( - file=str(shared_basepath / simulation_settings.output_filename), - reportInterval=simulation_settings.trajectory_interval.m, + file=str(shared_basepath / simulation_settings.production_trajectory_filename), + reportInterval=simulation_settings.trajectory_write_interval.m, atomSubset=selection_indices)) simulation.reporters.append(openmm.app.CheckpointReporter( - file=str(shared_basepath / simulation_settings.checkpoint_storage), + file=str(shared_basepath / simulation_settings.checkpoint_storage_filename), reportInterval=simulation_settings.checkpoint_interval.m)) simulation.reporters.append(openmm.app.StateDataReporter( str(shared_basepath / simulation_settings.log_output), @@ -547,7 +547,7 @@ def run(self, *, dry=False, verbose=True, ) # f. Save pdb of entire system - with open(shared_basepath / sim_settings.output_structure, "w") as f: + with open(shared_basepath / sim_settings.preminimized_structure, "w") as f: openmm.app.PDBFile.writeFile( stateA_topology, stateA_positions, file=f, keepIds=True ) @@ -592,12 +592,12 @@ def run(self, *, dry=False, verbose=True, if not dry: # pragma: no-cover return { - 'system_pdb': shared_basepath / sim_settings.output_filename, + 'system_pdb': shared_basepath / sim_settings.preminimized_structure, 'minimized_pdb': shared_basepath / sim_settings.minimized_structure, 'nvt_equil_pdb': shared_basepath / sim_settings.equil_NVT_structure, 'npt_equil_pdb': shared_basepath / sim_settings.equil_NPT_structure, - 'nc': shared_basepath / sim_settings.output_filename, - 'last_checkpoint': shared_basepath / sim_settings.checkpoint_storage, + 'nc': shared_basepath / sim_settings.production_trajectory_filename, + 'last_checkpoint': shared_basepath / sim_settings.checkpoint_storage_filename, } else: return {'debug': {'system': stateA_system}} diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 152fca2c4..aea2dd94f 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -389,13 +389,13 @@ class Config: :class:`IntegratorSettings.n_steps`. """ # reporter settings - output_filename = 'simulation.xtc' + production_trajectory_filename = 'simulation.xtc' """Path to the storage file for analysis. Default 'simulation.xtc'.""" - trajectory_interval = 5000 * unit.timestep + trajectory_write_interval = 5000 * unit.timestep """ Frequency to write the xtc file. Default 5000 * unit.timestep. """ - output_structure = 'system.pdb' + preminimized_structure = 'system.pdb' """Path to the pdb file of the full pre-minimized system. Default 'system.pdb'.""" minimized_structure = 'minimized.pdb' """Path to the pdb file of the system after minimization. @@ -406,7 +406,7 @@ class Config: equil_NPT_structure = 'equil_NPT.pdb' """Path to the pdb file of the system after NPT equilibration. Only the specified atom subset is saved. Default 'equil_NPT.pdb'.""" - checkpoint_storage = 'checkpoint.chk' + checkpoint_storage_filename = 'checkpoint.chk' """ Separate filename for the checkpoint file. Note, this should not be a full path, just a filename. Default 'checkpoint.chk'. From f21a7b1c93597ad2e54d8f21dc91d0ac9dc91428 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 8 Dec 2023 09:23:18 +0100 Subject: [PATCH 60/66] change results.json file with new settings names --- .../tests/data/openmm_md/md_results.json.gz | Bin 3009 -> 3055 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/openfe/tests/data/openmm_md/md_results.json.gz b/openfe/tests/data/openmm_md/md_results.json.gz index 7ee1acdb3ded2b83763fd101d59873a3d8ad3ad8..7fe11cf19637c2359048c7b46f37a97348996b47 100644 GIT binary patch literal 3055 zcmVAr!w!K!fSn!8Xj)Rs7U)811wvl-w{Le<+iJ^pxDwvJtn>0Lsk*wlj<2d)WuI!a zVQEen)%FHOn`PSrtt}#2Gm;kUq$5^!$;znAf|@pM#`wUg786F`#bq<*8PXTr8#?mDw>Z*x=yK;3y|qc2H(CYBRcD=2cnn zk-aRP1oU+tSjM2H6(ww-4<@vrHDPI440Mg_6}+KBOYKnV_q<5iQ&4dyw8PMJ{Up&X z-}V#Fv%G}=6VEk_$R9akdQ-M_MCU>>ZuP+(Ro0GAqsX^tV)`Mq{SomjEB5Ws8&O>! znWihGq(#-T23*lF?#tV78_f*C5pE4BVz?!&%mwk~OJjSW!+%1uhRz;N&t7d0UXEsu ztZx4|htAX9a6iSg6=U8xW9P~Irml1^THtxwf3?-0>E-k0og4CRJ@!S8;e%RO7BMCu zry20CPUH?S%teyfBk_A{1*s>Cv)DZ@PO|hxF3#fgxH!pQWOA0(;v{d8iF3G%6d^?i z^%qEen++7bM6-G{@$B`zBZz%7fzY=~5Y3)lz_jFWPBeSA0h8NjbXO_FiDu8oz|>n> z@R-I5TF(>W*|Zj&2sFFLC(LtJ&_*k0YMcZlqq0cKI;XXOr8aJL-M1b+KREoixYZ5# z)Q|CUx)swVs<-wATYpcyt?j`UpUP@d1Z5plLBEHf-)g3$qWt-$CD8f(_TX&rrGMi> zES^~6)1yc1hPW|Y_yl64|Ldi=hM;Vih))`UB^2}?A(s!sHavJ>Q3&I1q9cO?GZXLb z{Tie@NIry>V-V<2Nv?Y=oL(lRONwTbPdgG--3{^ z&p@2*{cb7^O)32~uv3`j5tM-r8g-Lb*z7021DA(fx4^|CG!aCixu5}YcabXqZ^C#4 zQVjnIu))Gch+Hr;ieCMSi^8Yd#VpcuxPWsL5EGiG@L&!Vk5F_YDgg8mNK8-`uA*0+ zG9D}yo9Gdkr+2|ZsT3wB7%#v>V#rywG_qD*33F^z4){PHgD4U)Lqf$}AlEaWj1$Sm zI`{bWkMDsR?shWJBx(~CMqS`A=E#aZ3^7JBIXiIjs?zWA-w?x( z;S$RdQmshqx7TnCE}>peK{X?vK>{oSXT{+ZQCPx7LP-AwbK^r7146g1E0YOI?otR1 z^Bq&+>awMSCv|reF<7!gC_|yA;L5pyTn~omTd*XFkN<+?19+V8DR*Wd8UqC#2Fa8- zHxW}(-nam8z!3Sa(*&b8uViT5PtqHYr2fBQ0@>2UADusdk34uNC8dflJxp)=uFSU| zz$4-F4v8epuPMHen+Gp6HHwicuiYHCrDX64?vl8u@9f`F5(X>l0)oRRQsSgr;1E5$ zaTC#NbYob42?Q!csKTh>ey2!?_#&Ez zt(|Rl$PwFPhQKuuVfyjo4n4*Q_Ej%dQ#}fU85(Km@B1_g7*3FKqw_$_& zcCBppkuyL9)zc>msO1>?cnNQbuHX@?yl)_)6T1IFivRf^upDldfJcJI5weinXtVnh zJb-wLpRV-RI9vv#FHail%hqFck+nrVv1HXPI|O;Ed&q5Nk6X7{W8L#j1<8c``RpkG zipq^oFYE<-Lq~y>_j>7Il*9^ZLu(=34}X<1_pbl*w9}c_5JlNu)x4(0%>Us8wrhfw z%u&YB35NY6kZ>ObQK7S~*taUpj;^oIAQY+w1^4pV#s-@_)(bsR7NS_65jjSGX?$&1 zIKKitA>*+;{ryLU^+OE4fBR3Aq@0a(u$!MSv4qM)ltfRhKmwB+J~35sKlb&<(ner? zv?@0mS*A7lqNL?zR@G&N{j)NmJG0(l*BR8#L=!N0+mcKS1_87dSN0R3T^QtL%!$Z1 zX%jl^+H}K4%**AxIUSa0uv>{ezC6vSP)d6&WL9O9vhzISiEYGCzD-jI*b zY`qFR(hZMgicMf`ij5~eP^XhL4&D(4FZ`0brg_N$!$ZVrvTvnLPvW|qP;N5UnbWPI zzF-^lu{4s0wBAh;!#?_X_&(3X8c1VlWmQ4LND-4dUN!d1y}Lo)W*n>}bNv4d7%s@&JcgBsU|UmIDJPR6Wo;~; zOy3)dzLM5*l92p{X2)r76{&;rjwLM@uV{l=nWWH)wF3=t)~xtZ{(YnIU2df zh;JOevi2+-0bzOBRNM#js}-4c-`3PLF5d@L$#^(ZlF0$cO#u2i@xV;@{LE}wDjxFu z<9S(>g~$UM2RUWa5`c4fa93Lxidt4C$&yzMK>y5*1u04{gVtZ=yjgvRa(td^W;~bb z9a*OXfVz+gNvbI$6^UamVuf$aOYsdYCTT%GT0`#w+ok$OY+fy+0pA0d=$0zL=C;rY z6z_8seK2RbR5e+eQ=UcDQiN!@qa9rBf{)s%-aofN?_&B}IUAIfm z?%)J<_G+cn#lWo-{3!T|EBrBY5s4T|C{Fjrun}Vp-uCD)BN8w|2NJ5 xP4j=#{NFVHH_iV|^MBL)ze2`tn*UeI!cFsk)BL~i=KtB*e*xsprR{M<001w7=_UXG literal 3009 zcmV;y3qJH8iwFoK3|?gb18rnqa%FRMY;BNa>uN*{p;z>hNiB^NrQg#k=Rbgo2X2y*I17_72t&uQ zE$aEMZw9gBNi9h%r+4+dP1sUMA*?>WtIIl|#ATu9F~&UG5h6ko5*nv8B#y8zaix^J zYUaEHS9DzX@+-Ja7KY%6u!dGC+)`eblKAYIxj!QCm(ikQi-(KL=li2)-r|wB^Z(7E ztF$-$KVsU3vEYiat7LvqSNiW;;8ohcv(&r~t0Zec68H87(m# zCW>C8*#nw<4u;;5#G#o$7+NKX*1#@cT64G{S_9jF$saQMn-mg6Yv5yG8Z0e&EOP^G z;0gI`*{e&=D8ba-UgZ;Cke@Hom!r4+H>u>}nJupY zJLNa!jp@S`hza@EGx-^VvT-J_ECEXx=fx)|h=FMM$ka);I-TgcL zR2ExG`U_xZFe{)a10Ce`lQ-BLWe}{|0)$e1LkP<$?xfzIwtqXXt#+pYMx)Busx(5<7l!dF1rdK9_rBVw$ z0`o{8ER@P%a)$8&JY=St)hZ)f-IXxMLFIrC^f9R>5i=y#+y!z2>-$Nj*m&=rkp1!* zP~*d12Aaeis>7%c9L5}3(T6d{NF`@)Z~#gYM^n%GY4bEjlmbc};39pnV*1m>Kn`j> zC4cQLGfgD_h&X})bST6m4I&lB6~Ll0U{#86VZD}tefSV0c$n;gH?$k=Lk9<-VBjKx zZ4ERjjK{bn+?(jZ^ub8;Ndl(F@g8_9(;T2yoK1}4^&5ogZ4Hyi3T9Ho08moEa)?T2 z8o#rqE^25dN*yd_m~t7+3=t7bM`ITDin6eeB+c|x_+|s|Z@Nl9r+-5XKZcJ~m5^pb zTEDu1WAG7^b_S}M(6G?8G_F2d zdU!JbKof%{JBBh82O6&2o5=NGjJ^d+vh?gPNIr(g<(_hP0ip>|&|#2FnR^Q{rQ}U2 z00#`wPkT);ddo_N)_tkG@kr_a6DE+OO#H(A1^DQLhe}ee`7*%tcJ8Ts`y4zHKkbo7 zzSz9%FbIK#yhQM=c0wTSC{YG3{GD*I>g2^_^DR?h|)} z2x@2FYoJzR=;JlKWkkaxR{79GL}%>I8zuhhf539MSpps@9!Ka(a-+@um+%1M8D4$m zuSvWP$WWa$wwJBP`XXD4dShE6b@ zFMx#mD2R&Pea*g2X?FU__6$OydQfmLUuJ(h=ASVgo!0o9ikL^S`8AI+;GKI#r-%WkClzU`e;*bbgE1n z>P5*Y>!NAv2K#4q#ts&}!@e_UT*@Y3^lDBEIT%IITHM%Agmz(6)~O&O-*ld_ao?sJ zcXD2-<}K;4M2FozaQ#p!v+ zS-EWm#nd4a_=<6!S5pjI6eR`iI-D_Ru%0d$*F@J1tAy}H*SQ-@LRm&-I&P#4aT9gg z?4!Jkl6lAL@~HQjcqvX;ma{_RT~q0?5PW=$CVqp><+{Cpf4gd6^+MiQjId(63L?^- zh-HpVU}J`jCo$0G(>#rSqa0rN6?eW)&A097U> zBuTW$lBic?8Mo}f8j~>rEbtj?%lZocKLLjK0J#Z3zak!(34cB_=RB7WW%>JgS=E)y z1C~Z5voSQ$#@yw0*UubP1Vg&Qkc)j|eq-^qE4`VQszD%UJ{uGKrLPDcQB zDHD>^&_)`Xrb5I9-&oh;8&*y8iv4a4y$5X9>KnOv%{hz29>8U{R0FoKg-xM&U!mxO zCDXO4$@7wlENa#w795br)T~lAES>&ArZ;RjLUX?TE=y25J$$MH_maL@4uV6`ir7PQ;=j2Ye@+`A(L6MLmrWHUrR!0aYaBZ)L71C_ezhHNSkxVl0<| zjJeeG%H?_Q5B_3E%<6yeBhS>A%bY>G-iOToR4hHb{?001?_Zx#QIggi)d;Ep|8^u% zC}hUr(~h;>;fz)loyBxsikn09^5TOFOm>G$mmhXBDh~0{+x#MFvHL~ja8Qy~Vf)cp zyPA4%IyvJJCeTM`>}vf0W1ElN*b~jzAD6v@lLVfI;p{H==XG0C zaU_C~HBuhXe zum7fXOM9%}sTX%}*tvYZQR-^owh8_thX8Ef;BiPDn4~b=n&0LZxpLnB&F_PcWb6)v zXbF}qVKM)s}+G!AH4h@|+Nq3F*U88;1Xx}y3ca8R4qkY$CUyIPL z(Y|Z6?;7p9M*FVOzH7Aa8tuDA`>xTxYqakg?Yl<%uF<}0v|mG`{qpj^0rw|SNkaet D98}r6 From 7f46b7f587511651a8f6212e0c61b67c95866cc8 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 8 Dec 2023 10:23:21 +0100 Subject: [PATCH 61/66] change run method to take temp and barostat_frequency instead of full settings --- .../protocols/openmm_md/plain_md_methods.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 1d0d78edc..27480c661 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -247,8 +247,8 @@ def __init__(self, *, def _run_MD(simulation: openmm.app.Simulation, positions: omm_unit.Quantity, simulation_settings: SimulationSettingsMD, - thermo_settings: settings.ThermoSettings, - integrator_settings: IntegratorSettings, + temperature: settings.ThermoSettings.temperature, + barostat_frequency: IntegratorSettings.barostat_frequency, equil_steps_nvt: int, equil_steps_npt: int, prod_steps: int, @@ -267,10 +267,10 @@ def _run_MD(simulation: openmm.app.Simulation, Initial positions for the system. simulation_settings : SimulationSettingsMD Settings for MD simulation - thermo_settings: settings.ThermoSettings - temperature settings - integrator_settings: IntegratorSettings - Settings for the MD integrator + temperature: settings.ThermoSettings.temperature + temperature setting + barostat_frequency: IntegratorSettings.barostat_frequency + Frequency for the barostat equil_steps_nvt: int number of steps for NVT equilibration equil_steps_npt: int @@ -328,7 +328,7 @@ def _run_MD(simulation: openmm.app.Simulation, x.setFrequency(0) simulation.context.setVelocitiesToTemperature( - to_openmm(thermo_settings.temperature)) + to_openmm(temperature)) t0 = time.time() simulation.step(equil_steps_nvt) @@ -355,12 +355,12 @@ def _run_MD(simulation: openmm.app.Simulation, if verbose: logger.info("Running NPT equilibration") simulation.context.setVelocitiesToTemperature( - to_openmm(thermo_settings.temperature)) + to_openmm(temperature)) # Enable the barostat for NPT for x in simulation.context.getSystem().getForces(): if x.getName() == 'MonteCarloBarostat': - x.setFrequency(integrator_settings.barostat_frequency.m) + x.setFrequency(barostat_frequency.m) t0 = time.time() simulation.step(equil_steps_npt) @@ -561,7 +561,7 @@ def run(self, *, dry=False, verbose=True, integrator = openmm.LangevinMiddleIntegrator( to_openmm(thermo_settings.temperature), to_openmm(integrator_settings.collision_rate), - to_openmm(integrator_settings.timestep), + to_openmm(timestep), ) simulation = openmm.app.Simulation( @@ -577,8 +577,8 @@ def run(self, *, dry=False, verbose=True, self._run_MD(simulation, stateA_positions, sim_settings, - thermo_settings, - integrator_settings, + thermo_settings.temperature, + integrator_settings.barostat_frequency, equil_steps_nvt, equil_steps_npt, prod_steps, From f8ec4713c20c4fe93221eb667a0320cb7f2d140a Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 8 Dec 2023 10:26:39 +0100 Subject: [PATCH 62/66] add slow test --- .../test_openmm_plain_md_protocols.py | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index 0cc724ebf..2bd832853 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -297,30 +297,30 @@ def test_dry_run_ligand_tip4p(benzene_system, tmpdir): assert system -# @pytest.mark.slow -# def test_dry_run_complex(benzene_complex_system, tmpdir): -# # this will be very time consuming -# settings = PlainMDProtocol.default_settings() -# -# protocol = PlainMDProtocol( -# settings=settings, -# ) -# dag = protocol.create( -# stateA=benzene_complex_system, -# stateB=benzene_complex_system, -# mapping=None, -# ) -# dag_unit = list(dag.protocol_units)[0] -# -# with tmpdir.as_cwd(): -# sim = dag_unit.run(dry=True)['debug']['system'] -# assert not ThermodynamicState(sim, temperature= -# to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic -# assert isinstance(ThermodynamicState(sim, temperature= -# to_openmm(protocol.settings.thermo_settings.temperature)).barostat, -# MonteCarloBarostat) -# assert ThermodynamicState(sim, temperature= -# to_openmm(protocol.settings.thermo_settings.temperature)).pressure == 1 * omm_unit.bar +@pytest.mark.slow +def test_dry_run_complex(benzene_complex_system, tmpdir): + # this will be very time consuming + settings = PlainMDProtocol.default_settings() + + protocol = PlainMDProtocol( + settings=settings, + ) + dag = protocol.create( + stateA=benzene_complex_system, + stateB=benzene_complex_system, + mapping=None, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sim = dag_unit.run(dry=True)['debug']['system'] + assert not ThermodynamicState(sim, temperature= + to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic + assert isinstance(ThermodynamicState(sim, temperature= + to_openmm(protocol.settings.thermo_settings.temperature)).barostat, + MonteCarloBarostat) + assert ThermodynamicState(sim, temperature= + to_openmm(protocol.settings.thermo_settings.temperature)).pressure == 1 * omm_unit.bar def test_hightimestep(benzene_vacuum_system, tmpdir): From 39803e50d0ab5c605dd8409009fb8d1dcca71ca4 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 8 Dec 2023 13:05:20 +0100 Subject: [PATCH 63/66] fix test --- openfe/tests/protocols/test_openmm_plain_md_protocols.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index 2bd832853..d29d08790 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -5,10 +5,10 @@ import pytest from unittest import mock from openff.units import unit - +from openmm import unit as omm_unit from openff.units.openmm import to_openmm from openmmtools.states import ThermodynamicState - +from openmm import MonteCarloBarostat from openfe.protocols.openmm_md.plain_md_methods import ( PlainMDProtocol, PlainMDProtocolUnit, PlainMDProtocolResult, ) @@ -314,7 +314,7 @@ def test_dry_run_complex(benzene_complex_system, tmpdir): with tmpdir.as_cwd(): sim = dag_unit.run(dry=True)['debug']['system'] - assert not ThermodynamicState(sim, temperature= + assert ThermodynamicState(sim, temperature= to_openmm(protocol.settings.thermo_settings.temperature)).is_periodic assert isinstance(ThermodynamicState(sim, temperature= to_openmm(protocol.settings.thermo_settings.temperature)).barostat, From b9c924b4f62991a1b5a2c26fc809157ab46d9b19 Mon Sep 17 00:00:00 2001 From: Hannah Baumann <43765638+hannahbaumann@users.noreply.github.com> Date: Tue, 12 Dec 2023 11:37:57 +0100 Subject: [PATCH 64/66] Update openfe/protocols/openmm_md/plain_md_methods.py Co-authored-by: Irfan Alibay --- openfe/protocols/openmm_md/plain_md_methods.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 27480c661..eb0e7e1dc 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -78,8 +78,8 @@ def get_traj_filename(self) -> list[pathlib.PurePath]: Returns ------- - traj : list - list of paths (pathlib.PurePath) to the simulation trajectory + traj : list[pathlib.Path] + list of paths (pathlib.Path) to the simulation trajectory """ traj = [pus[0].outputs['nc'] for pus in self.data.values()] From 4704336c3f5ce87179ac87ab01b1453de04bebed Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 12 Dec 2023 11:40:05 +0100 Subject: [PATCH 65/66] Change purepath to path --- openfe/tests/protocols/test_openmm_plain_md_protocols.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/test_openmm_plain_md_protocols.py index d29d08790..b96bd94a4 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/test_openmm_plain_md_protocols.py @@ -443,10 +443,10 @@ def test_get_traj_filename(self, protocolresult): traj = protocolresult.get_traj_filename() assert isinstance(traj, list) - assert isinstance(traj[0], pathlib.PurePath) + assert isinstance(traj[0], pathlib.Path) def test_get_pdb_filename(self, protocolresult): pdb = protocolresult.get_pdb_filename() assert isinstance(pdb, list) - assert isinstance(pdb[0], pathlib.PurePath) + assert isinstance(pdb[0], pathlib.Path) From 64109d8ed4737624b457c325c06570cf61244390 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 13 Dec 2023 15:52:48 +0100 Subject: [PATCH 66/66] Change PurePath to Path --- openfe/protocols/openmm_md/plain_md_methods.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index eb0e7e1dc..0feab1339 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -72,7 +72,7 @@ def get_uncertainty(self): return None - def get_traj_filename(self) -> list[pathlib.PurePath]: + def get_traj_filename(self) -> list[pathlib.Path]: """ Get a list of trajectory paths @@ -85,14 +85,14 @@ def get_traj_filename(self) -> list[pathlib.PurePath]: return traj - def get_pdb_filename(self) -> list[pathlib.PurePath]: + def get_pdb_filename(self) -> list[pathlib.Path]: """ Get a list of paths to the pdb files of the pre-minimized system. Returns ------- - pdbs : list - list of paths (pathlib.PurePath) to the pdb files + pdbs : list[pathlib.Path] + list of paths (pathlib.Path) to the pdb files """ pdbs = [pus[0].outputs['system_pdb'] for pus in self.data.values()]