Enhance stability of minimization

openmmtools/multistate/multistatesampler.py

@@ -49,7 +49,8 @@
 from openmmtools.multistate.utils import SimulationNaNError
 from openmmtools.multistate.pymbar import ParameterError
 
-from openmmtools.integrators import FIREMinimizationIntegrator
+from openmmtools.integrators import FIREMinimizationIntegrator, GradientDescentMinimizationIntegrator
+from openmmtools.utils import get_fastest_platform
 
 logger = logging.getLogger(__name__)
 
@@ -1354,47 +1355,84 @@ def _minimize_replica(self, replica_id, tolerance, max_iterations):
         thermodynamic_state_id = self._replica_thermodynamic_states[replica_id]
         thermodynamic_state = self._thermodynamic_states[thermodynamic_state_id]
         sampler_state = self._sampler_states[replica_id]
-
-        # Use the FIRE minimizer
-        integrator = FIREMinimizationIntegrator(tolerance=tolerance)
+        # NOTE: At this point, sampler_state.potential_energy is None. Initialize below.
+
+        logger.debug('Sampler state {}/{}:'.format(replica_id + 1, self.n_replicas))
+
+        # Temporarily disable the barostat during minimization by setting the
+        # pressure to None. (Otherwise, the minimizer will modify the box
+        # vectors and may cause instabilities which are very difficult to debug.)
+        pressure = thermodynamic_state.pressure
+        thermodynamic_state.pressure = None
+
+        # Use Gradient Descent first for numerical stability
+        # This initial step size is 10 times smaller than the default, but
+        # 1. the goal here is stability, not convergence rate and
+        # 2. the step size is adaptive, so it should quickly increase anyway.
+        initial_step_size = 0.001 * unit.angstroms
+        integrator_grad = GradientDescentMinimizationIntegrator(initial_step_size)
+        # platform = self.energy_context_cache.platform # This is None, causes:
+        # "ValueError: To set platform_properties, you need to also specify the platform."
+        platform = get_fastest_platform(minimum_precision='mixed')
+        if 'Precision' in platform.getPropertyNames():
+            platform.setPropertyDefaultValue('Precision', 'mixed')
+        context_grad = thermodynamic_state.create_context(integrator_grad, platform)
+        # Initialize sampler_state.potential_energy using a context
+        sampler_state.apply_to_context(context_grad)
+        sampler_state.update_from_context(context_grad)
+
+        # Compute the energy of the system for logging.
+        energy = thermodynamic_state.reduced_potential(sampler_state)
+        volume = sampler_state.volume.value_in_unit(unit.nanometer**3)
+        bv = sampler_state.box_vectors
+        [bxx, byy, bzz] = [v.value_in_unit(unit.nanometer) for v in [bv[0][0], bv[1][1], bv[2][2]]]
+        logger.debug('energy {:8.3f}kT box_volume {:6.4f}nm**3 bxx {:2.4f}nm  byy {:2.4f}nm  bzz {:2.4f}nm'.format(
+            energy, volume, bxx, byy, bzz))
+
+        # If pressure is not None, use one timestep less than the barostat
+        # frequency so we absolutely do not apply any barostat moves.
+        bar_freq = thermodynamic_state.barostat.getFrequency() - 1 if thermodynamic_state.pressure is not None else 25 - 1
+        logger.debug('Using Gradient Descent: num_iterations {}'.format(bar_freq))
+        integrator_grad.step(bar_freq)
+        sampler_state.update_from_context(context_grad)
+        del context_grad
+
+        # Compute the energy of the system for logging.
+        energy = thermodynamic_state.reduced_potential(sampler_state)
+        volume = sampler_state.volume.value_in_unit(unit.nanometer**3)
+        bv = sampler_state.box_vectors
+        [bxx, byy, bzz] = [v.value_in_unit(unit.nanometer) for v in [bv[0][0], bv[1][1], bv[2][2]]]
+        logger.debug('energy {:8.3f}kT box_volume {:6.4f}nm**3 bxx {:2.4f}nm  byy {:2.4f}nm  bzz {:2.4f}nm'.format(
+            energy, volume, bxx, byy, bzz))
 
         # 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, integrator_grad)
         # inform of platform used in current context
         logger.debug(f"{type(integrator).__name__}: Minimize using {context.getPlatform().getName()} platform.")
 
         # Set initial positions and box vectors.
         sampler_state.apply_to_context(context)
 
-        # Compute the initial energy of the system for logging.
-        initial_energy = thermodynamic_state.reduced_potential(context)
-        logger.debug('Replica {}/{}: initial energy {:8.3f}kT'.format(
-            replica_id + 1, self.n_replicas, initial_energy))
+        # Use L-BFGS minimizer because (if pressure is not None) FIRE modifies
+        # the box vectors and can be unstable.
+        logger.debug('Using L-BFGS: tolerance {} max_iterations {}'.format(tolerance, max_iterations))
+        openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations)
+        sampler_state.update_from_context(context)
+
 
-        # 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
+        # Restore the barostat
+        thermodynamic_state.pressure = pressure
 
         # 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))
+        # Compute the energy of the system for logging.
+        energy = thermodynamic_state.reduced_potential(sampler_state)
+        volume = sampler_state.volume.value_in_unit(unit.nanometer**3)
+        bv = sampler_state.box_vectors
+        [bxx, byy, bzz] = [v.value_in_unit(unit.nanometer) for v in [bv[0][0], bv[1][1], bv[2][2]]]
+        logger.debug('energy {:8.3f}kT box_volume {:6.4f}nm**3 bxx {:2.4f}nm  byy {:2.4f}nm  bzz {:2.4f}nm'.format(
+            energy, volume, bxx, byy, bzz))
         # TODO if energy > 0, use slower openmm minimizer
 
         # Clean up the integrator