Forward model fields to RSP (#66)

* Add radial velocity field

* Add overdensity plot

* Flip velocities too

* Add field calculations

* Add RSP mapping

* Add potential in RSP

* Add projected field plotting

csiborgtools/field/__init__.py

@@ -18,7 +18,8 @@ try:
     import MAS_library as MASL  # noqa
 
     from .density import DensityField, PotentialField, VelocityField  # noqa
-    from .interp import evaluate_cartesian, evaluate_sky, make_sky  # noqa
+    from .interp import (evaluate_cartesian, evaluate_sky, field2rsp,  # noqa
+                         make_sky)
     from .utils import nside2radec, smoothen_field  # noqa
 except ImportError:
     warn("MAS_library not found, `DensityField` will not be available", UserWarning)  # noqa

csiborgtools/field/density.py

@@ -20,8 +20,10 @@ TODO:
 """
 from abc import ABC
 
-import MAS_library as MASL
 import numpy
+
+import MAS_library as MASL
+from numba import jit
 from tqdm import trange
 
 from ..read.utils import real2redshift
@@ -220,6 +222,36 @@ class VelocityField(BaseField):
         self.box = box
         self.MAS = MAS
 
+    @staticmethod
+    @jit(nopython=True)
+    def radial_velocity(rho_vel):
+        """
+        Calculate the radial velocity field around the observer in the centre
+        of the box.
+
+        Parameters
+        ----------
+        rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
+            Velocity field along each axis.
+
+        Returns
+        -------
+        radvel : 3-dimensional array of shape `(grid, grid, grid)`.
+            Radial velocity field.
+        """
+        grid = rho_vel.shape[1]
+        radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
+        for i in range(grid):
+            px = i - 0.5 * (grid - 1)
+            for j in range(grid):
+                py = j - 0.5 * (grid - 1)
+                for k in range(grid):
+                    pz = k - 0.5 * (grid - 1)
+                    vx, vy, vz = rho_vel[:, i, j, k]
+                    radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
+                                       / numpy.sqrt(px**2 + py**2 + pz**2))
+        return radvel
+
     def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30,
                  verbose=True):
         """
@@ -245,7 +277,7 @@ class VelocityField(BaseField):
 
         Returns
         -------
-        rho_vel : 3-dimensional array of shape `(3, grid, grid, grid)`.
+        rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
             Velocity field along each axis.
 
         References
@@ -272,6 +304,7 @@ class VelocityField(BaseField):
             mass = force_single_precision(mass, "particle_mass")
             if flip_xz:
                 pos[:, [0, 2]] = pos[:, [2, 0]]
+                vel[:, [0, 2]] = vel[:, [2, 0]]
             vel *= mass.reshape(-1, 1) / mpart
 
             for i in range(3):

csiborgtools/field/interp.py

@@ -17,9 +17,11 @@ Tools for interpolating 3D fields at arbitrary positions.
 """
 import MAS_library as MASL
 import numpy
+from numba import jit
+from scipy.ndimage import gaussian_filter
 from tqdm import trange
 
-from ..read.utils import radec_to_cartesian
+from ..read.utils import radec_to_cartesian, real2redshift
 from .utils import force_single_precision
 
 
@@ -137,3 +139,99 @@ def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True):
                 evaluate_sky(field, pos=dir_loop, box=box, isdeg=True))
     out *= dx
     return out
+
+
+@jit(nopython=True)
+def divide_nonzero(field0, field1):
+    """
+    Divide two fields where the second one is not zero. If the second field
+    is zero, the first one is left unchanged. Operates in-place.
+
+    Parameters
+    ----------
+    field0 : 3-dimensional array of shape `(grid, grid, grid)`
+        Field to be divided.
+    field1 : 3-dimensional array of shape `(grid, grid, grid)`
+        Field to divide by.
+
+    Returns
+    -------
+    field0 : 3-dimensional array of shape `(grid, grid, grid)`
+        Field divided by the second one.
+    """
+    assert field0.shape == field1.shape
+
+    imax, jmax, kmax = field0.shape
+    for i in range(imax):
+        for j in range(jmax):
+            for k in range(kmax):
+                if field1[i, j, k] != 0:
+                    field0[i, j, k] /= field1[i, j, k]
+
+
+def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
+              verbose=True):
+    """
+    Forward model real space scalar field 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 regular grid by NGP
+    interpolation. This by definition produces a discontinuous field.
+
+    Parameters
+    ----------
+    field : 3-dimensional array of shape `(grid, grid, grid)`
+        Real space field to be evolved to redshift space.
+    parts_pos : 2-dimensional array of shape `(n_parts, 3)`
+        Particle positions in real space.
+    parts_vel : 2-dimensional array of shape `(n_parts, 3)`
+        Particle velocities in real space.
+    box : :py:class:`csiborgtools.read.CSiBORGBox`
+        The simulation box information and transformations.
+    nbatch : int, optional
+        Number of batches to use when moving particles to redshift space.
+        Particles are assumed to be lazily loaded to memory.
+    flip_partsxz : bool, optional
+        Whether to flip the x and z coordinates of the particles. This is
+        because of a BORG bug.
+    init_value : float, optional
+        Initial value of the RSP field on the grid.
+    verbose : bool, optional
+        Verbosity flag.
+
+    Returns
+    -------
+    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.
+    nparts = parts.shape[0]
+    batch_size = nparts // nbatch
+    start = 0
+    for k in trange(nbatch + 1) if verbose else range(nbatch + 1):
+        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.
+        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.
+        MASL.MA(rsp_pos, cellcounts, 1., "NGP")
+        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

csiborgtools/read/paths.py

@@ -363,7 +363,8 @@ class Paths:
         Parameters
         ----------
         kind : str
-            Field type. Must be one of: `density`, `velocity`, `potential`.
+            Field type. Must be one of: `density`, `velocity`, `potential`,
+            `radvel`.
         MAS : str
            Mass-assignment scheme.
         grid : int
@@ -378,7 +379,7 @@ class Paths:
         path : str
         """
         fdir = join(self.postdir, "environment")
-        assert kind in ["density", "velocity", "potential"]
+        assert kind in ["density", "velocity", "potential", "radvel"]
         if not isdir(fdir):
             makedirs(fdir)
             warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)

scripts/field_derived.py

@@ -1,78 +0,0 @@
-# Copyright (C) 2022 Richard Stiskalek
-# This program is free software; you can redistribute it and/or modify it
-# under the terms of the GNU General Public License as published by the
-# Free Software Foundation; either version 3 of the License, or (at your
-# option) any later version.
-#
-# This program is distributed in the hope that it will be useful, but
-# WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
-# Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""
-MPI script to calculate density field-derived fields in the CSiBORG
-simulations' final snapshot.
-"""
-from argparse import ArgumentParser
-from datetime import datetime
-from distutils.util import strtobool
-
-import numpy
-from mpi4py import MPI
-
-try:
-    import csiborgtools
-except ModuleNotFoundError:
-    import sys
-    sys.path.append("../")
-    import csiborgtools
-
-comm = MPI.COMM_WORLD
-rank = comm.Get_rank()
-nproc = comm.Get_size()
-verbose = nproc == 1
-
-parser = ArgumentParser()
-parser.add_argument("--ics", type=int, nargs="+", default=None,
-                    help="IC realisations. If `-1` processes all simulations.")
-parser.add_argument("--kind", type=str, choices=["potential", "velocity"],
-                    help="What derived field to calculate?")
-parser.add_argument("--MAS", type=str, choices=["NGP", "CIC", "TSC", "PCS"])
-parser.add_argument("--grid", type=int, help="Grid resolution.")
-parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
-                    help="Calculate from the RSP density field?")
-args = parser.parse_args()
-paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-
-if args.ics is None or args.ics[0] == -1:
-    ics = paths.get_ics("csiborg")
-else:
-    ics = args.ics
-
-
-for i in csiborgtools.fits.split_jobs(len(ics), nproc)[rank]:
-    nsim = ics[i]
-    if verbose:
-        print(f"{datetime.now()}: rank {rank} working on simulation {nsim}.",
-              flush=True)
-    nsnap = max(paths.get_snapshots(nsim))
-    box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
-    density_gen = csiborgtools.field.DensityField(box, args.MAS)
-
-    rho = numpy.load(paths.field("density", args.MAS, args.grid, nsim,
-                                      args.in_rsp))
-    rho = density_gen.overdensity_field(rho)
-
-    if args.kind == "potential":
-        gen = csiborgtools.field.PotentialField(box, args.MAS)
-    else:
-        raise RuntimeError(f"Field {args.kind} is not implemented yet.")
-
-    field = gen(rho)
-    fout = paths.field("potential", args.MAS, args.grid, nsim,
-                            args.in_rsp)
-    print(f"{datetime.now()}: rank {rank} saving output to `{fout}`.")
-    numpy.save(fout, field)

