Renaming and moving some of the protocol settings

openfe/protocols/openmm_afe/base.py

@@ -52,8 +52,9 @@
 )
 from openfe.protocols.openmm_afe.equil_afe_settings import (
     SolvationSettings,
-    AlchemicalSamplerSettings, OpenMMEngineSettings,
-    IntegratorSettings, SimulationSettings, LambdaSettings,
+    MultiStateSimulationSettings, OpenMMEngineSettings,
+    IntegratorSettings, LambdaSettings, OutputSettings,
+    ThermoSettings,
 )
 from openfe.protocols.openmm_rfe._rfe_utils import compute
 from ..openmm_utils import (
@@ -234,10 +235,11 @@ def _handle_settings(self):
           * solvation_settings : SolvationSettings
           * alchemical_settings : AlchemicalSettings
           * lambda_settings : LambdaSettings
-          * sampler_settings : AlchemicalSamplerSettings
+          * sampler_settings : MultiStateSimulationSettings
           * engine_settings : OpenMMEngineSettings
           * integrator_settings : IntegratorSettings
           * simulation_settings : SimulationSettings
+          * output_settings: OutputSettings
 
         Settings may change depending on what type of simulation you are
         running. Cherry pick them and return them to be available later on.
@@ -270,15 +272,14 @@ def _get_system_generator(
         system_generator : openmmforcefields.generator.SystemGenerator
           System Generator to parameterise this unit.
         """
-        ffcache = settings['simulation_settings'].forcefield_cache
+        ffcache = settings['output_settings'].forcefield_cache
         if ffcache is not None:
             ffcache = self.shared_basepath / ffcache
 
         system_generator = system_creation.get_system_generator(
             forcefield_settings=settings['forcefield_settings'],
-            thermo_settings=settings['thermo_settings'],
             integrator_settings=settings['integrator_settings'],
-            system_settings=settings['system_settings'],
+            thermo_settings=settings['thermo_settings'],
             cache=ffcache,
             has_solvent=solvent_comp is not None,
         )
@@ -537,7 +538,7 @@ def _get_reporter(
         self,
         topology: app.Topology,
         positions: openmm.unit.Quantity,
-        simulation_settings: SimulationSettings,
+        output_settings: OutputSettings,
     ) -> multistate.MultiStateReporter:
         """
         Get a MultistateReporter for the simulation you are running.
@@ -546,8 +547,8 @@ def _get_reporter(
         ----------
         topology : app.Topology
           A Topology of the system being created.
-        simulation_settings : SimulationSettings
-          Settings for the simulation.
+        output_settings: OutputSettings
+          Output settings for the simulations
 
         Returns
         -------
@@ -557,16 +558,16 @@ def _get_reporter(
         mdt_top = mdt.Topology.from_openmm(topology)
 
         selection_indices = mdt_top.select(
-                simulation_settings.output_indices
+                output_settings.output_indices
         )
 
-        nc = self.shared_basepath / simulation_settings.output_filename
-        chk = simulation_settings.checkpoint_storage
+        nc = self.shared_basepath / output_settings.output_filename
+        chk = output_settings.checkpoint_storage_filename
 
         reporter = multistate.MultiStateReporter(
             storage=nc,
             analysis_particle_indices=selection_indices,
-            checkpoint_interval=simulation_settings.checkpoint_interval.m,
+            checkpoint_interval=output_settings.checkpoint_interval.m,
             checkpoint_storage=chk,
         )
 
@@ -577,7 +578,7 @@ def _get_reporter(
                 mdt_top.subset(selection_indices),
             )
             traj.save_pdb(
-                self.shared_basepath / simulation_settings.output_structure
+                self.shared_basepath / output_settings.output_structure
             )
 
         return reporter
@@ -614,38 +615,46 @@ def _get_ctx_caches(
 
         return energy_context_cache, sampler_context_cache
 
+    @staticmethod
     def _get_integrator(
-        self,
-        integrator_settings: IntegratorSettings
+        integrator_settings: IntegratorSettings,
+        simulation_settings: MultiStateSimulationSettings
     ) -> openmmtools.mcmc.LangevinDynamicsMove:
         """
         Return a LangevinDynamicsMove integrator
 
         Parameters
         ----------
         integrator_settings : IntegratorSettings
+        simulation_settings : MultiStateSimulationSettings
 
         Returns
         -------
         integrator : openmmtools.mcmc.LangevinDynamicsMove
           A configured integrator object.
         """
+        # TODO: Check this is correct
+        tpi_fs = simulation_settings.time_per_iteration.to(unit.femtosecond).m
+        ts_fs = integrator_settings.timestep.to(unit.femtosecond).m
+        steps_per_iteration = int(round(tpi_fs / ts_fs))
+
         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,
+            collision_rate=to_openmm(integrator_settings.langevin_collision_rate),
+            n_steps=steps_per_iteration,
             reassign_velocities=integrator_settings.reassign_velocities,
             n_restart_attempts=integrator_settings.n_restart_attempts,
             constraint_tolerance=integrator_settings.constraint_tolerance,
         )
 
         return integrator
 
+    @staticmethod
     def _get_sampler(
-        self,
         integrator: openmmtools.mcmc.LangevinDynamicsMove,
         reporter: openmmtools.multistate.MultiStateReporter,
-        sampler_settings: AlchemicalSamplerSettings,
+        simulation_settings: MultiStateSimulationSettings,
+        thermo_settings: ThermoSettings,
         cmp_states: list[ThermodynamicState],
         sampler_states: list[SamplerState],
         energy_context_cache: openmmtools.cache.ContextCache,
@@ -660,8 +669,10 @@ def _get_sampler(
           The simulation integrator.
         reporter : openmmtools.multistate.MultiStateReporter
           The reporter to hook up to the sampler.
-        sampler_settings : AlchemicalSamplerSettings
+        simulation_settings : MultiStateSimulationSettings
           Settings for the alchemical sampler.
+        thermo_settings : ThermoSettings
+          Thermodynamic settings
         cmp_states : list[ThermodynamicState]
           A list of thermodynamic states to sample.
         sampler_states : list[SamplerState]
@@ -676,30 +687,37 @@ def _get_sampler(
         sampler : multistate.MultistateSampler
           A sampler configured for the chosen sampling method.
         """
+        rta_its, rta_min_its = settings_validation.convert_real_time_analysis_iterations(
+            simulation_settings=simulation_settings,
+        )
+        et_target_err = settings_validation.convert_target_error(
+            thermo_settings=thermo_settings,
+            simulation_settings=simulation_settings,
+        )
 
         # Select the right sampler
         # Note: doesn't need else, settings already validates choices
-        if sampler_settings.sampler_method.lower() == "repex":
+        if simulation_settings.sampler_method.lower() == "repex":
             sampler = multistate.ReplicaExchangeSampler(
                 mcmc_moves=integrator,
-                online_analysis_interval=sampler_settings.online_analysis_interval,
-                online_analysis_target_error=sampler_settings.online_analysis_target_error.m,
-                online_analysis_minimum_iterations=sampler_settings.online_analysis_minimum_iterations
+                online_analysis_interval=rta_its,
+                online_analysis_target_error=et_target_err,
+                online_analysis_minimum_iterations=rta_min_its
             )
-        elif sampler_settings.sampler_method.lower() == "sams":
+        elif simulation_settings.sampler_method.lower() == "sams":
             sampler = multistate.SAMSSampler(
                 mcmc_moves=integrator,
-                online_analysis_interval=sampler_settings.online_analysis_interval,
-                online_analysis_minimum_iterations=sampler_settings.online_analysis_minimum_iterations,
-                flatness_criteria=sampler_settings.flatness_criteria,
-                gamma0=sampler_settings.gamma0,
+                online_analysis_interval=rta_its,
+                online_analysis_minimum_iterations=rta_min_its,
+                flatness_criteria=simulation_settings.sams_flatness_criteria,
+                gamma0=simulation_settings.sams_gamma0,
             )
-        elif sampler_settings.sampler_method.lower() == 'independent':
+        elif simulation_settings.sampler_method.lower() == 'independent':
             sampler = multistate.MultiStateSampler(
                 mcmc_moves=integrator,
-                online_analysis_interval=sampler_settings.online_analysis_interval,
-                online_analysis_target_error=sampler_settings.online_analysis_target_error.m,
-                online_analysis_minimum_iterations=sampler_settings.online_analysis_minimum_iterations
+                online_analysis_interval=rta_its,
+                online_analysis_target_error=et_target_err,
+                online_analysis_minimum_iterations=rta_min_its,
             )
 
         sampler.create(
@@ -741,7 +759,10 @@ def _run_simulation(
           if not a dry run.
         """
         # Get the relevant simulation steps
-        mc_steps = settings['integrator_settings'].n_steps.m
+        mc_steps = settings_validation.convert_steps_per_iteration(
+            simulation_settings=settings['simulation_settings'],
+            integrator_settings=settings['integrator_settings'],
+        )
 
         equil_steps = settings_validation.get_simsteps(
             sim_length=settings['simulation_settings'].equilibration_length,
@@ -793,8 +814,8 @@ def _run_simulation(
             reporter.close()
 
             # clean up the reporter file
-            fns = [self.shared_basepath / settings['simulation_settings'].output_filename,
-                   self.shared_basepath / settings['simulation_settings'].checkpoint_storage]
+            fns = [self.shared_basepath / settings['output_settings'].output_filename,
+                   self.shared_basepath / settings['output_settings'].checkpoint_storage_filename]
             for fn in fns:
                 os.remove(fn)
 
@@ -878,7 +899,7 @@ def run(self, dry=False, verbose=True,
         # 11. Create the multistate reporter & create PDB
         reporter = self._get_reporter(
             omm_topology, positions,
-            settings['simulation_settings'],
+            settings['output_settings'],
         )
 
         # Wrap in try/finally to avoid memory leak issues
@@ -889,11 +910,15 @@ def run(self, dry=False, verbose=True,
             )
 
             # 13. Get integrator
-            integrator = self._get_integrator(settings['integrator_settings'])
+            integrator = self._get_integrator(
+                settings['integrator_settings'],
+                settings['simulation_settings'],
+            )
 
             # 14. Get sampler
             sampler = self._get_sampler(
-                integrator, reporter, settings['sampler_settings'],
+                integrator, reporter, settings['simulation_settings'],
+                settings['thermo_settings'],
                 cmp_states, sampler_states,
                 energy_ctx_cache, sampler_ctx_cache
             )
@@ -925,8 +950,8 @@ def run(self, dry=False, verbose=True,
                 del integrator, sampler
 
         if not dry:
-            nc = self.shared_basepath / settings['simulation_settings'].output_filename
-            chk = settings['simulation_settings'].checkpoint_storage
+            nc = self.shared_basepath / settings['output_settings'].output_filename
+            chk = settings['output_settings'].checkpoint_storage_filename
             return {
                 'nc': nc,
                 'last_checkpoint': chk,