Moving environment to RS (#70)

* Drag several fields at once

* add env in RSP

* Add docstrings

* Make __main__

* Fix bug

* Fix plotting little bug

scripts/match_all.py

@@ -16,12 +16,14 @@ Script to match all pairs of CSiBORG simulations. Mathches main haloes whose
 mass is above 1e12 solar masses.
 """
 from argparse import ArgumentParser
-from datetime import datetime
 from distutils.util import strtobool
 from itertools import combinations
 from random import Random
 
 from mpi4py import MPI
+from taskmaster import work_delegation
+
+from match_singlematch import pair_match
 
 try:
     import csiborgtools
@@ -31,28 +33,16 @@ except ModuleNotFoundError:
     sys.path.append("../")
     import csiborgtools
 
-from taskmaster import master_process, worker_process
-
-from match_singlematch import pair_match
-
-# Argument parser
-parser = ArgumentParser()
-parser.add_argument("--sigma", type=float, default=None)
-parser.add_argument("--smoothen", type=lambda x: bool(strtobool(x)),
-                    default=None)
-parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
-                    default=False)
-args = parser.parse_args()
-
-comm = MPI.COMM_WORLD
-rank = comm.Get_rank()
-nproc = comm.Get_size()
-
 
 def get_combs():
     """
     Get the list of all pairs of simulations, then permute them with a known
     seed to minimise loading the same files simultaneously.
+
+    Returns
+    -------
+    combs : list
+        List of pairs of simulations.
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     ics = paths.get_ics("csiborg")
@@ -62,18 +52,31 @@ def get_combs():
 
 
 def do_work(comb):
+    """
+    Match a pair of simulations.
+
+    Parameters
+    ----------
+    comb : tuple
+        Pair of simulations.
+
+    Returns
+    -------
+    None
+    """
     nsim0, nsimx = comb
     pair_match(nsim0, nsimx, args.sigma, args.smoothen, args.verbose)
 
 
-if nproc > 1:
-    if rank == 0:
-        combs = get_combs()
-        master_process(combs, comm, verbose=True)
-    else:
-        worker_process(do_work, comm, verbose=False)
-else:
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--sigma", type=float, default=None)
+    parser.add_argument("--smoothen", type=lambda x: bool(strtobool(x)),
+                        default=None)
+    parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
+                        default=False)
+    args = parser.parse_args()
+    comm = MPI.COMM_WORLD
+
     combs = get_combs()
-    for comb in combs:
-        print(f"{datetime.now()}: completing task `{comb}`.", flush=True)
-        do_work(comb)
+    work_delegation(do_work, combs, comm, master_verbose=True)