Moving environment to RS (#70)

* Drag several fields at once

* add env in RSP

* Add docstrings

* Make __main__

* Fix bug

* Fix plotting little bug

csiborgtools/field/interp.py

@@ -168,17 +168,17 @@ def divide_nonzero(field0, field1):
                     field0[i, j, k] /= field1[i, j, k]
 
 
-def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
+def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
               verbose=True):
     """
-    Forward model real space scalar field to redshift space. Attaches field
+    Forward model real space scalar fields to redshift space. Attaches field
     values to particles, those are then moved to redshift space and from their
     positions reconstructs back the field on a grid by NGP interpolation.
 
     Parameters
     ----------
-    field : 3-dimensional array of shape `(grid, grid, grid)`
-        Real space field to be evolved to redshift space.
+    fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
+        Real space fields to be evolved to redshift space.
     parts : 2-dimensional array of shape `(n_parts, 6)`
         Particle positions and velocities in real space. Must be organised as
         `x, y, z, vx, vy, vz`.
@@ -199,39 +199,50 @@ def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
     -------
     rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
     """
-    rsp_field = numpy.full(field.shape, init_value, dtype=numpy.float32)
-    cellcounts = numpy.zeros(rsp_field.shape, dtype=numpy.float32)
-    # We iterate over the fields and in the inner loop over the particles. This
-    # is slower than iterating over the particles and in the inner loop over
-    # the fields, but it is more memory efficient. Typically we will only have
-    # one field.
+    nfields = len(fields)
+    # Check that all fields have the same shape.
+    if nfields > 1:
+        assert all(fields[0].shape == fields[i].shape
+                   for i in range(1, nfields))
+
+    rsp_fields = [numpy.full(field.shape, init_value, dtype=numpy.float32)
+                  for field in fields]
+    cellcounts = numpy.zeros(rsp_fields[0].shape, dtype=numpy.float32)
+
     nparts = parts.shape[0]
     batch_size = nparts // nbatch
     start = 0
-    for k in trange(nbatch + 1) if verbose else range(nbatch + 1):
+    for __ in trange(nbatch + 1) if verbose else range(nbatch + 1):
+        # We first load the batch of particles into memory and flip x and z.
         end = min(start + batch_size, nparts)
         pos = parts[start:end]
         pos, vel = pos[:, :3], pos[:, 3:6]
         if flip_partsxz:
             pos[:, [0, 2]] = pos[:, [2, 0]]
             vel[:, [0, 2]] = vel[:, [2, 0]]
-        # Evaluate the field at the particle positions in real space.
-        values = evaluate_cartesian(field, pos=pos)
-        # Move particles to redshift space.
+        # Then move the particles to redshift space.
         rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
                                 in_box_units=True, periodic_wrap=True,
                                 make_copy=True)
-        # Assign particles' values to the grid.
-        MASL.MA(rsp_pos, rsp_field, 1., "NGP", W=values)
-        # Count the number of particles in each grid cell.
+        # ... and count the number of particles in each grid cell.
         MASL.MA(rsp_pos, cellcounts, 1., "NGP")
+
+        # Now finally we evaluate the field at the particle positions in real
+        # space and then assign the values to the grid in redshift space.
+        for i in range(nfields):
+            values = evaluate_cartesian(fields[i], pos=pos)
+            MASL.MA(rsp_pos, rsp_fields[i], 1., "NGP", W=values)
         if end == nparts:
             break
         start = end
 
-    # Finally divide by the number of particles in each cell and smooth.
-    divide_nonzero(rsp_field, cellcounts)
-    return rsp_field
+    # We divide by the number of particles in each cell.
+    for i in range(len(fields)):
+        divide_nonzero(rsp_fields[i], cellcounts)
+
+    if len(fields) == 1:
+        return rsp_fields[0]
+    return rsp_fields
 
 
 @jit(nopython=True)

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)

scripts_plots/plot_data.py

@@ -322,15 +322,24 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
                 ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
 
             frad = 155.5 / 677.7
-            if not highres_only and 0.5 - frad < slice_find < 0.5 + frad:
-                theta = numpy.linspace(0, 2 * numpy.pi, 100)
+            R = 155.5 / 677.7 * grid
+            if slice_find is None:
+                rad = R
+                plot_circle = True
+            elif (not highres_only and 0.5 - frad < slice_find < 0.5 + frad):
                 z = (slice_find - 0.5) * grid
-                R = 155.5 / 677.7 * grid
                 rad = R * numpy.sqrt(1 - z**2 / R**2)
+                plot_circle = True
+            else:
+                plot_circle = False
+
+            if not highres_only and plot_circle:
+                theta = numpy.linspace(0, 2 * numpy.pi, 100)
                 ax[i].plot(rad * numpy.cos(theta) + grid // 2,
                            rad * numpy.sin(theta) + grid // 2,
                            lw=0.75 * plt.rcParams["lines.linewidth"], zorder=1,
                            c="red", ls="--")
+
             ax[i].set_title(labels[i])
 
         if highres_only:
@@ -551,7 +560,7 @@ if __name__ == "__main__":
                               plot_halos=5e13, volume_weight=False)
 
     if True:
-        kind = "environment"
+        kind = "density"
         grid = 256
         smooth_scale = 0
         # plot_projected_field("overdensity", 7444, grid, in_rsp=True,

scripts_plots/plot_knn.py

@@ -41,6 +41,10 @@ def plot_knn(runname):
     Parameters
     ----------
     runname : str
+
+    Returns
+    -------
+    None
     """
     print(f"Plotting kNN CDF for {runname}.")
     cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]