Handling unsampled states in offline/realtime analysis. (#593)

choderalab · Jun 7, 2022 · cea9d2c · cea9d2c
1 parent 83dd2be
commit cea9d2c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/docs/releasehistory.rst b/docs/releasehistory.rst
@@ -8,6 +8,7 @@ Bugfixes
 --------
 - Bug in statistical inefficiency computation -- where self.max_n_iterations wasn't being used -- was fixed (`#577 <https://github.com/choderalab/openmmtools/pull/577>`_).
 - Bug in estimated performance in realtime yaml file -- fixed by iterating through all MCMC moves (`#578 <https://github.com/choderalab/openmmtools/pull/578>`_)
+- Bug in handling unsampled states in realtime/offline analysis -- fixed by using ``MultiStateSampler._unsampled_states`` to build the mbar estimate array. Issue `#592 <https://github.com/choderalab/openmmtools/issues/592>`_ (`#593 <https://github.com/choderalab/openmmtools/pull/593>`_)
 
 
 0.21.3 - Bugfix release

diff --git a/openmmtools/multistate/multistatesampler.py b/openmmtools/multistate/multistatesampler.py
@@ -1514,9 +1514,14 @@ def _offline_analysis(self):
         bump_error_counter = False
         # Set up analyzer
         # Unbias restraints is False because this will quickly accumulate large time to re-read the trajectories
-        # and unbias the restraints every online-analysis. Current model is to use the last_mbar_f_k as a
-        # hot-start to the analysis. Once we store unbias info as part of a CustomCVForce, we can revisit this choice.
-        analysis = MultiStateSamplerAnalyzer(self._reporter, analysis_kwargs={'initial_f_k': self._last_mbar_f_k},
+        # and unbias the restraints every online-analysis. Current model is to initialize the last_mbar_f_k_offline
+        # attribute to zeros the first time, and then using its last result to produce the subsequent ones.
+        # Once we store unbias info as part of a CustomCVForce, we can revisit this choice.
+        # Initialize with zeros on first run including unsampled states
+        if not hasattr(self, "_last_mbar_f_k_offline"):
+            self._last_mbar_f_k_offline = np.zeros(len(self._thermodynamic_states) + len(self._unsampled_states))
+        analysis = MultiStateSamplerAnalyzer(self._reporter,
+                                             analysis_kwargs={'initial_f_k': self._last_mbar_f_k_offline},
                                              unbias_restraint=False)
 
         # Indices for online analysis, "i'th index, j'th index"
@@ -1532,14 +1537,15 @@ def _offline_analysis(self):
         except ParameterError as e:
             # We don't update self._last_err_free_energy here since if it
             # wasn't below the target threshold before, it won't stop MultiStateSampler now.
+            logger.debug(f"ParameterError computing MBAR. {e}.")
             bump_error_counter = True
             self._online_error_bank.append(e)
             if len(self._online_error_bank) > 6:
                 # Cache only the last set
                 self._online_error_bank.pop(0)
             free_energy = None
         else:
-            self._last_mbar_f_k = mbar.f_k
+            self._last_mbar_f_k_offline = mbar.f_k
             free_energy = free_energy[idx, jdx]
             self._last_err_free_energy = err_free_energy[idx, jdx]
             logger.debug("Current Free Energy Estimate is {} +- {} kT".format(free_energy,
@@ -1564,7 +1570,7 @@ def _offline_analysis(self):
 
         # Write out the numbers
         self._reporter.write_online_data_dynamic_and_static(self._iteration,
-                                                            f_k=self._last_mbar_f_k,
+                                                            f_k_offline=self._last_mbar_f_k_offline,
                                                             free_energy=(free_energy, self._last_err_free_energy))
         # Write real time offline analysis YAML file
         # TODO: Specify units