Fix a bug in the OXPY_REMD/reshuffle.py script, and strengthen the re…

…md.py script

examples/OXPY_REMD/README.md

@@ -39,7 +39,7 @@ To run a REMD simulation execute the following commands:
 ```shell
 ./clean.sh
 
-mpirun -n 4  python3  remd.py  input_md
+mpirun -n 4 python3 remd.py input_md
 
 python3 reshuffle.py input_md history.json reshuffled
 ```

examples/OXPY_REMD/remd.py

@@ -8,6 +8,8 @@
 from mpi4py import MPI
 from random import random
 
+import argparse
+
 import oxpy
 
 try:
@@ -33,42 +35,61 @@ def get_order(i, rank, locations, mpi_nprocs):
         return im_responsible, im_rank, talk_to 
     else:
         return None    
+
+
+def remd_log(pid, *args):
+    if pid == rank:
+        print("REMD_LOG:", pid, ":", *args)
 
 
 comm = MPI.COMM_WORLD
-rank = comm.Get_rank()  # process id  
-mpi_nprocs = comm.Get_size()  # number of processes running 
+rank = comm.Get_rank()  # process id
+mpi_nprocs = comm.Get_size()  # number of processes running
 locations = list(range(mpi_nprocs))  # curent location of the exchange temperature
 exchange = np.zeros(mpi_nprocs)  # keep track of exchanges
 exchange_tries = np.zeros(mpi_nprocs)  # and attempts
 rates = np.zeros(mpi_nprocs)
 
-if rank == 0:
-    history = []
-
+parser = argparse.ArgumentParser(description="Run an REMD simulation.")
+parser.add_argument('input_file', type=str, nargs=1, help="The input file of the simulation")
+args = None
+try:
+    if rank == 0:
+        args = parser.parse_args()
+finally:
+    args = comm.bcast(args, root=0)
 
-def remd_log(pid, *args):
-    if pid == rank:
-        print("x", pid, ":", *args)
+if args == None:
+    exit(1)
+else:
+    input_file = args.input_file[0]
 
+if rank == 0:
+    history = []
 
 with oxpy.Context():
     input = oxpy.InputFile()  # load up conf and input
-    input.init_from_filename(sys.argv[1])  # as 1st is the name of the script
+    input.init_from_filename(input_file)  # as 1st is the name of the script
+
     # have separate files for every running instance
-
     input["energy_file"] = f"energy_{rank}.dat"
     input["lastconf_file"] = f"last_conf_{rank}.dat"
     input["CUDA_device"] = str(rank)
     # figure out steps to run and our temperature
     pt_move_every = int(input["pt_move_every"])
-    steps_to_run = int(int(input["steps"]) / pt_move_every) 
+    steps_to_run = int(int(input["steps"]) / pt_move_every)
     # setup temperatures
     temperatures = input["pt_temp_list"].split(",")
 
     input["T"] = temperatures[rank]  # set process temperature
     input["trajectory_file"] = f"trajectory_{rank}.dat"
 
+    # check that the number of processes is equal to the number of temperatures
+    if rank == 0:
+        if len(temperatures) != mpi_nprocs:
+            print(f"The number of temperatures ({len(temperatures)}) should match the number of MPI processes ({mpi_nprocs})", file=sys.stderr)
+            comm.Abort(errorcode=1)
+
     # convert temperaures to ox units
     if "K" in input["T"]:
         temperatures = [oxpy.utils.Kelvin_to_oxDNA(t.replace("K", "")) for t in temperatures]   
@@ -86,10 +107,10 @@ def remd_log(pid, *args):
         locations = comm.bcast(locations, root=0)
         # let's wait for everyone to get started 
         swap = False
-        resut = get_order(i, rank , locations, mpi_nprocs)
-        if resut:
+        result = get_order(i, rank , locations, mpi_nprocs)
+        if result:
             # unpack result
-            im_responsible, im_rank, talk_to = resut
+            im_responsible, im_rank, talk_to = result
             if im_responsible:
                 irresp_energy = comm.recv(source=talk_to)
 
@@ -127,32 +148,32 @@ def remd_log(pid, *args):
 
         # handle updates to locations
         if rank == 0:
-            swaped = [] 
+            swapped = [] 
             # handle 0
-            if  resut and im_responsible:
+            if  result and im_responsible:
                 exchange_tries[0] += 1
                 exchange_tries[talk_to] += 1 
                 if swap:
                     locations[0] , locations[talk_to] = locations[talk_to] , locations[0]
