fix merge

Author: ajfriedman22

ensemble_md/analysis/analyze_free_energy.py

@@ -178,7 +178,7 @@ def _calculate_df_adjacent(estimators):
     return df_adjacent, df_err_adjacent
 
 
-def _combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent=None, err_type="propagate"):
+def _combine_df_adjacent(df_adjacent, state_ranges, df_err_adjacent=None, err_type="propagate"):
     """
     An internal function used in :func:`calculate_free_energy` to combine the free energy differences between
     adjacent states in different state ranges using either simple means or inverse-variance weighted means.
@@ -215,6 +215,7 @@ def _combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent=None,
     --------
     :func:`calculate_free_energy`
     """
+    n_tot = state_ranges[-1][-1] + 1
     df, df_err, overlap_bool = [], [], []
     for i in range(n_tot - 1):
         # df_list is a list of free energy difference between sates i and i+1 in different replicas
@@ -248,7 +249,7 @@ def _combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent=None,
     return df, df_err, overlap_bool
 
 
-def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="propagate", n_bootstrap=None, seed=None, MTREXEE=False):
+def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="propagate", n_bootstrap=None, seed=None, MTREXEE=False):  # noqa: E501
     """
     Caculates the averaged free energy profile with the chosen method given :math:`u_{nk}` or :math:`dH/dλ` data
     obtained from all replicas of the REXEE simulation. Available methods include TI, BAR, and MBAR. TI
@@ -307,8 +308,11 @@ def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="prop
     else:
         n_tot = state_ranges[-1] + 1
     estimators = _apply_estimators(data, df_method)
-    df_adjacent, df_err_adjacent = _calculate_df_adjacent(estimators)
-    df, df_err, overlap_bool = _combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent, err_type='propagate')
+    if MTREXEE is False:
+        df_adjacent, df_err_adjacent = _calculate_df_adjacent(estimators)
+        df, df_err, overlap_bool = _combine_df_adjacent(df_adjacent, state_ranges, df_err_adjacent, err_type='propagate')  # noqa: E501
+    else:
+        df, df_err = _calculate_df_adjacent(estimators)
 
     if err_method == 'bootstrap':
         if seed is not None:
@@ -320,15 +324,18 @@ def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="prop
         for b in range(n_bootstrap):
             sampled_data = [sampled_data_all[i].iloc[b * len(data[i]):(b + 1) * len(data[i])] for i in range(n_sim)]
             bootstrap_estimators = _apply_estimators(sampled_data, df_method)
-            df_adjacent, df_err_adjacent = _calculate_df_adjacent(bootstrap_estimators)
-            df_sampled, _, overlap_bool = _combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent, err_type='propagate')  # doesn't matter what value err_type here is # noqa: E501
+            if MTREXEE is False:
+                df_adjacent, df_err_adjacent = _calculate_df_adjacent(bootstrap_estimators)
+                df_sampled, _, overlap_bool = _combine_df_adjacent(df_adjacent, state_ranges, df_err_adjacent, err_type='propagate')  # doesn't matter what value err_type here is # noqa: E501
+            else:
+                df_sampled, _ = _calculate_df_adjacent(bootstrap_estimators)
             df_bootstrap.append(df_sampled)
         error_bootstrap = np.std(df_bootstrap, axis=0, ddof=1)
 
         # Replace the value in df_err with value in error_bootstrap if df_err corresponds to
         # the df between overlapping states
         for i in range(n_tot - 1):
-            if overlap_bool[i] is True:
+            if MTREXEE is True or overlap_bool[i] is True:
                 print(f'Replaced the propagated error with the bootstrapped error for states {i} and {i + 1}: {df_err[i]:.5f} -> {error_bootstrap[i]:.5f}.')  # noqa: E501
                 df_err[i] = error_bootstrap[i]
     elif err_method == 'propagate':

ensemble_md/cli/run_REXEE.py

@@ -322,7 +322,7 @@ def main():
                             gro = f'{REXEE.working_dir}/sim_{j}/iteration_{i-1}/confout.gro'
                             if os.path.exists(gro_backup):
                                 os.rename(gro_backup, gro)
-                                
+
                         for j in range(len(swap_list)):
                             print('\nModifying the coordinates of the following output GRO files ...')
                             # gro_1 and gro_2 are the simlation outputs (that we want to back up) and the inputs to modify_coords  # noqa: E501

ensemble_md/tests/test_analyze_free_energy.py

@@ -186,26 +186,25 @@ def test_combine_df_adjacent():
     df_adjacent = [[1, 3], [4, 6]]
     df_err_adjacent = [[0.1, 0.1], [0.1, 0.1]]
     state_ranges = [[0, 1, 2], [1, 2, 3]]
-    n_tot = state_ranges[-1][-1] + 1
 
     # Test 1: df_err_adjacent is None (in which case err_type is ignored)
     # Note that this test would lead to two harmless RuntimWarnings due to calculations like np.std([1], ddof=1), which return NaN  # noqa: E501
-    results = analyze_free_energy._combine_df_adjacent(df_adjacent, state_ranges, n_tot, None, "propagate")
+    results = analyze_free_energy._combine_df_adjacent(df_adjacent, state_ranges, None, "propagate")
     assert results[0] == [1, 3.5, 6]
     assert math.isnan(results[1][0])
     assert results[1][1] == np.std([3, 4], ddof=1)
     assert math.isnan(results[1][2])
     assert results[2] == [False, True, False]
 
     # Test 2: df_err_adjacent is not None and err_type is "std"
-    results = analyze_free_energy._combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent, "std")
+    results = analyze_free_energy._combine_df_adjacent(df_adjacent, state_ranges, df_err_adjacent, "std")
     assert results[0] == [1, 3.5, 6]
     np.testing.assert_array_almost_equal(results[1], [0.1, np.std([3, 4], ddof=1), 0.1])
     assert results[2] == [False, True, False]
 
     # Test 3: df_err_adjacent is not None and err_type is "propagate"
     df_err_adjacent = [[0.1, 0.1], [0.2, 0.1]]  # make the errs different so that the weighted mean will not be equal to simple mean  # noqa: E501
-    results = analyze_free_energy._combine_df_adjacent(df_adjacent, state_ranges, n_tot, df_err_adjacent, "propagate")
+    results = analyze_free_energy._combine_df_adjacent(df_adjacent, state_ranges, df_err_adjacent, "propagate")
     assert results[0] == [1, utils.weighted_mean([3, 4], [0.1, 0.2])[0], 6]
     assert results[1] == [0.1, utils.weighted_mean([3, 4], [0.1, 0.2])[1], 0.1]
     assert results[2] == [False, True, False]