From 21139efa4174018213774c0a580ef83b0a3fc4f2 Mon Sep 17 00:00:00 2001 From: ajfriedman22 Date: Tue, 15 Oct 2024 12:22:26 -0600 Subject: [PATCH] Add weights to checkpoint files for variable weight REXEE --- ensemble_md/cli/run_REXEE.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/ensemble_md/cli/run_REXEE.py b/ensemble_md/cli/run_REXEE.py index 7f53788..7ca4b89 100644 --- a/ensemble_md/cli/run_REXEE.py +++ b/ensemble_md/cli/run_REXEE.py @@ -44,6 +44,12 @@ def initialize(args): help='The file path of the NPY file containing the time series of the whole-range\ alchemical weights. This file is a necessary input if one wants to update the \ file when extending a weight-updating simulation. (Default: g_vecs.npy)') + parser.add_argument('-e', + '--equil', + type=str, + default='equil.npy', + help='The file path of the NPY file containing the equilibration times for all simulations \ + when completing a variable weight REXEE simulation. (Default: equil.npy)') parser.add_argument('-o', '--output', type=str, @@ -124,11 +130,15 @@ def main(): shutil.rmtree(f'{REXEE.working_dir}/sim_{i}/iteration_{j}') # Read g_vecs.npy and rep_trajs.npy so that new data can be appended, if any. + # Read equil.npy if running variable weight REXEE simulations # Note that these two arrays are created in rank 0 and should always be operated in rank 0, # or broadcasting is required. REXEE.rep_trajs = [list(i) for i in ckpt_data] if os.path.isfile(args.g_vecs) is True: REXEE.g_vecs = [list(i) for i in np.load(args.g_vecs)] + if REXEE.fixed_weights is not True and os.path.isfile(args.equil) is True: + REXEE.equil = np.load(args.equil) + print(REXEE.equil) else: start_idx = None @@ -322,16 +332,20 @@ def main(): if (i + 1) % REXEE.n_ckpt == 0: if len(REXEE.g_vecs) != 0: # Save g_vec as a function of time if weight combination was used. - np.save('g_vecs.npy', REXEE.g_vecs) + np.save(args.g_vecs, REXEE.g_vecs) print('\n----- Saving .npy files to checkpoint the simulation ---') - np.save('rep_trajs.npy', REXEE.rep_trajs) + np.save(args.ckpt, REXEE.rep_trajs) + if REXEE.fixed_weights is not True: + np.save(args.equil, REXEE.equil) # Save the npy files at the end of the simulation anyway. if rank == 0: if len(REXEE.g_vecs) != 0: # The length will be 0 only if there is no weight combination. - np.save('g_vecs.npy', REXEE.g_vecs) - np.save('rep_trajs.npy', REXEE.rep_trajs) + np.save(args.g_vecs, REXEE.g_vecs) + np.save(args.ckpt, REXEE.rep_trajs) + if REXEE.fixed_weights is not True: + np.save(args.equil, REXEE.equil) # Step 5: Write a summary for the simulation ensemble if rank == 0: