Merge pull request #31 from wehs7661/refactor_w_combine

Refactor weight combination schemes

docs/simulations.rst

@@ -277,34 +277,21 @@ include parameters for data analysis here.
       This could be useful for REXEE simulations for multiple serial mutations, where we enforce exchanges between states 4 and 5 (and 14 and 15) and perform
       coordinate manipulation.
   - :code:`proposal`: (Optional, Default: :code:`exhaustive`)
-      The method for proposing simulations to be swapped. Available options include :code:`single`, :code:`exhaustive`, :code:`neighboring`, and :code:`multiple`.
+      The method for proposing simulations to be swapped. Available options include :code:`single`, :code:`neighboring`, and :code:`exhaustive`.
       For more details, please refer to :ref:`doc_proposal`.
   - :code:`acceptance`: (Optional, Default: :code:`metropolis`)
       The Monte Carlo method for swapping simulations. Available options include :code:`same-state`/:code:`same_state`, :code:`metropolis`, and :code:`metropolis-eq`/:code:`metropolis_eq`. 
       For more details, please refer to :ref:`doc_acceptance`.
-  - :code:`w_combine`: (Optional, Default: :code:`None`)
-      The type of weights to be combined across multiple replicas in a weight-updating REXEE simulation. The following options are available:
-
-        - :code:`None`: No weight combination.
-        - :code:`final`: Combine the final weights.
-        - :code:`avg`: Combine the weights averaged over from last time the Wang-Landau incrementor was updated.
-
-      For more details about weight combination, please refer to :ref:`doc_w_schemes`.
-
-  - :code:`rmse_cutoff`: (Optional, Default: :code:`None`)
-      The cutoff for the root-mean-square error (RMSE) between the weights of the current iteration 
-      and the weights averaged over from the last time the Wang-Landau incrementor was updated.
-      For each replica, the RMSE between the averaged weights and the current weights will be calculated.
-      When :code:`rmse_cutoff` is specified, weight combination will be performed only if the maximum RMSE across all replicas
-      is smaller than the cutoff. Otherwise, weight combination is deactivated (even if :code:`w_combine` is specified)
-      because a larger RMSE indicates that the weights are noisy and should not be combined.
-      The default value is infinity, which means that weight combination will always be performed if :code:`w_combine` is specified.
-      The units of the cutoff are :math:`k_B T`.
+  - :code:`w_combine`: (Optional, Default: :code:`False`)
+      Whether to perform weight combination or not. Note that weights averaged over from the last time the Wang-Landau incrementor was updated (instead of 
+      final weights) will be used for weight combination. For more details about weight combination, please refer to :ref:`doc_w_schemes`.
+  - :code:`w_mean_type`: (Optional, Default: code:`simple`)
+      The type of mean to use when combining weights. Available options include :code:`simple` and :code:`weighted`.
+      For the later case, inverse-variance weighted means are used. 
   - :code:`N_cutoff`: (Optional, Default: 1000)
-      The histogram cutoff. -1 means that no histogram correction will be performed.
-  - :code:`n_ex`: (Optional, Default: 1)
-      The number of attempts swap during an exchange interval. This option is only relevant if the option :code:`proposal` is :code:`multiple`.
-      Otherwise, this option is ignored. For more details, please refer to :ref:`doc_multiple_swaps`.
+      The histogram cutoff for weight corrections. -1 means that no histogram correction will be performed.
+  - :code:`hist_corr` (Optional, Default: :code:`False`)
+      Whether to perform histogram correction. 
   - :code:`mdp_args`: (Optional, Default: :code:`None`)
       MDP parameters differing across replicas provided in a dictionary. For each key in the dictionary, the value should
       always be a list of length of the number of replicas. For example, :code:`{'ref_p': [1.0, 1.01, 1.02, 1.03]}` means that the
@@ -389,10 +376,10 @@ infinity internally.
     add_swappables: null
     proposal: 'exhaustive'
     acceptance: 'metropolis' 
-    w_combine: null
-    rmse_cutoff: null
+    w_combine: False
+    w_mean_type: 'simple'
     N_cutoff: 1000
-    n_ex: 1
+    hist_corr: False
     mdp_args: null
     grompp_args: null
     runtime_args: null

docs/theory.rst

@@ -1,6 +1,8 @@
 .. _doc_basic_idea:
 
-.. note:: This page is still a work in progress. Please refer to our paper for more details before this page is completed.
+.. note:: This page is still a work in progress. Please check `Issue 33`_ for the current progress.
+
+.. _`Issue 33`: https://github.com/wehs7661/ensemble_md/issues/33
 
 1. Basic idea
 =============

ensemble_md/cli/run_REXEE.py

@@ -165,53 +165,81 @@ def main():
                 if wl_delta != [None for i in range(REXEE.n_sim)]:  # weight-updating
                     print(f'\nCurrent Wang-Landau incrementors: {wl_delta}\n')
 
-                # (1) First we prepare the weights to be combined.
+                # (1) First we prepare the time-averaged weights to be combined, if needed.
                 # Note that although averaged weights are sometimes used for weight correction/weight combination,
                 # the final weights are always used for calculating the acceptance ratio.
                 if REXEE.N_cutoff != -1 or REXEE.w_combine is not None:
                     # Only when weight correction/weight combination is needed.
                     weights_avg, weights_err = REXEE.get_averaged_weights(log_files)
