Minor updates to flow validation

scripts/flow_validation.py

@@ -30,7 +30,7 @@ from numpyro.infer import MCMC, NUTS
 from taskmaster import work_delegation  # noqa
 
 
-def get_model(args, nsim):
+def get_model(args, nsim_iterator):
     """
     Load the data and create the NumPyro model.
 
@@ -38,8 +38,8 @@ def get_model(args, nsim):
     ----------
     args : argparse.Namespace
         Command line arguments.
-    nsim : int
-        Simulation index.
+    nsim_iterator : int
+        Simulation index, not the IC index. Ranges from 0, ... .
 
     Returns
     -------
@@ -58,8 +58,8 @@ def get_model(args, nsim):
     Omega_m = csiborgtools.simname2Omega_m(args.simname)
 
     # Read in the data from the loader.
-    los_overdensity = loader.los_density[:, nsim, :]
-    los_velocity = loader.los_radial_velocity[:, nsim, :]
+    los_overdensity = loader.los_density[:, nsim_iterator, :]
+    los_velocity = loader.los_radial_velocity[:, nsim_iterator, :]
 
     RA = loader.cat["RA"]
     dec = loader.cat["DEC"]
@@ -186,7 +186,7 @@ if __name__ == "__main__":
     nsims = paths.get_ics(args.simname)
 
     nsteps = 5000
-    nchains = 1
+    nchains = 4
 
     # Create the dumping folder.
     if comm.Get_rank() == 0:
@@ -198,12 +198,13 @@ if __name__ == "__main__":
         dump_folder = None
     dump_folder = comm.bcast(dump_folder, root=0)
 
-    def main(nsim):
-        model = get_model(args, nsim)
-        run_model(model, nsteps, nchains, nsim, dump_folder,
+    def main(i):
+        model = get_model(args, i)
+        run_model(model, nsteps, nchains, nsims[i], dump_folder,
                   show_progress=size == 1)
 
-    work_delegation(main, nsims, comm, master_verbose=True)
+    work_delegation(main, [i for i in range(len(nsims))], comm,
+                    master_verbose=True)
     comm.Barrier()
 
     if rank == 0:

scripts/flow_validation.sh

@@ -1,13 +1,14 @@
 memory=4
 on_login=${1}
 nthreads=${2}
+ksmooth=${3}
+
 queue="berg"
 env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
 file="flow_validation.py"
 
 catalogue="A2"
 simname="Carrick2015"
-ksmooth=2
 
 
 pythoncm="$env $file --catalogue $catalogue --simname $simname --ksmooth $ksmooth"