Fix MTREXEE FE issue

Author: ajfriedman22

ensemble_md/analysis/analyze_free_energy.py

@@ -247,7 +247,7 @@ def _combine_df_adjacent(df_adjacent, state_ranges, df_err_adjacent=None, err_ty
     return df, df_err, overlap_bool
 
 
-def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="propagate", n_bootstrap=None, seed=None):
+def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="propagate", n_bootstrap=None, seed=None, MTREXEE=False):
     """
     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
@@ -275,6 +275,8 @@ def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="prop
     seed : int, Optional
         The random seed for bootstrapping. Only relevant when :code:`err_method` is :code:`"bootstrap"`.
         The default is :code:`None`.
+    MTREXEE : bool
+        Whether this is a MT-REXEE simulation or not
 
     Returns
     -------
@@ -299,7 +301,10 @@ def calculate_free_energy(data, state_ranges, df_method="MBAR", err_method="prop
         >>> f, _, _ = analyze_free_energy.calculate_free_energy(data_list, state_ranges, "MBAR", "propagate")
     """
     n_sim = len(data)
-    n_tot = state_ranges[-1][-1] + 1
+    if MTREXEE is False:
+        n_tot = state_ranges[-1][-1] + 1
+    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, df_err_adjacent, err_type='propagate')

ensemble_md/cli/analyze_REXEE.py

@@ -182,7 +182,7 @@ def main():
     rmse = analyze_traj.calc_hist_rmse(hist_data, REXEE.state_ranges)
     print(f'The RMSE of accumulated histogram counts of the state index: {rmse:.0f}')
 
-    if REXEE.proposal != 'forced_random' and REXEE.proposal != 'forced_swap':  # Need to FIX THIS FOR FORCED-RANDOM
+    if REXEE.proposal != 'random_range':  # Need to FIX THIS FOR RANDOM-Range
         # 2-4. Stitch the time series of state index for different replicas
         if os.path.isfile(args.state_trajs_for_sim) is True:
             print('\n2-4. Reading in the stitched time series of state index for different replicas ...')
@@ -485,7 +485,7 @@ def main():
                         pickle.dump(data_all, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
                 # 4-2. Calculate the free energy profile
-                f, f_err, estimators = analyze_free_energy.calculate_free_energy(data_list, REXEE.state_ranges[sim], REXEE.df_method, REXEE.err_method, REXEE.n_bootstrap, REXEE.seed)  # noqa: E501
+                f, f_err, estimators = analyze_free_energy.calculate_free_energy(data_list, REXEE.state_ranges[sim], REXEE.df_method, REXEE.err_method, REXEE.n_bootstrap, REXEE.seed, MTREXEE=True)  # noqa: E501
 
                 print('Plotting the full-range free energy profile ...')
                 analyze_free_energy.plot_free_energy(f, f_err, f'{args.dir}/free_energy_profile_{sim}.png')