The current DAG does not exactly implement the REXEE workflow implemented in `ensemble_md`

[weitse-hsu]

Currently, the DAG in replex.py does not seem to implement the REXEE workflow implemented in ensemble_md. Specifically, in the ensemble_md implementation, the workflow of a REXEE simulation is as follows:

• Step 1: Validate and print MDP/YAML parameters.
• Step 2: Take in the template MDP file and customize the MDP file for each replica using initialize_MDP
• Step 3: Run the first iteration using the GRO and TOP files in the home directory of the simulation, as well as the customized MDP files.
• Step 4: Parse the DHDL files, figure out MDP parameters to update in the next iteration, then customize MDP files for the next iteration using update_MDP.
• Step 5: Run another iteration using (1) the GRO files from the previous iteration (considering the swap pattern), (2) the TOP file in the home directory, and (3) the customized MDP files.
• Step 6: Repeat Steps 4 to 5 until the desired number of iteration is reached.

With #7, we have implemented all the methods required for Steps 2 to 5. However, to my understanding, the current DAG repeats Steps 2 to 5 instead of Steps 4 to 5 since TriggerDagRunOperator self-triggers the DAG that includes Steps 2 to 5. As what initialize_MDP does is pretty different from update_MDP, I wonder whether it would be beneficial to write another DAG specifically for iterations starting from the second iteration. That is, we can have the current DAG (with some modifications) trigger a child DAG that only repeats Steps 4 and 5 and triggers itself.

[ejjordan]

This sounds like a good approach to me. I am currently trying to get things ready to be able to test on actual HPC resources so it might be a little while before I can get to this. At any rate it should not be too much work since it is just moving things around that already exist.

I think this issue still has one question that I could use a little more clarification on though.

• How much work is outstanding for step 1 above, validating the MDP parameters?

[weitse-hsu]

Hi @ejjordan, so Step 1 mainly involves the functions set_params and print_params in replica_exchange_EE.py in ensemble_md.

• set_params: This function mainly checks if the YAML parameters (REXEE parameters specified in an input YAML file) are well-defined. Since currently all the REXEE parameters are hardcoded in replex.py, I think there might not be an urgent need to translate this function since we can easily make sure the hardcoded parameters are properly set. A major part of the parameter validation is for weight-updating REXEE simulations, so I think it would make sense to translate this function as we get to the point when we want to test weight-updating simulations. Currently the simulation method implemented in replex.py is the fixed-weight analog. The translation work should not be much.
• print_params: This function print parameters to run_REXEE_log.txt and can be easily translated to airflowHPC. One thing is that I am not sure if we would like to gather all the information in a file. In ensemble_md, everything is printed to run_REXEE_log.txt on the fly, including simulation parameters and what happened in each iteration. It could be pretty handy for a user to get a sense about what the simulation was doing after the simulation is done. However, since in the current airflowHPC implementation, such information is printed to the log file for each task, I am not sure if having such a text file would be standard. But if this is desired, I can help translate this function and a bunch of other print statements I have for logging information during the simulation.

[ejjordan]

It might be good to just put some calls to logging.info into the relevant functions in the replica exchange functions so that the important information gets to the task logs. Not having a sense for what can go wrong in such a simulation workflow I don't have a good idea how consolidated or not the log needs to be. But it is still nice to record relevant info that might have some debugging or sanity-checking utility.

[weitse-hsu]

It might be good to just put some calls to logging.info into the relevant functions in the replica exchange functions

Sounds good. I'll work on that soon.

[weitse-hsu]

@ejjordan, PR #10, which I have just opened, should address the issue.

I don't have a good idea how consolidated or not the log needs to be.

So in ensemble_md, the content of run_REXEE_log.txt loos something like below:

Current time: 19/06/2023 04:08:37
Command line: /jet/home/wehs7661/.local/python3.9/bin/run_EEXE

