Make MBAR bootstrapping deterministic independent of verbose flag (#551)

* Add failing test

* Makes MBAR deterministic independent of verbose flag

* Minor clean up

* Fix windoze tests

* run black

---------

Co-authored-by: Mike Henry <11765982+mikemhenry@users.noreply.github.com>

Author: badisa

pymbar/mbar.py

@@ -186,6 +186,10 @@ def __init__(
             We usually just do steps of adaptive sampling without. "robust" would be the backup.
             Default: dict(method="adaptive", options=dict(min_sc_iter=0)),
 
+        rseed: int or None, optional, default=None
+            Seed to use when constructing a new RNGState. If rseed is None, will use the global
+            np.random RNG to generate a seed.
+
         Notes
         -----
         The reduced potential energy ``u_kn[k,n] = u_k(x_{ln})``, where the reduced potential energy ``u_l(x)`` is
@@ -266,30 +270,34 @@ def __init__(
         # verbosity level -- if True, will print extra debug information
         self.verbose = verbose
 
+        if rseed is None:
+            rseed = np.random.randint(np.iinfo(np.int32).max)
+        self.rng = np.random.default_rng(rseed)
+
         # perform consistency checks on the data.
 
         self.samestates = []
+        # if, for any set of data, all reduced potential energies are the same,
+        # they are probably the same state.
+        #
+        # This can take quite a while, so we do it on just
+        # the first few data points.
+        #
+        # determine if there are less than 50 points to compare energies.
+        # If so, use that number instead of 50.
+        # (the number 50 is pretty arbitrary)
+
+        maxpoint = 50
+        if self.N < maxpoint:
+            maxpoint = self.N
+        # pick random indices
+        # indices = np.arange(maxpoint) # can uncomment this if need to remove random choices in testing.
+        # Done outside of the verbose if statement to ensure deterministic results of bootstrapping, independent of
+        # the Verbose flag
+        indices = self.rng.choice(np.arange(self.N), maxpoint)
+        # this could possibly be made faster with np.unique(axis=0,return_indices=True)
+        # but not clear if needed.
         if self.verbose:
-            # if, for any set of data, all reduced potential energies are the same,
-            # they are probably the same state.
-            #
-            # This can take quite a while, so we do it on just
-            # the first few data points.
-            #
-            # determine if there are less than 50 points to compare energies.
-            # If so, use that number instead of 50.
-            # (the number 50 is pretty arbitrary)
-
-            maxpoint = 50
-            if self.N < maxpoint:
-                maxpoint = self.N
-
-            # this could possibly be made faster with np.unique(axis=0,return_indices=True)
-            # but not clear if needed.
-
-            # pick random indices
-            # indices = np.arange(maxpoint) # can uncomment this if need to remove random choices in testing.
-            indices = np.random.choice(np.arange(self.N), maxpoint)
             for k in range(K):
                 for l in range(k):
                     diffsum = 0
@@ -402,11 +410,6 @@ def __init__(
             elif pname == "bootstrap_solver_protocol":
                 bootstrap_solver_protocol = prot
 
-        if rseed != None:
-            self.rstate = np.random.get_state()
-        else:
-            np.random.seed(rseed)
-
         self.f_k = mbar_solvers.solve_mbar_for_all_states(
             self.u_kn, self.N_k, self.f_k, self.states_with_samples, solver_protocol
         )
@@ -426,7 +429,7 @@ def __init__(
                     # which of the indices are equal to K
                     k_indices = np.where(self.x_kindices == k)[0]
                     # pick new random ones, selected of these K.
-                    new_kindices = k_indices[np.random.randint(int(N_k[k]), size=int(N_k[k]))]
+                    new_kindices = k_indices[self.rng.integers(int(N_k[k]), size=int(N_k[k]))]
                     rints[k_indices] = new_kindices
                 # If we initialized with BAR, then BAR, starting from the provided initial_f_k as well.
                 if initialize == "BAR":

pymbar/tests/test_mbar.py

@@ -528,3 +528,18 @@ def test_mbar_compute_expectations_inner(mbar_and_test):
     u_n = u_kn[:2, :]
     state_map = np.array([[0, 0], [1, 0], [2, 0], [2, 1]], int)
     _ = mbar.compute_expectations_inner(A_in, u_n, state_map)
+
+
+@pytest.mark.parametrize("n_bootstrap", [1, 100])
+def test_mbar_bootstrap_deterministic_given_same_seed(fixed_harmonic_sample, n_bootstrap):
+    """Verify that providing a seed to the mbar bootstrap will produce the same values"""
+    _, u_kn, _, _ = fixed_harmonic_sample.sample(N_k, mode="u_kn")
+
+    mbar_a = MBAR(u_kn, N_k, verbose=True, n_bootstraps=n_bootstrap, rseed=814)
+    ref_fe_diff = mbar_a.compute_free_energy_differences(uncertainty_method="bootstrap")
+
+    # The verbose flag should have no impact on the determinism
+    mbar_b = MBAR(u_kn, N_k, verbose=False, n_bootstraps=n_bootstrap, rseed=814)
+    test_fe_diff = mbar_b.compute_free_energy_differences(uncertainty_method="bootstrap")
+    np.testing.assert_equal(ref_fe_diff["Delta_f"], test_fe_diff["Delta_f"])
+    np.testing.assert_equal(ref_fe_diff["dDelta_f"], test_fe_diff["dDelta_f"])