Add NVT and NPT equilibration

OpenFreeEnergy · Nov 20, 2023 · e160600 · e160600
1 parent 22eec31
commit e160600
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 7 deletions.
diff --git 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
@@ -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'."""