Change get_simsteps to handle NVT equilibration

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
         )
 

openfe/protocols/openmm_utils/settings_validation.py

@@ -37,47 +37,38 @@ 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
       Number of integration timesteps between MCMC moves.
 
     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