Tweaked combinne_weights and modified run_EEXE.py

wehs7661 · Aug 29, 2023 · 9ec50c6 · 9ec50c6
1 parent 453e059
commit 9ec50c6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/ensemble_md/cli/run_EEXE.py b/ensemble_md/cli/run_EEXE.py
@@ -180,6 +180,8 @@ def main():
                         # Then only histogram correction will be performed
                         print('Note: Weight combination is deactivated because the weights are too noisy.')
                         weights = EEXE.histogram_correction(weights, counts)
+                        w_for_printing = EEXE.combine_weights(weights, print_weights=False)[1]
+                        print(f'The alchemical weights of all states (after histogram correction): \n  {list(np.round(w_for_printing, decimals=3))}')  # noqa: E501
                     else:
                         weights_preprocessed = EEXE.histogram_correction(weights_input, counts)
                         weights, g_vec = EEXE.combine_weights(weights_preprocessed)  # inverse-variance weighting seems worse  # noqa: E501
@@ -189,15 +191,19 @@ def main():
                     print('Note: No histogram correction will be performed.')
                     if weights_input is None:
                         print('Note: Weight combination is deactivated because the weights are too noisy.')
+                        w_for_printing = EEXE.combine_weights(weights, print_weights=False)[1]
+                        print(f'The alchemical weights of all states: \n  {list(np.round(w_for_printing, decimals=3))}')  # noqa: E501
                     else:
                         weights, g_vec = EEXE.combine_weights(weights_input)  # inverse-variance weighting seems worse
                         EEXE.g_vecs.append(g_vec)
                 elif EEXE.N_cutoff != -1 and EEXE.w_combine is None:
                     # only perform histogram correction
                     print('Note: No weight combination will be performed.')
                     weights = EEXE.histogram_correction(weights_input, counts)
+                    w_for_printing = EEXE.combine_weights(weights, print_weights=False)[1]
+                    print(f'The alchemical weights of all states: \n  {list(np.round(w_for_printing, decimals=3))}')  # noqa: E501
                 else:
-                    w_for_printing = EEXE.combine_weights(weights)[1]
+                    w_for_printing = EEXE.combine_weights(weights, print_weights=False)[1]
                     print('Note: No histogram correction will be performed.')
                     print('Note: No weight combination will be performed.')
                     print(f'The alchemical weights of all states: \n  {list(np.round(w_for_printing, decimals=3))}')

diff --git a/ensemble_md/ensemble_EXE.py b/ensemble_md/ensemble_EXE.py
@@ -1216,7 +1216,7 @@ def prepare_weights(self, weights_avg, weights_final):
 
         return weights_output
 
-    def combine_weights(self, weights, weights_err=None):
+    def combine_weights(self, weights, weights_err=None, print_weights=True):
         """
         Combine alchemical weights across multiple replicas. Note that if
         :code:`weights_err` is provided, inverse-variance weighting will be used.
@@ -1234,13 +1234,15 @@ def combine_weights(self, weights, weights_err=None):
             A list of modified Wang-Landau weights of ALL simulations.
         g_vec : np.ndarray
             An array of alchemical weights of the whole range of states.
+        print_weights : bool
+            Whether to print the original and combined weights for each replica.
         """
         if self.verbose is True:
             print('Performing weight combination ...')
         else:
             print('Performing weight combination ...', end='')
 
-        if self.verbose is True:
+        if print_weights is True:
             w = np.round(weights, decimals=3).tolist()  # just for printing
             print('  Original weights:')
             for i in range(len(w)):
@@ -1275,7 +1277,7 @@ def combine_weights(self, weights, weights_err=None):
                 weights[i] = self.equilibrated_weights[i]
         weights = np.round(weights, decimals=5).tolist()
 
-        if self.verbose is True:
+        if print_weights is True:
             w = np.round(weights, decimals=3).tolist()  # just for printing
             print('\n  Combined weights:')
             for i in range(len(w)):