Skip to content

Commit

Permalink
Removed prepare_weights and modified run_REXEE.py accordingly
Browse files Browse the repository at this point in the history
  • Loading branch information
@wehs7661
wehs7661 committed Nov 1, 2023
1 parent 16b9e34 commit 563e845
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 59 deletions.
55 changes: 25 additions & 30 deletions ensemble_md/cli/run_REXEE.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,52 +165,47 @@ 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
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='')
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 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='')
counts, weights, g_vec = REXEE.combine_weights(counts_, weights_avg)
REXEE.g_vecs.append(g_vec)
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)
weights = REXEE.weights_correction(weights_avg, counts)
_ = REXEE.combine_weights(counts_, 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.')
# Note that in this case, the final weights will be used in the next iteration.
_ = REXEE.combine_weights(counts_, 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)
Expand Down
29 changes: 0 additions & 29 deletions ensemble_md/replica_exchange_EE.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1149,35 +1149,6 @@ def get_averaged_weights(self, log_files):

return weights_avg, weights_err

def prepare_weights(self, weights_avg, weights_final):
"""
Prepared weights to be combined by the function :code:`combine_weights`.
Parameters
----------
weights_avg : list
A list of lists of weights averaged since the last update of the Wang-Landau
incrementor. The length of the list should be the number of replicas.
weights_final : list
A list of lists of final weights of all simulations. The length of the list should
be the number of replicas.
Returns
-------
weights_output : list
A list of lists of weights to be combined.
"""
rmse_list = [utils.calc_rmse(weights_avg[i], weights_final[i]) for i in range(self.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')

if self.w_combine is True:
weights_output = weights_avg
else:
weights_output = None

return weights_output

def combine_weights(self, hist, weights, weights_err=None, print_values=True):
"""
Combine alchemical weights across multiple replicas and adjusts the histogram counts
Expand Down

0 comments on commit 563e845

Please sign in to comment.