-                    weights_input = REXEE.prepare_weights(weights_avg, weights)  # weights_input is for weight combination  # noqa: E501
+
+                    # Calculate the RMSE between the averaged weights and the final weights by the way.
+                    rmse_list = [utils.calc_rmse(weights_avg[i], weights[i]) for i in range(REXEE.n_sim)]
+                    rmse_str = ', '.join([f'{i:.2f}' for i in rmse_list])
+                    print(f'RMSE between the final weights and time-averaged weights for each replica: {rmse_str} kT')
 
                 # (2) Now we perform weight correction/weight combination.
                 # The product of this step should always be named as "weights" to be used in update_MDP
-                if REXEE.N_cutoff != -1 and REXEE.w_combine is not None:
-                    # perform both
-                    if weights_input is None:
-                        # Then only weight correction will be performed
-                        print('Note: Weight combination is deactivated because the weights are too noisy.')
-                        weights = REXEE.weight_correction(weights, counts)
-                        _ = REXEE.combine_weights(counts_, weights, print_values=False)[1]  # just to print the combiend weights  # noqa: E501
+                if REXEE.N_cutoff != -1 and REXEE.w_combine is True:
+                    # Perform both weight correction and weight combination
+                    if REXEE.verbose is True:
+                        print('Performing weight correction ...')
+                    else:
+                        print('Performing weight correction ...', end='')
+                    weights_preprocessed = REXEE.weight_correction(weights_avg, counts)
+
+                    if REXEE.verbose is True:
+                        print('Performing weight combination ...')
                     else:
-                        weights_preprocessed = REXEE.weight_correction(weights_input, counts)
-                        if REXEE.verbose is True:
-                            print('Performing weight combination ...')
-                        else:
-                            print('Performing weight combination ...', end='')
-                        counts, weights, g_vec = REXEE.combine_weights(counts_, weights_preprocessed)  # inverse-variance weighting seems worse  # noqa: E501
-                        REXEE.g_vecs.append(g_vec)
-                elif REXEE.N_cutoff == -1 and REXEE.w_combine is not None:
-                    # only perform weight combination
+                        print('Performing weight combination ...', end='')
+                    if REXEE.w_mean_type == 'simple':
+                        weights, g_vec = REXEE.combine_weights(weights_preprocessed)  # simple means
+                    else:
+                        # Note that here weights_err are acutally not the uncertainties for weights_prepocessed
+                        # but weights_avg ... We might need to disable this feature in the future.
+                        weights, g_vec = REXEE.combine_weights(weights_preprocessed, weights_err)  # inverse-variance weighting  # noqa: E501
+                    REXEE.g_vecs.append(g_vec)
+
+                    # Check if histogram correction is needed after weight combination
+                    if REXEE.hist_corr is True:
+                        print('Performing histogram correction ...')
+                        counts = REXEE.histogram_correction(counts_)
+                    else:
+                        print('Note: No histogram correction will be performed.')
+
+                elif REXEE.N_cutoff == -1 and REXEE.w_combine is True:
+                    # Only perform weight combination
                     print('Note: No weight correction will be performed.')
-                    if weights_input is None:
-                        print('Note: Weight combination is deactivated because the weights are too noisy.')
-                        _ = REXEE.combine_weights(counts_, weights, print_values=False)[1]  # just to print the combined weights  # noqa: E501
+                    if REXEE.verbose is True:
+                        print('Performing weight combination ...')
                     else:
-                        if REXEE.verbose is True:
-                            print('Performing weight combination ...')
-                        else:
-                            print('Performing weight combination ...', end='')
-                        counts, weights, g_vec = REXEE.combine_weights(counts_, weights_input)  # inverse-variance weighting seems worse  # noqa: E501
-                        REXEE.g_vecs.append(g_vec)
-                elif REXEE.N_cutoff != -1 and REXEE.w_combine is None:
-                    # only perform weight correction
+                        print('Performing weight combination ...', end='')
+                    if REXEE.w_mean_type == 'simple':
+                        weights, g_vec = REXEE.combine_weights(weights_avg)  # simple means
+                    else:
+                        weights, g_vec = REXEE.combine_weights(weights_avg, weights_err)  # inverse-variance weighting
+                    REXEE.g_vecs.append(g_vec)
+
+                    # Check if histogram correction is needed after weight combination
+                    if REXEE.hist_corr is True:
+                        print('Performing histogram correction ...')
+                        counts = REXEE.histogram_correction(counts_)
+                    else:
+                        print('Note: No histogram correction will be performed.')
+
+                elif REXEE.N_cutoff != -1 and REXEE.w_combine is False:
+                    # Only perform weight correction
                     print('Note: No weight combination will be performed.')
-                    weights = REXEE.weights_correction(weights_input, counts)
-                    _ = REXEE.combine_weights(counts_, weights, print_values=False)[1]  # just to print the combined weights  # noqa: E501
+                    if REXEE.verbose is True:
+                        print('Performing weight correction ...')
+                    else:
+                        print('Performing weight correction ...', end='')
+                    weights = REXEE.weights_correction(weights_avg, counts)
+                    _ = REXEE.combine_weights(weights, print_values=False)[1]  # just to print the combined weights  # noqa: E501
                 else:
                     print('Note: No weight correction will be performed.')
                     print('Note: No weight combination will be performed.')
-                    _ = REXEE.combine_weights(counts_, weights, print_values=False)[1]  # just to print the combiend weights  # noqa: E501
+                    # Note that in this case, the final weights will be used in the next iteration.
+                    _ = REXEE.combine_weights(weights, print_values=False)[1]  # just to print the combiend weights  # noqa: E501
 
                 # 3-5. Modify the MDP files and swap out the GRO files (if needed)
                 # Here we keep the lambda range set in mdp the same across different iterations in the same folder but swap out the gro file  # noqa: E501