Important parameters of EXEE
============================
Python version: 3.9.12 (main, Feb 28 2023, 14:59:18) 
[GCC 10.2.0]
GROMACS executable: /jet/home/wehs7661/pkgs/gromacs/2022.5/bin/gmx
GROMACS version: 2022.5-dev-20230428-fdf57150ad
ensemble_md version: 0.6.0+40.gba0d9e6.dirty
Simulation inputs: anthracene.gro, anthracene.top, expanded.mdp
Verbose log file: True
Proposal scheme: exhaustive
Acceptance scheme for swapping simulations: metropolis
Whether to perform weight combination: False
Histogram cutoff: -1
Number of replicas: 4
Number of iterations: 12500
Number of attempted swaps in one exchange interval: N^3
Length of each replica: 4.0 ps
Frequency for checkpointing: 100 iterations
Total number of states: 8
Additional grompp arguments: None
Additional runtime arguments: {'-nt': '16', '-ntmpi': '1'}
Alchemical ranges of each replica in EEXE:
  - Replica 0: States [0, 1, 2, 3, 4]
  - Replica 1: States [1, 2, 3, 4, 5]
  - Replica 2: States [2, 3, 4, 5, 6]
  - Replica 3: States [3, 4, 5, 6, 7]

Iteration 0:  0.0 -  4.0 ps
===========================
Generating a TPR file on rank 0 ...
Running an EXE simulation on rank 0 ...

Below are the final states being visited:
  Simulation 0: Global state 1, (coul, vdw) =                 (0.18,)
  Simulation 1: Global state 2, (coul, vdw) =                 (0.42,)
  Simulation 2: Global state 3, (coul, vdw) =                 (0.57,)
  Simulation 3: Global state 4, (coul, vdw) =                 (0.68,)

Parsing sim_0/iteration_0/md.log ...
Parsing sim_1/iteration_0/md.log ...
Parsing sim_2/iteration_0/md.log ...
Parsing sim_3/iteration_0/md.log ...

Swappable pairs: [(0, 1), (1, 2), (2, 3)]

Proposed swap: (2, 3)
  Proposing a move from (x^i_m, x^j_n) to (x^i_n, x^j_m) ...
  U^i_n - U^i_m = 1.91 kT, U^j_m - U^j_n = 6.93 kT, Total dU: 8.84 kT
  g^i_n - g^i_m = 0.43 kT, g^j_m - g^j_n = -0.43 kT, Total dg: 0.00 kT
  Acceptance rate: 0.000 / Random number drawn: 0.212
  Swap rejected! 
  Current list of configurations: [0, 1, 2, 3]

Remaining swappable pairs: [(0, 1)]

Proposed swap: (0, 1)
  Proposing a move from (x^i_m, x^j_n) to (x^i_n, x^j_m) ...
  U^i_n - U^i_m = 5.12 kT, U^j_m - U^j_n = 1.53 kT, Total dU: 6.65 kT
  g^i_n - g^i_m = 4.94 kT, g^j_m - g^j_n = -4.94 kT, Total dg: 0.00 kT
  Acceptance rate: 0.001 / Random number drawn: 0.958
  Swap rejected! 
  Current list of configurations: [0, 1, 2, 3]

The finally adopted swap pattern: [0, 1, 2, 3]
The list of configurations sampled in each replica in the next iteration: [0, 1, 2, 3]

Note: No histogram correction will be performed.
Note: No weight combination will be performed.

Iteration 1:  4.0 -  8.0 ps
===========================
...

Iteration 12499:  49996.0 -  50000.0 ps
=======================================
Generating a TPR file on rank 0 ...
Running an EXE simulation on rank 0 ...

----- Saving .npy files to checkpoint the simulation ---

Summary of the simulation ensemble
==================================
Simulation status:
- Rep 0: The weights were fixed throughout the simulation.
- Rep 1: The weights were fixed throughout the simulation.
- Rep 2: The weights were fixed throughout the simulation.
- Rep 3: The weights were fixed throughout the simulation.

283 out of 12500, or 2.3% iterations had an empty list of swappable pairs.
10926 out of 16658, or 65.6% of attempted exchanges were rejected.

Time elapsed: 10 hour(s) 57 minute(s) 26 second(s)

By inserting calls of logging.info in different places, all information like this should be available in the log file of the corresponding task in a corresponding instance of DAG. However, I think it might still be useful to have a single file that gathers all useful information from all replicas so that when a user tries to look up some information about a certain iteration after the simulation is done, there is no need to identify the right task of the right DAG. If having a file like run_REXEE_log.txt makes sense, I can work on translating the relevant codes from ensemble_md to airflowHPC, preferentially in a separate PR though.