From 63b6cdbe7265b864dd8c6dffc8b4d7e03bef8c7a Mon Sep 17 00:00:00 2001 From: Richard Stiskalek Date: Mon, 5 Jun 2023 17:24:20 +0200 Subject: [PATCH] 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 | 3 +- csiborgtools/field/density.py | 37 +++++++- csiborgtools/field/interp.py | 100 ++++++++++++++++++- csiborgtools/read/paths.py | 5 +- scripts/field_derived.py | 78 --------------- scripts/field_prop.py | 169 +++++++++++++++++++++++++++++++++ scripts_plots/plot_data.py | 75 ++++++++++++--- 7 files changed, 371 insertions(+), 96 deletions(-) delete mode 100644 scripts/field_derived.py create mode 100644 scripts/field_prop.py diff --git a/csiborgtools/field/__init__.py b/csiborgtools/field/__init__.py index 65cb2c0..55dd462 100644 --- a/csiborgtools/field/__init__.py +++ b/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 diff --git a/csiborgtools/field/density.py b/csiborgtools/field/density.py index a2e0c16..af4df05 100644 --- a/csiborgtools/field/density.py +++ b/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): diff --git a/csiborgtools/field/interp.py b/csiborgtools/field/interp.py index e6f4929..5b86aa0 100644 --- a/csiborgtools/field/interp.py +++ b/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 diff --git a/csiborgtools/read/paths.py b/csiborgtools/read/paths.py index c77529c..d692b1c 100644 --- a/csiborgtools/read/paths.py +++ b/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) diff --git a/scripts/field_derived.py b/scripts/field_derived.py deleted file mode 100644 index 6185172..0000000 --- a/scripts/field_derived.py +++ /dev/null @@ -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) diff --git a/scripts/field_prop.py b/scripts/field_prop.py new file mode 100644 index 0000000..fe2f6db --- /dev/null +++ b/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) diff --git a/scripts_plots/plot_data.py b/scripts_plots/plot_data.py index 7cd27f3..fc70b28 100644 --- a/scripts_plots/plot_data.py +++ b/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,7 +244,12 @@ 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) - field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=False) + 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) dist = numpy.linspace(dmin, dmax, 500) @@ -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, - dmin=50, dmax=100, plot_halos=5e13, - volume_weight=True) + # plot_sky_distribution("radvel", 7444, 256, nside=64, + # plot_groups=False, dmin=50, dmax=100, + # plot_halos=5e13, volume_weight=False) + + plot_projected_field("potential", 7444, 256, in_rsp=True)