Merge pull request #720 from OpenFreeEnergy/afe-md

Add MD equilibration for AFE calculations

devtools/data/gen-serialized-results.py

@@ -91,8 +91,13 @@ def generate_md_json(smc):
 
 def generate_ahfe_json(smc):
     settings = AbsoluteSolvationProtocol.default_settings()
+    settings.solvent_equil_simulation_settings.equilibration_length_nvt = 10 * unit.picosecond
+    settings.solvent_equil_simulation_settings.equilibration_length = 10 * unit.picosecond
+    settings.solvent_equil_simulation_settings.production_length = 10 * unit.picosecond
     settings.solvent_simulation_settings.equilibration_length = 10 * unit.picosecond
     settings.solvent_simulation_settings.production_length = 500 * unit.picosecond
+    settings.vacuum_equil_simulation_settings.equilibration_length = 10 * unit.picosecond
+    settings.vacuum_equil_simulation_settings.production_length = 10 * unit.picosecond
     settings.vacuum_simulation_settings.equilibration_length = 10 * unit.picosecond
     settings.vacuum_simulation_settings.production_length = 1000 * unit.picosecond
     settings.lambda_settings.lambda_elec = [0.0, 0.25, 0.5, 0.75, 1.0, 1.0,

openfe/protocols/openmm_afe/base.py

@@ -60,6 +60,7 @@
     ThermoSettings, OpenFFPartialChargeSettings,
 )
 from openfe.protocols.openmm_rfe._rfe_utils import compute