scripts/field_prop.py

@@ -0,0 +1,169 @@
+# Copyright (C) 2022 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+MPI script to calculate density field-derived fields in the CSiBORG
+simulations' final snapshot.
+"""
+from argparse import ArgumentParser
+from datetime import datetime
+from distutils.util import strtobool
+
+import numpy
+from mpi4py import MPI
+
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
+    import csiborgtools
+
+from taskmaster import work_delegation
+
+from utils import get_nsims
+
+###############################################################################
+#                            Density field                                    #
+###############################################################################
+
+
+def density_field(nsim, parser_args):
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsnap = max(paths.get_snapshots(nsim))
+    box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
+    parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
+
+    gen = csiborgtools.field.DensityField(box, parser_args.MAS)
+    field = gen(parts, parser_args.grid, in_rsp=parser_args.in_rsp,
+                verbose=parser_args.verbose)
+
+    fout = paths.field("density", parser_args.MAS, parser_args.grid,
+                       nsim, parser_args.in_rsp)
+    print(f"{datetime.now()}: saving output to `{fout}`.")
+    numpy.save(fout, field)
+
+
+###############################################################################
+#                            Velocity field                                   #
+###############################################################################
+
+
+def velocity_field(nsim, parser_args):
+    if parser_args.in_rsp:
+        raise NotImplementedError("Velocity field in RSP is not implemented.")
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    mpart = 1.1641532e-10  # Particle mass in CSiBORG simulations.
+    nsnap = max(paths.get_snapshots(nsim))
+    box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
+    parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
+
+    gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
+    field = gen(parts, parser_args.grid, mpart, verbose=parser_args.verbose)
+
+    fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
+                       nsim, in_rsp=False)
+    print(f"{datetime.now()}: saving output to `{fout}`.")
+    numpy.save(fout, field)
+
+
+###############################################################################
+#                          Potential field                                    #
+###############################################################################
+
+
+def potential_field(nsim, parser_args):
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsnap = max(paths.get_snapshots(nsim))
+    box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
+
+    # Load the real space overdensity field
+    density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
+    rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
+                                 nsim, in_rsp=False))
+    rho = density_gen.overdensity_field(rho)
+    # Calculate the real space potentiel field
+    gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
+    field = gen(rho)
+
+    if parser_args.in_rsp:
+        parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
+        field = csiborgtools.field.field2rsp(field, parts=parts, box=box,
+                                             verbose=parser_args.verbose)
+    fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
+                       nsim, parser_args.in_rsp)
+    print(f"{datetime.now()}: saving output to `{fout}`.")
+    numpy.save(fout, field)
+
+
+###############################################################################
+#                          Radial velocity field                              #
+###############################################################################
+
+
+def radvel_field(nsim, parser_args):
+    if parser_args.in_rsp:
+        raise NotImplementedError("Radial vel. 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)
+
+    vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid,
+                                 nsim, parser_args.in_rsp))
+    gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
+    field = gen.radial_velocity(vel)
+
+    fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
+                       nsim, parser_args.in_rsp)
+    print(f"{datetime.now()}: saving output to `{fout}`.")
+    numpy.save(fout, field)
+
+
+###############################################################################
+#                          Command line interface                             #
+###############################################################################
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--nsims", type=int, nargs="+", default=None,
+                        help="IC realisations. `-1` for all simulations.")
+    parser.add_argument("--kind", type=str,
+                        choices=["density", "velocity", "radvel", "potential"],
+                        help="What derived field to calculate?")
+    parser.add_argument("--MAS", type=str,
+                        choices=["NGP", "CIC", "TSC", "PCS"])
+    parser.add_argument("--grid", type=int, help="Grid resolution.")
+    parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
+                        help="Calculate in RSP?")
+    parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
+                        help="Verbosity flag for reading in particles.")
+    parser_args = parser.parse_args()
+    comm = MPI.COMM_WORLD
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(parser_args, paths)
+
+    def main(nsim):
+        if parser_args.kind == "density":
+            density_field(nsim, parser_args)
+        elif parser_args.kind == "velocity":
+            velocity_field(nsim, parser_args)
+        elif parser_args.kind == "radvel":
+            radvel_field(nsim, parser_args)
+        elif parser_args.kind == "potential":
+            potential_field(nsim, parser_args)
+        else:
+            raise RuntimeError(f"Field {parser_args.kind} is not implemented.")
+
+    work_delegation(main, nsims, comm, master_verbose=True)

scripts_plots/plot_data.py

@@ -162,29 +162,73 @@ def plot_hmf(pdf=False):
         plt.close()
 
 
-###############################################################################
-#                           Sky distribution                                  #
-###############################################################################
-
 @cache_to_disk(7)
 def load_field(kind, nsim, grid, MAS, in_rsp=False):
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    print(paths.field(kind, MAS, grid, nsim, in_rsp=in_rsp))
     return numpy.load(paths.field(kind, MAS, grid, nsim, in_rsp=in_rsp))
 
 
+###############################################################################
+#                             Projected field                                 #
+###############################################################################
+
+
+def plot_projected_field(kind, nsim, grid, in_rsp, MAS="PCS", pdf=False):
+    print(f"Plotting projected field `{kind}`. ", flush=True)
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsnap = max(paths.get_snapshots(nsim))
+    box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
+
+    if kind == "overdensity":
+        field = load_field("density", nsim, grid, MAS=MAS, in_rsp=in_rsp)
+        density_gen = csiborgtools.field.DensityField(box, MAS)
+        field = density_gen.overdensity_field(field) + 2
+    else:
+        field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp)
+
+    print(field)
+
+    with plt.style.context(utils.mplstyle):
+        fig, ax = plt.subplots(figsize=(3.5 * 2, 2.625), ncols=3, sharey=True,
+                               sharex=True)
+        fig.subplots_adjust(hspace=0, wspace=0)
+        for i in range(3):
+            ax[i].imshow(numpy.sum(field, axis=i))
+
+        fig.tight_layout(h_pad=0, w_pad=0)
+        for ext in ["png"] if pdf is False else ["png", "pdf"]:
+            fout = join(utils.fout, f"field_{kind}_{nsim}_rsp{in_rsp}.{ext}")
+            print(f"Saving to `{fout}`.")
+            fig.savefig(fout, dpi=utils.dpi, bbox_inches="tight")
+        plt.close()
+
+###############################################################################
+#                             Sky distribution                                #
+###############################################################################
+
+
 def get_sky_label(kind, volume_weight):
     if volume_weight:
         if kind == "density":
-            label = r"$\log \int_{0}^{R} r^2 \delta(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
+            label = r"$\log \int_{0}^{R} r^2 \rho(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
+        if kind == "overdensity":
+            label = r"$\log \int_{0}^{R} r^2 \left[\delta(r, \mathrm{RA}, \mathrm{dec}) + 2\right] \mathrm{d} r$"  # noqa
         elif kind == "potential":
             label = r"$\int_{0}^{R} r^2 \phi(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
+        elif kind == "radvel":
+            label = r"$\int_{0}^{R} r^2 v_r(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
         else:
             label = None
     else:
         if kind == "density":
-            label = r"$\log \int_{0}^{R} \delta(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
+            label = r"$\log \int_{0}^{R} \rho(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
+        if kind == "overdensity":
+            label = r"$\log \int_{0}^{R} \left[\delta(r, \mathrm{RA}, \mathrm{dec}) + 2\right] \mathrm{d} r$"  # noqa
         elif kind == "potential":
             label = r"$\int_{0}^{R} \phi(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
+        elif kind == "radvel":
+            label = r"$\int_{0}^{R} v_r(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$"  # noqa
         else:
             label = None
     return label
@@ -200,6 +244,11 @@ def plot_sky_distribution(kind, nsim, grid, nside, MAS="PCS", plot_groups=True,
     nsnap = max(paths.get_snapshots(nsim))
     box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
 
+    if kind == "overdensity":
+        field = load_field("density", nsim, grid, MAS=MAS, in_rsp=False)
+        density_gen = csiborgtools.field.DensityField(box, MAS)
+        field = density_gen.overdensity_field(field) + 2
+    else:
         field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=False)
 
     angpos = csiborgtools.field.nside2radec(nside)
@@ -209,7 +258,7 @@ def plot_sky_distribution(kind, nsim, grid, nside, MAS="PCS", plot_groups=True,
 
     with plt.style.context(utils.mplstyle):
         label = get_sky_label(kind, volume_weight)
-        if kind == "density":
+        if kind in ["density", "overdensity"]:
             out = numpy.log10(out)
         healpy.mollview(out, fig=0, title="", unit=label)
 
@@ -248,7 +297,7 @@ if __name__ == "__main__":
     parser.add_argument('-c', '--clean', action='store_true')
     args = parser.parse_args()
 
-    cached_funcs = []
+    cached_funcs = ["load_field"]
     if args.clean:
         for func in cached_funcs:
             print(f"Cleaning cache for function {func}.")
@@ -258,6 +307,8 @@ if __name__ == "__main__":
     # plot_mass_vs_normcells(7444 + 24 * 4, pdf=False)
     # plot_mass_vs_ncells(7444, pdf=True)
     # plot_hmf(pdf=True)
-    plot_sky_distribution("potential", 7444, 256, nside=64, plot_groups=False,
+    # plot_sky_distribution("radvel", 7444, 256, nside=64,
-                          dmin=50, dmax=100, plot_halos=5e13,
+    #                       plot_groups=False, dmin=50, dmax=100,
-                          volume_weight=True)
+    #                       plot_halos=5e13, volume_weight=False)
+
+    plot_projected_field("potential", 7444, 256, in_rsp=True)