Fix real time mbar disabling appending uncorrelated endstates (#979)

choderalab · Apr 12, 2022 · d3649d1 · d3649d1
1 parent 54b0a3d
commit d3649d1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git a/perses/app/setup_relative_calculation.py b/perses/app/setup_relative_calculation.py
@@ -660,7 +660,10 @@ def run_setup(setup_options, serialize_systems=True, build_samplers=True):
                 if phase == 'vacuum':
                     endstates = False
                 else:
-                    endstates = True
+                    # TODO: Make this True when dispersed.utils.create_endstates is fixed
+                    _logger.info("Disabling creation of endstates with expanded cutoffs in order to support new "
+                                 "alchemical factories.")
+                    endstates = False
 
                 if setup_options['fe_type'] == 'fah':
                     _logger.info('SETUP FOR FAH DONE')

diff --git a/perses/samplers/multistate.py b/perses/samplers/multistate.py
@@ -5,8 +5,6 @@
 from perses.annihilation.lambda_protocol import RelativeAlchemicalState, LambdaProtocol
 
 from openmmtools.multistate import sams, replicaexchange
-from openmmtools import cache, utils
-from perses.dispersed.utils import configure_platform
 from openmmtools.states import CompoundThermodynamicState, SamplerState, ThermodynamicState
 from perses.dispersed.utils import create_endstates
 
@@ -29,11 +27,11 @@ def __init__(self, *args, hybrid_factory=None, **kwargs):
         self._hybrid_factory = hybrid_factory
         super(HybridCompatibilityMixin, self).__init__(*args, **kwargs)
 
+    # TODO: Should this overload the create() method from parent instead of breing setup()?
     def setup(self, n_states, temperature, storage_file, minimisation_steps=100,
               n_replicas=None, lambda_schedule=None,
               lambda_protocol=LambdaProtocol(), endstates=True):
 
-
         from perses.dispersed import feptasks
 
         hybrid_system = self._factory.hybrid_system
@@ -47,8 +45,6 @@ def setup(self, n_states, temperature, storage_file, minimisation_steps=100,
         thermodynamic_state_list = []
         sampler_state_list = []
 
-        context_cache = cache.ContextCache()
-
         if n_replicas is None:
             _logger.info(f'n_replicas not defined, setting to match n_states, {n_states}')
             n_replicas = n_states
@@ -66,16 +62,14 @@ def setup(self, n_states, temperature, storage_file, minimisation_steps=100,
             difference = np.diff(lambda_schedule)
             assert ( all(i >= 0. for i in difference ) ), 'lambda_schedule must be monotonicly increasing'
 
-        #starting with the initial positions generated py geometry.py
-        sampler_state =  SamplerState(positions, box_vectors=hybrid_system.getDefaultPeriodicBoxVectors())
+        # starting with the initial positions generated py geometry.py
+        sampler_state = SamplerState(positions, box_vectors=hybrid_system.getDefaultPeriodicBoxVectors())
         for lambda_val in lambda_schedule:
             compound_thermodynamic_state_copy = copy.deepcopy(compound_thermodynamic_state)
             compound_thermodynamic_state_copy.set_alchemical_parameters(lambda_val,lambda_protocol)
             thermodynamic_state_list.append(compound_thermodynamic_state_copy)
-
-             # now generating a sampler_state for each thermodyanmic state, with relaxed positions
-            context, context_integrator = context_cache.get_context(compound_thermodynamic_state_copy)
-            feptasks.minimize(compound_thermodynamic_state_copy,sampler_state)
+            # now generating a sampler_state for each thermodynamic state
+            feptasks.minimize(compound_thermodynamic_state_copy, sampler_state, max_iterations=minimisation_steps)
             sampler_state_list.append(copy.deepcopy(sampler_state))
 
         reporter = storage_file