+from openfe.protocols.openmm_md.plain_md_methods import PlainMDProtocolUnit
 from ..openmm_utils import (
     settings_validation, system_creation,
     multistate_analysis, charge_generation
@@ -152,37 +153,113 @@ def _get_alchemical_indices(omm_top: openmm.Topology,
 
         return atom_ids
 
-    @staticmethod
-    def _pre_minimize(system: openmm.System,
-                      positions: omm_unit.Quantity) -> npt.NDArray:
+    def _pre_equilibrate(
+        self,
+        system: openmm.System,
+        topology: openmm.app.Topology,
+        positions: omm_unit.Quantity,
+        settings: dict[str, SettingsBaseModel],
+        dry: bool
+    ) -> omm_unit.Quantity:
         """
-        Short CPU minization of System to avoid GPU NaNs
+        Run a non-alchemical equilibration to get a stable system.
 
         Parameters
         ----------
         system : openmm.System
-          An OpenMM System to minimize.
-        positionns : openmm.unit.Quantity
+          An OpenMM System to equilibrate.
+        topology : openmm.app.Topology
+          OpenMM Topology of the System.
+        positions : openmm.unit.Quantity
           Initial positions for the system.
+        settings : dict[str, SettingsBaseModel]
+          A dictionary of settings objects. Expects the
+          following entries:
+          * `engine_settings`
+          * `thermo_settings`
+          * `integrator_settings`
+          * `equil_simulation_settings`
+          * `equil_output_settings`
+        dry: bool
+          Whether or not this is a dry run.
 
         Returns
         -------
-        minimized_positions : npt.NDArray
-          Minimized positions
+        equilibrated_positions : npt.NDArray
+          Equilibrated system positions
         """
-        integrator = openmm.VerletIntegrator(0.001)
-        context = openmm.Context(
-            system, integrator,
-            openmm.Platform.getPlatformByName('CPU'),
+        # Prep the simulation object
+        platform = compute.get_openmm_platform(
+            settings['engine_settings'].compute_platform
         )
-        context.setPositions(positions)
-        # Do a quick 100 steps minimization, usually avoids NaNs
-        openmm.LocalEnergyMinimizer.minimize(
-            context, maxIterations=100
+
+        integrator = openmm.LangevinMiddleIntegrator(
+            to_openmm(settings['thermo_settings'].temperature),
+            to_openmm(settings['integrator_settings'].langevin_collision_rate),
+            to_openmm(settings['integrator_settings'].timestep),
         )
-        state = context.getState(getPositions=True)
-        minimized_positions = state.getPositions(asNumpy=True)
-        return minimized_positions
+
+        simulation = openmm.app.Simulation(
+            topology=topology,
+            system=system,
+            integrator=integrator,
+            platform=platform,
+        )
+
+        # Get the necessary number of steps
+        if settings['equil_simulation_settings'].equilibration_length_nvt is not None:
+            equil_steps_nvt = settings_validation.get_simsteps(
+                sim_length=settings[
+                    'equil_simulation_settings'].equilibration_length_nvt,
+                timestep=settings['integrator_settings'].timestep,
+                mc_steps=1,
+            )
+        else:
+            equil_steps_nvt = None
+
+        equil_steps_npt = settings_validation.get_simsteps(
+            sim_length=settings['equil_simulation_settings'].equilibration_length,
+            timestep=settings['integrator_settings'].timestep,
+            mc_steps=1,
+        )
+
+        prod_steps_npt = settings_validation.get_simsteps(
+            sim_length=settings['equil_simulation_settings'].production_length,
+            timestep=settings['integrator_settings'].timestep,
+            mc_steps=1,
+        )
+
+        if self.verbose:
+            logger.info("running non-alchemical equilibration MD")
+
+        # Don't do anything if we're doing a dry run
+        if dry:
+            return positions
+
+        # Use the _run_MD method from the PlainMDProtocolUnit
+        # Should in-place modify the simulation
+        PlainMDProtocolUnit._run_MD(
+            simulation=simulation,
+            positions=positions,
+            simulation_settings=settings['equil_simulation_settings'],
+            output_settings=settings['equil_output_settings'],
+            temperature=settings['thermo_settings'].temperature,
+            barostat_frequency=settings['integrator_settings'].barostat_frequency,
+            timestep=settings['integrator_settings'].timestep,
+            equil_steps_nvt=equil_steps_nvt,
+            equil_steps_npt=equil_steps_npt,
+            prod_steps=prod_steps_npt,
+            verbose=self.verbose,
+            shared_basepath=self.shared_basepath,
+        )
+
+        state = simulation.context.getState(getPositions=True)
+        equilibrated_positions = state.getPositions(asNumpy=True)
+
+        # cautiously delete out contexts & integrator
+        del simulation.context, integrator
+
+        return equilibrated_positions
 
     def _prepare(
         self, verbose: bool,
@@ -241,6 +318,8 @@ def _handle_settings(self):
           * integrator_settings : IntegratorSettings
           * simulation_settings : MultiStateSimulationSettings
           * output_settings: OutputSettings
+          * equil_simulation_settings: MDSimulationSettings
+          * equil_output_settings: MDOutputSettings
 
         Settings may change depending on what type of simulation you are
         running. Cherry pick them and return them to be available later on.
@@ -914,8 +993,10 @@ def run(self, dry=False, verbose=True,
             system_modeller, system_generator, list(smc_comps.values())
         )
 
-        # 6. Pre-minimize System (Test + Avoid NaNs)
-        positions = self._pre_minimize(omm_system, positions)
+        # 6. Pre-equilbrate System (Test + Avoid NaNs + get stable system)
+        positions = self._pre_equilibrate(
+            omm_system, omm_topology, positions, settings, dry
+        )
 
         # 7. Get lambdas
         lambdas = self._get_lambda_schedule(settings)

openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -28,6 +28,8 @@
     IntegratorSettings,
     OpenFFPartialChargeSettings,
     OutputSettings,
+    MDSimulationSettings,
+    MDOutputSettings,
 )
 import numpy as np
 
@@ -173,20 +175,41 @@ def must_be_positive(cls, v):
     """
 
     # Simulation run settings
+    vacuum_equil_simulation_settings: MDSimulationSettings
+    """
+    Pre-alchemical vacuum simulation control settings.
+
+    Notes
+    -----
+    The `NVT` equilibration should be set to 0 * unit.nanosecond
+    as it will not be run.
+    """
     vacuum_simulation_settings: MultiStateSimulationSettings
     """
     Simulation control settings, including simulation lengths
     for the vacuum transformation.
     """
+    solvent_equil_simulation_settings: MDSimulationSettings
+    """
+    Pre-alchemical solvent simulation control settings.
+    """
     solvent_simulation_settings: MultiStateSimulationSettings
     """
     Simulation control settings, including simulation lengths
     for the solvent transformation.
     """
+    vacuum_equil_output_settings: MDOutputSettings
+    """
+    Simulation output settings for the vacuum non-alchemical equilibration.
+    """
     vacuum_output_settings: OutputSettings
     """
     Simulation output settings for the vacuum transformation.
     """
+    solvent_equil_output_settings: MDOutputSettings
+    """
+    Simulation output settings for the solvent non-alchemical equilibration.
+    """
     solvent_output_settings: OutputSettings
     """
     Simulation output settings for the solvent transformation.

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -51,6 +51,7 @@
 from openfe.protocols.openmm_afe.equil_afe_settings import (
     AbsoluteSolvationSettings,
     OpenMMSolvationSettings, AlchemicalSettings, LambdaSettings,
+    MDSimulationSettings, MDOutputSettings,
     MultiStateSimulationSettings, OpenMMEngineSettings,
     IntegratorSettings, OutputSettings,
     OpenFFPartialChargeSettings,
@@ -421,6 +422,17 @@ def _default_settings(cls):
             vacuum_engine_settings=OpenMMEngineSettings(),
             solvent_engine_settings=OpenMMEngineSettings(),
             integrator_settings=IntegratorSettings(),
+            solvent_equil_simulation_settings=MDSimulationSettings(
+                equilibration_length_nvt=0.1 * unit.nanosecond,
+                equilibration_length=0.2 * unit.nanosecond,
+                production_length=0.5 * unit.nanosecond,
+            ),
+            solvent_equil_output_settings=MDOutputSettings(
+                equil_nvt_structure='equil_nvt_structure.pdb',
+                equil_npt_structure='equil_npt_structure.pdb',
+                production_trajectory_filename='production_equil.xtc',
+                log_output='equil_simulation.log',
+            ),
             solvent_simulation_settings=MultiStateSimulationSettings(
                 n_replicas=14,
                 equilibration_length=1.0 * unit.nanosecond,
@@ -430,6 +442,17 @@ def _default_settings(cls):
                 output_filename='solvent.nc',
                 checkpoint_storage_filename='solvent_checkpoint.nc',
             ),
+            vacuum_equil_simulation_settings=MDSimulationSettings(
+                equilibration_length_nvt=None,
+                equilibration_length=0.2 * unit.nanosecond,
+                production_length=0.5 * unit.nanosecond,
+            ),
+            vacuum_equil_output_settings=MDOutputSettings(
+                equil_nvt_structure=None,
+                equil_npt_structure='equil_structure.pdb',
+                production_trajectory_filename='production_equil.xtc',
+                log_output='equil_simulation.log',
+            ),
             vacuum_simulation_settings=MultiStateSimulationSettings(
                 n_replicas=14,
                 equilibration_length=0.5 * unit.nanosecond,
@@ -636,6 +659,13 @@ def _create(
                       "passed")
             raise ValueError(errmsg)
 
+        # Check vacuum equilibration MD settings is 0 ns
+        nvt_time = self.settings.vacuum_equil_simulation_settings.equilibration_length_nvt
+        if nvt_time is not None:
+            if not np.allclose(nvt_time, 0 * unit.nanosecond):
+                errmsg = "NVT equilibration cannot be run in vacuum simulation"
+                raise ValueError(errmsg)
+
         # Get the name of the alchemical species
         alchname = alchem_comps['stateA'][0].name
 
@@ -749,6 +779,8 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
             * lambda_settings : LambdaSettings
             * engine_settings : OpenMMEngineSettings
             * integrator_settings : IntegratorSettings
+            * equil_simulation_settings : MDSimulationSettings
+            * equil_output_settings : MDOutputSettings
             * simulation_settings : SimulationSettings
             * output_settings: OutputSettings
         """
@@ -763,6 +795,8 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
         settings['lambda_settings'] = prot_settings.lambda_settings
         settings['engine_settings'] = prot_settings.vacuum_engine_settings
         settings['integrator_settings'] = prot_settings.integrator_settings
+        settings['equil_simulation_settings'] = prot_settings.vacuum_equil_simulation_settings
+        settings['equil_output_settings'] = prot_settings.vacuum_equil_output_settings
         settings['simulation_settings'] = prot_settings.vacuum_simulation_settings
         settings['output_settings'] = prot_settings.vacuum_output_settings
 
@@ -834,6 +868,8 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
             * lambda_settings : LambdaSettings
             * engine_settings : OpenMMEngineSettings
             * integrator_settings : IntegratorSettings
+            * equil_simulation_settings : MDSimulationSettings
+            * equil_output_settings : MDOutputSettings
             * simulation_settings : MultiStateSimulationSettings
             * output_settings: OutputSettings
         """
@@ -848,6 +884,8 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
         settings['lambda_settings'] = prot_settings.lambda_settings
         settings['engine_settings'] = prot_settings.solvent_engine_settings
         settings['integrator_settings'] = prot_settings.integrator_settings
+        settings['equil_simulation_settings'] = prot_settings.solvent_equil_simulation_settings
+        settings['equil_output_settings'] = prot_settings.solvent_equil_output_settings
         settings['simulation_settings'] = prot_settings.solvent_simulation_settings
         settings['output_settings'] = prot_settings.solvent_output_settings