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/field_prop.py

@@ -193,7 +193,7 @@ def potential_field(nsim, parser_args, to_save=True):
 
     if parser_args.in_rsp:
         parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
-        field = csiborgtools.field.field2rsp(field, parts=parts, box=box,
+        field = csiborgtools.field.field2rsp(*field, parts=parts, box=box,
                                              verbose=parser_args.verbose)
     if to_save:
         fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
@@ -268,8 +268,6 @@ def environment_field(nsim, parser_args, to_save=True):
     -------
     env : 3-dimensional array
     """
-    if parser_args.in_rsp:
-        raise NotImplementedError("Env. field in RSP not implemented.")
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     nsnap = max(paths.get_snapshots(nsim))
     box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
@@ -292,7 +290,22 @@ def environment_field(nsim, parser_args, to_save=True):
     del rho
     collect()
 
-    # TODO: Optionally drag the field to RSP.
+    # Optionally drag the field to RSP.
+    if parser_args.in_rsp:
+        parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
+        fields = (tensor_field.T00, tensor_field.T11, tensor_field.T22,
+                  tensor_field.T01, tensor_field.T02, tensor_field.T12)
+
+        T00, T11, T22, T01, T02, T12 = csiborgtools.field.field2rsp(
+            *fields, parts=parts, box=box, verbose=parser_args.verbose)
+        tensor_field.T00[...] = T00
+        tensor_field.T11[...] = T11
+        tensor_field.T22[...] = T22
+        tensor_field.T01[...] = T01
+        tensor_field.T02[...] = T02
+        tensor_field.T12[...] = T12
+        del T00, T11, T22, T01, T02, T12
+        collect()
 
     # Calculate the eigenvalues of the tidal tensor field, delete tensor field.
     if parser_args.verbose:

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)