From 88bd0bf81faa2875bdac55c1f92f993b19b43783 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 28 Mar 2023 17:47:21 -0400
Subject: [PATCH 1/3] Replacing FireMinimizer with LocalEnergyMinimizer

---
 openmmtools/multistate/multistatesampler.py | 25 +++------------------
 1 file changed, 3 insertions(+), 22 deletions(-)

diff --git a/openmmtools/multistate/multistatesampler.py b/openmmtools/multistate/multistatesampler.py
index d259a65d..f2dd645d 100644
--- a/openmmtools/multistate/multistatesampler.py
+++ b/openmmtools/multistate/multistatesampler.py
@@ -1356,11 +1356,8 @@ def _minimize_replica(self, replica_id, tolerance, max_iterations):
         thermodynamic_state = self._thermodynamic_states[thermodynamic_state_id]
         sampler_state = self._sampler_states[replica_id]
 
-        # Use the FIRE minimizer
-        integrator = FIREMinimizationIntegrator(tolerance=tolerance)
-
         # Get context and bound integrator from energy_context_cache
-        context, integrator = self.energy_context_cache.get_context(thermodynamic_state, integrator)
+        context, integrator = self.energy_context_cache.get_context(thermodynamic_state)
         # inform of platform used in current context
         logger.debug(f"{type(integrator).__name__}: Minimize using {context.getPlatform().getName()} platform.")
 
@@ -1373,31 +1370,15 @@ def _minimize_replica(self, replica_id, tolerance, max_iterations):
             replica_id + 1, self.n_replicas, initial_energy))
 
         # Minimize energy.
-        try:
-            if max_iterations == 0:
-                logger.debug('Using FIRE: tolerance {} minimizing to convergence'.format(tolerance))
-                while integrator.getGlobalVariableByName('converged') < 1:
-                    integrator.step(50)
-            else:
-                logger.debug('Using FIRE: tolerance {} max_iterations {}'.format(tolerance, max_iterations))
-                integrator.step(max_iterations)
-        except Exception as e:
-            if str(e) == 'Particle coordinate is nan':
-                logger.debug('NaN encountered in FIRE minimizer; falling back to L-BFGS after resetting positions')
-                sampler_state.apply_to_context(context)
-                openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations)
-            else:
-                raise e
+        openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations)
 
         # Get the minimized positions.
         sampler_state.update_from_context(context)
 
         # Compute the final energy of the system for logging.
         final_energy = thermodynamic_state.reduced_potential(sampler_state)
-        logger.debug('Replica {}/{}: final energy {:8.3f}kT'.format(
-            replica_id + 1, self.n_replicas, final_energy))
+        logger.debug(f'Replica {replica_id + 1}/{self.n_replicas}: final energy {final_energy:8.3f}kT')
         # TODO if energy > 0, use slower openmm minimizer
-
         # Clean up the integrator
         del context
 

From 0782749fbf419b470c8878a8ffd650d84ee3cf7f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Wed, 29 Mar 2023 17:23:07 -0400
Subject: [PATCH 2/3] Removing unnecessary import

---
 openmmtools/multistate/multistatesampler.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/openmmtools/multistate/multistatesampler.py b/openmmtools/multistate/multistatesampler.py
index f2dd645d..ddb7ac9b 100644
--- a/openmmtools/multistate/multistatesampler.py
+++ b/openmmtools/multistate/multistatesampler.py
@@ -49,8 +49,6 @@
 from openmmtools.multistate.utils import SimulationNaNError
 from openmmtools.multistate.pymbar import ParameterError
 
-from openmmtools.integrators import FIREMinimizationIntegrator
-
 logger = logging.getLogger(__name__)
 
 

From cc49fbc3b9ea1b9c6b019152c399351bf4d64ca4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Thu, 6 Apr 2023 17:47:59 -0400
Subject: [PATCH 3/3] shoudn't need to delete the context here

---
 openmmtools/multistate/multistatesampler.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/openmmtools/multistate/multistatesampler.py b/openmmtools/multistate/multistatesampler.py
index ddb7ac9b..93d03990 100644
--- a/openmmtools/multistate/multistatesampler.py
+++ b/openmmtools/multistate/multistatesampler.py
@@ -1376,9 +1376,6 @@ def _minimize_replica(self, replica_id, tolerance, max_iterations):
         # Compute the final energy of the system for logging.
         final_energy = thermodynamic_state.reduced_potential(sampler_state)
         logger.debug(f'Replica {replica_id + 1}/{self.n_replicas}: final energy {final_energy:8.3f}kT')
-        # TODO if energy > 0, use slower openmm minimizer
-        # Clean up the integrator
-        del context
 
         # Return minimized positions.
         return sampler_state.positions