Enable MCMC Integrators to do a .reinitilize action on contexts as …

…a last-chance for restart

Cut 0.13.4

docs/releasehistory.rst

@@ -1,6 +1,21 @@
 Release History
 ===============
 
+0.13.4 - Barostat/External Force Bugfix, Restart Robustness
+===========================================================
+
+Bug fixes
+---------
+- Fixed implementation bug where ``CustomExternalForce`` restraining atoms to absolute coordinates caused an issue
+  when a Barostat was used (`#310 <https://github.com/choderalab/openmmtools/issues/310>`_)
+
+Enhancements
+------------
+- MCMC Integrators now attempt to re-initialize the ``Context`` object on the last restart attempt when NaN's are
+  encountered. This has internally been shown to correct some instances where normally resetting positions does
+  not work around the NaN's. This is a slow step relative to just resetting positions, but better than simulation
+  crashing.
+
 0.13.3 - Critical Bugfix to SamplerState Context Manipulation
 =============================================================
 
@@ -86,8 +101,8 @@ New features
 - Add ``restrain_atoms`` to restrain molecule conformation through an harmonic restrain
   (`#255 <https://github.com/choderalab/openmmtools/issues/255>`_)
 
-Bugfixes
---------
+Bug fixes
+---------
 
 - Bugfix for ``testsystems`` that use implicit solvent (`#250 <https://github.com/choderalab/openmmtools/issues/250>`_)
 - Bugfix for ``ContextCache``: two consecutive calls retrieve the same ``Context`` with same thermodynamic state and no

openmmtools/mcmc.py

@@ -577,8 +577,10 @@ class BaseIntegratorMove(object):
     restart_attempts : int, optional
         When greater than 0, if after the integration there are NaNs in energies,
         the move will restart. When the integrator has a random component, this
-        may help recovering. An IntegratorMoveError is raised after the given
-        number of attempts if there are still NaNs.
+        may help recovering. On the last attempt, the ``Context`` is
+        re-initialized in a slower process, but better than the simulation
+        crashing. An IntegratorMoveError is raised after the given number of
+        attempts if there are still NaNs.
 
     Attributes
     ----------
@@ -695,8 +697,14 @@ def apply(self, thermodynamic_state, sampler_state):
                 err_msg = ('Potential energy is NaN after {} attempts of integration '
                            'with move {}'.format(attempt_counter, self.__class__.__name__))
 
+                # If we are on our last chance before crash, try to re-initialize context
+                if attempt_counter == self.n_restart_attempts - 1:
+                    logger.error(err_msg + ' Trying to reinitialize Context as a last-resort restart attempt...')
+                    context.reinitialize()
+                    sampler_state.apply_to_context(context)
+                    thermodynamic_state.apply_to_context(context)
                 # If we have hit the number of restart attempts, raise an exception.
-                if attempt_counter == self.n_restart_attempts:
+                elif attempt_counter == self.n_restart_attempts:
                     # Restore the context to the state right before the integration.
                     sampler_state.apply_to_context(context)
                     logger.error(err_msg)

setup.py

@@ -15,7 +15,7 @@
 
 ########################
 VERSION = "0.13.4"
-ISRELEASED = False
+ISRELEASED = True
 __version__ = VERSION
 ########################
 CLASSIFIERS = """\