Repository: choderalab/perses
Issue: Reset nonequilibrium integrators before use in nonequilibrium tasklet

Nonequilibrium integrators need to call `integrator.reset()` before switching in order to work since they are being retrieved from a cached `Context`.

Changed files:
perses/distributed/relative_setup.py, perses/distributed/feptasks.py

diff --git a/perses/distributed/feptasks.py b/perses/distributed/feptasks.py
index 9eb72699..2398304a 100644
--- a/perses/distributed/feptasks.py
+++ b/perses/distributed/feptasks.py
@@ -48,50 +48,21 @@ class NonequilibriumSwitchTask(celery.Task):
         platform = openmm.Platform.getPlatformByName("OpenCL")
         self._cache = cache.ContextCache(platform=platform)
 
-def restore_properties(integrator, measure_shadow_work=False, measure_heat=False, metropolized_integrator=False):
-    """
-    This is a temporary utility function to restore the properties of the AlchemicalNonequilibriumLangevin...
-    integrator so that we can avoid exceptions.
-
-    Parameters
-    ----------
-    integrator : AlchemicalNone...
-        integrator to which the attributes should be added
-    measure_shadow_work : bool, default False
-        whether to tell it to measure shadow work
-    measure_heat : bool, default False
-        whether to tell it to measure heat
-    metropolized_integrator : bool, default False
-        whether the integrator is metropolized
-
-    Returns
-    -------
-    integrator : AlchemicalNone...
-        integrator with attributes added
-    """
-    integrator._measure_shadow_work = measure_shadow_work
-    integrator._measure_heat = measure_heat
-    integrator._metropolized_integrator = metropolized_integrator
-    return integrator
-
-
 @app.task(bind=True, base=NonequilibriumSwitchTask, serializer="pickle")
 def run_protocol(self, starting_positions, nsteps, thermodynamic_state, integrator):
     """
-    Perform a switching protocol and return the nonequilibrium switching weight
+    Perform a nonequilibrium switching protocol and return the nonequilibrium protocol work.
 
     Returns
     -------
-    weight : float64
-        The nonequilibrium switching weight
+    work : float
+        The dimensionless nonequilibrium protocol work.
     """
-    integrator = restore_properties(integrator)
     switching_ctx, integrator_neq = self._cache.get_context(thermodynamic_state, integrator)
-    integrator_neq = restore_properties(integrator_neq)
     switching_ctx.setPositions(starting_positions)
+    integrator_neq.reset()
     integrator_neq.step(nsteps)
     work = integrator_neq.get_protocol_work(dimensionless=True)
-    integrator_neq.reset()
     return work
 
 @app.task(bind=True, base=NonequilibriumSwitchTask, serializer="pickle")
@@ -103,7 +74,6 @@ def run_equilibrium(self, starting_positions, nsteps, lambda_state, functions, t
     -------
     positions : [n, 3] np.ndarray quantity
     """
-    integrator = restore_properties(integrator)
     equilibrium_ctx, integrator = self._cache.get_context(thermodynamic_state, integrator)
     equilibrium_ctx.setPositions(starting_positions)
     for parm in functions.keys():
@@ -115,8 +85,6 @@ def run_equilibrium(self, starting_positions, nsteps, lambda_state, functions, t
 
 @app.task(bind=True, base=NonequilibriumSwitchTask, serializer="pickle")
 def minimize(self, starting_positions, nsteps_max, lambda_state, functions, thermodynamic_state, integrator):
-    integrator = restore_properties(integrator)
-
     equilibrium_ctx, integrator = self._cache.get_context(thermodynamic_state, integrator)
 
     equilibrium_ctx.setPositions(starting_positions)
@@ -133,4 +101,4 @@ def minimize(self, starting_positions, nsteps_max, lambda_state, functions, ther
 @app.task(bind=True, base=NonequilibriumSwitchTask, serializer="pickle")
 def dummy_task(self, integrator):
     rdb.set_trace()
-    return integrator
\\ No newline at end of file
+    return integrator
diff --git a/perses/distributed/relative_setup.py b/perses/distributed/relative_setup.py
index 67346e51..4abecbc9 100644
--- a/perses/distributed/relative_setup.py
+++ b/perses/distributed/relative_setup.py
@@ -464,6 +464,3 @@ if __name__=="__main__":
     ne_fep.run_nonequilibrium_task()
     ne_fep.run_equilibrium()
     ne_fep.collect_ne_work()
-
-
-