From 916a8323005dab202c5e00c11c42b9e773dbb65a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Thu, 12 Jan 2023 21:28:37 -0500
Subject: [PATCH] more intuitive handling of minimization steps kwarg

---
 perses/dispersed/feptasks.py  | 4 +++-
 perses/samplers/multistate.py | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/perses/dispersed/feptasks.py b/perses/dispersed/feptasks.py
index b2475d21a..141493bfe 100644
--- a/perses/dispersed/feptasks.py
+++ b/perses/dispersed/feptasks.py
@@ -659,7 +659,6 @@ def check_EquilibriumFEPTask(task):
 
 def minimize(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState,
              max_iterations: int=100) -> states.SamplerState:
-    # TODO: set max iterations to 1000 and check tyk2 benchmarks
     """
     Minimize the given system and state, up to a maximum number of steps.
     This does not return a copy of the samplerstate; it is an update-in-place.
@@ -687,6 +686,9 @@ def minimize(thermodynamic_state: states.ThermodynamicState, sampler_state: stat
         context, integrator = cache.global_context_cache.get_context(thermodynamic_state)
     sampler_state.apply_to_context(context, ignore_velocities = True)
     # TODO: Set logging for minimization
+    # If max_iterations = None, it means run as many iterations as needed to reach tolerance.
+    if max_iterations is None:
+        max_iterations = 0
     openmm.LocalEnergyMinimizer.minimize(context, maxIterations = max_iterations)
     sampler_state.update_from_context(context)
 
diff --git a/perses/samplers/multistate.py b/perses/samplers/multistate.py
index 179cfa954..a2d4c0239 100644
--- a/perses/samplers/multistate.py
+++ b/perses/samplers/multistate.py
@@ -84,7 +84,7 @@ def setup(self, n_states, temperature, storage_file, minimisation_steps=100,
         sampler_state = SamplerState(positions, box_vectors=hybrid_system.getDefaultPeriodicBoxVectors())
 
         for lambda_val in lambda_schedule:
-           # Create a compound thermodynamic for lambda_val and set alchemical parameters
+            # Create a compound thermodynamic for lambda_val and set alchemical parameters
             compound_thermodynamic_state_copy = copy.deepcopy(compound_thermodynamic_state)
             if factory_name == 'HybridTopologyFactory':
                 compound_thermodynamic_state_copy.set_alchemical_parameters(lambda_val,lambda_protocol)
@@ -93,7 +93,8 @@ def setup(self, n_states, temperature, storage_file, minimisation_steps=100,
             thermodynamic_state_list.append(compound_thermodynamic_state_copy)
 
             # Generate a sampler_state for each thermodynamic state
-            feptasks.minimize(compound_thermodynamic_state_copy, sampler_state, max_iterations=minimisation_steps)
+            if minimisation_steps != 0:  # only minimize if there are specified steps != 0
+                feptasks.minimize(compound_thermodynamic_state_copy, sampler_state, max_iterations=minimisation_steps)
             sampler_state_list.append(copy.deepcopy(sampler_state))
 
         reporter = storage_file