-                    swaped.append(0)
-                    swaped.append(talk_to)    
+                    swapped.append(0)
+                    swapped.append(talk_to)    
                     exchange[0] += 1
                     exchange[talk_to] += 1
                     rates [0] = exchange[0] / exchange_tries[0]
                     rates [talk_to] = exchange[talk_to] / exchange_tries[talk_to]
 
             # receive from everyone but me the updates
             for j in range(1, mpi_nprocs):
-                resut = comm.recv(source=j)
-                if resut:
-                    swap, im_responsible, im_rank, talk_to = resut
+                result = comm.recv(source=j)
+                if result:
+                    swap, im_responsible, im_rank, talk_to = result
                     if im_responsible:
                         exchange_tries[j] += 1
                         exchange_tries[talk_to] += 1
                         if swap: 
                             locations[j] , locations[talk_to] = locations[talk_to] , locations[j]
-                            swaped.append(j)
-                            swaped.append(talk_to)
+                            swapped.append(j)
+                            swapped.append(talk_to)
                             exchange[j] += 1
                             exchange[talk_to] += 1
                             rates [j] = exchange[0] / exchange_tries[j]
@@ -166,7 +187,7 @@ def remd_log(pid, *args):
             remd_log(0, "rates:", rates)
         else:
             # notify 0 of update
-            if resut:
+            if result:
                 comm.send((swap, im_responsible, im_rank, talk_to), 0)
             else:
                 comm.send(None, 0)  # nothing happened for this round

examples/OXPY_REMD/reshuffle.py

@@ -1,22 +1,22 @@
 # reintegrate trajectory
 from json import loads
 from oxDNA_analysis_tools.UTILS.RyeReader import describe, get_confs, get_input_parameter, write_conf
+import shutil, sys
 
 import argparse
 
 if __name__ == "__main__":
-    #handle commandline arguments
-    parser = argparse.ArgumentParser(description="Retrive temperature trajectories from an REMD run.")
-    parser.add_argument('input_file', type=str, nargs=1, help="The inputfile used to run the simulation")
-    parser.add_argument('history_file', type=str, nargs=1, help="The json history file keeping the remd run information.")
-    parser.add_argument('output_file', type=str, nargs=1, help="name of the file to save the output trajectories to")
+    # handle commandline arguments
+    parser = argparse.ArgumentParser(description="Retrieve temperature trajectories and last configurations from an REMD run.")
+    parser.add_argument('input_file', type=str, nargs=1, help="The input file used to run the simulation")
+    parser.add_argument('history_file', type=str, nargs=1, help="The json history file keeping the REMD run information.")
+    parser.add_argument('output_file', type=str, nargs=1, help="Prefix used to build the name of the file to save the output trajectories to")
 
     args = parser.parse_args()
     input_file = args.input_file[0]
     history_file = args.history_file[0]
     out_file = args.output_file[0]
 
-
     # parse the input file to retrieve the temperature settings
     pt_temp_list = None
     with open(input_file) as file:
@@ -29,21 +29,25 @@
         history = loads(file.read())
 
     top_name = get_input_parameter(input_file, 'topology')
-    traj_name = "trajectory"#get_input_parameter(input_file, 'trajectory_file')
-    #traj_base = traj_name.split('.')[:-1]
-    traj_ext = "dat"#traj_name.split('.')[-1]
+    traj_name = "trajectory"
+    traj_ext = "dat"
     tis = []
     dis = []
     for i in range(len(pt_temp_list)):
-        ti, di = describe(top_name, f'{traj_name}_{i}.{traj_ext}')
+        ti, di = describe(top_name, f"trajectory_{i}.dat")
         tis.append(ti)
         dis.append(di)
-
-    out_files = [f"./{out_file}_{T}.dat" for T in pt_temp_list]
 
+    # generate the ordered trajectories        
+    out_files = [f"./{out_file}_{T}.dat" for T in pt_temp_list]
     read_poses = [0 for _ in range(len(pt_temp_list))]
-    for i,locations in enumerate(history):
-        for j,l in enumerate(locations):
-            configuration = get_confs(tis[l], dis[l], read_poses[l], 1)[0]
-            read_poses[l] += 1
+    for i, locations in enumerate(history):
+        for j, l in enumerate(locations):
+            configuration = get_confs(tis[j], dis[j], read_poses[j], 1)[0]
+            read_poses[j] += 1
             write_conf(out_files[l], configuration, append=True, include_vel=dis[0].incl_v)
+
+    # print the last_conf files
+    last_files = [f"./last_conf_{T}.dat" for T in pt_temp_list]
+    for i, l in enumerate(history[-1]):
+        shutil.copy(f"last_conf_{i}.dat", last_files[l])