From de7def61d57795c23327ba9f27395df325648144 Mon Sep 17 00:00:00 2001 From: Richard Stiskalek Date: Mon, 26 Jun 2023 20:41:07 +0100 Subject: [PATCH] New RSP density distinction (#72) * Edit docs * Fix kernel units * Add NGP interpolation * Add Vobs but not implemented * Add BORG density option * Add BORG mcmc path * Organise imports * Add new density field distinction --- csiborgtools/__init__.py | 2 + csiborgtools/field/__init__.py | 5 ++- csiborgtools/field/density.py | 19 +++++---- csiborgtools/field/interp.py | 48 +++++++++++++++++---- csiborgtools/read/paths.py | 44 +++++++++++++++++++- scripts/field_prop.py | 76 ++++++++++++---------------------- scripts_plots/plot_data.py | 21 +++++++--- 7 files changed, 141 insertions(+), 74 deletions(-) diff --git a/csiborgtools/__init__.py b/csiborgtools/__init__.py index e456cfe..65a40f3 100644 --- a/csiborgtools/__init__.py +++ b/csiborgtools/__init__.py @@ -17,6 +17,8 @@ from csiborgtools import clustering, field, fits, match, read # noqa # Arguments to csiborgtools.read.Paths. paths_glamdring = {"srcdir": "/mnt/extraspace/hdesmond/", "postdir": "/mnt/extraspace/rstiskalek/CSiBORG/", + "BORG_dir": "/mnt/extraspace/rstiskalek/BORG/", + "BORG_final_density": "/users/hdesmond/BORG_final/", "quijote_dir": "/mnt/extraspace/rstiskalek/Quijote", } diff --git a/csiborgtools/field/__init__.py b/csiborgtools/field/__init__.py index dcd3773..772b911 100644 --- a/csiborgtools/field/__init__.py +++ b/csiborgtools/field/__init__.py @@ -17,9 +17,10 @@ from warnings import warn try: import MAS_library as MASL # noqa - from .density import DensityField, PotentialField, VelocityField, TidalTensorField # noqa + from .density import (DensityField, PotentialField, # noqa + TidalTensorField, VelocityField) from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa - make_sky, fill_outside) # noqa + fill_outside, make_sky, observer_vobs) 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 1039766..edc482b 100644 --- a/csiborgtools/field/density.py +++ b/csiborgtools/field/density.py @@ -17,13 +17,13 @@ Density field and cross-correlation calculations. """ from abc import ABC -import numpy - import MAS_library as MASL +import numpy from numba import jit from tqdm import trange from ..read.utils import real2redshift +from .interp import divide_nonzero from .utils import force_single_precision @@ -249,7 +249,7 @@ class VelocityField(BaseField): / numpy.sqrt(px**2 + py**2 + pz**2)) return radvel - def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30, + def __call__(self, parts, grid, flip_xz=True, nbatch=30, verbose=True): """ Calculate the velocity field using a Pylians routine [1, 2]. @@ -263,8 +263,6 @@ class VelocityField(BaseField): Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`. grid : int Grid size. - mpart : float - Particle mass. flip_xz : bool, optional Whether to flip the `x` and `z` coordinates. nbatch : int, optional @@ -287,6 +285,7 @@ class VelocityField(BaseField): rho_vely = numpy.zeros((grid, grid, grid), dtype=numpy.float32) rho_velz = numpy.zeros((grid, grid, grid), dtype=numpy.float32) rho_vel = [rho_velx, rho_vely, rho_velz] + cellcounts = numpy.zeros((grid, grid, grid), dtype=numpy.float32) nparts = parts.shape[0] batch_size = nparts // nbatch @@ -302,16 +301,22 @@ class VelocityField(BaseField): if flip_xz: pos[:, [0, 2]] = pos[:, [2, 0]] vel[:, [0, 2]] = vel[:, [2, 0]] - vel *= mass.reshape(-1, 1) / mpart + vel *= mass.reshape(-1, 1) for i in range(3): MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[:, i], verbose=False) + + MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass, + verbose=False) if end == nparts: break start = end - return numpy.stack(rho_vel) + for i in range(3): + divide_nonzero(rho_vel[i], cellcounts) + + return numpy.stack([rho_velx, rho_vely, rho_velz]) ############################################################################### diff --git a/csiborgtools/field/interp.py b/csiborgtools/field/interp.py index b9a3632..8d92f05 100644 --- a/csiborgtools/field/interp.py +++ b/csiborgtools/field/interp.py @@ -24,10 +24,9 @@ from ..read.utils import radec_to_cartesian, real2redshift from .utils import force_single_precision -def evaluate_cartesian(*fields, pos): +def evaluate_cartesian(*fields, pos, interp="CIC"): """ - Evaluate a scalar field at Cartesian coordinates using CIC - interpolation. + Evaluate a scalar field at Cartesian coordinates. Parameters ---------- @@ -36,19 +35,29 @@ def evaluate_cartesian(*fields, pos): pos : 2-dimensional array of shape `(n_samples, 3)` Positions to evaluate the density field. Assumed to be in box units. + interp : str, optional + Interpolation method. Can be either `CIC` or `NGP`. Returns ------- interp_fields : (list of) 1-dimensional array of shape `(n_samples,). """ + assert interp in ["CIC", "NGP"] boxsize = 1. pos = force_single_precision(pos, "pos") nsamples = pos.shape[0] interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32) for __ in range(len(fields))] - for i, field in enumerate(fields): - MASL.CIC_interp(field, boxsize, pos, interp_fields[i]) + + if interp == "CIC": + for i, field in enumerate(fields): + MASL.CIC_interp(field, boxsize, pos, interp_fields[i]) + else: + pos = numpy.floor(pos * fields[0].shape[0]).astype(numpy.int32) + for i, field in enumerate(fields): + for j in range(nsamples): + interp_fields[i][j] = field[pos[j, 0], pos[j, 1], pos[j, 2]] if len(fields) == 1: return interp_fields[0] @@ -168,8 +177,30 @@ def divide_nonzero(field0, field1): field0[i, j, k] /= field1[i, j, k] -def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0., - verbose=True): +def observer_vobs(velocity_field): + """ + Calculate the observer velocity from a velocity field. Assumes the observer + is in the centre of the box. + + Parameters + ---------- + velocity_field : 4-dimensional array of shape `(3, grid, grid, grid)` + Velocity field to calculate the observer velocity from. + + Returns + ------- + vobs : 1-dimensional array of shape `(3,)` + Observer velocity in units of `velocity_field`. + """ + pos = numpy.asanyarray([0.5, 0.5, 0.5]).reshape(1, 3) + vobs = numpy.full(3, numpy.nan, dtype=numpy.float32) + for i in range(3): + vobs[i] = evaluate_cartesian(velocity_field[i, ...], pos=pos)[0] + return vobs + + +def field2rsp(*fields, parts, vobs, box, nbatch=30, flip_partsxz=True, + init_value=0., verbose=True): """ Forward model real space scalar fields to redshift space. Attaches field values to particles, those are then moved to redshift space and from their @@ -182,6 +213,8 @@ def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0., 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`. + vobs : 1-dimensional array of shape `(3,)` + Observer velocity in units matching `parts`. box : :py:class:`csiborgtools.read.CSiBORGBox` The simulation box information and transformations. nbatch : int, optional @@ -199,6 +232,7 @@ def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0., ------- rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)` """ + raise NotImplementedError("Figure out what to do with Vobs.") nfields = len(fields) # Check that all fields have the same shape. if nfields > 1: diff --git a/csiborgtools/read/paths.py b/csiborgtools/read/paths.py index 51d9f6b..e7de6bb 100644 --- a/csiborgtools/read/paths.py +++ b/csiborgtools/read/paths.py @@ -31,16 +31,21 @@ class Paths: Path to the folder where the RAMSES outputs are stored. postdir: str, optional Path to the folder where post-processed files are stored. + borg_dir : str, optional + Path to the folder where BORG MCMC chains are stored. quiote_dir : str, optional Path to the folder where Quijote simulations are stored. """ _srcdir = None _postdir = None + _borg_dir = None _quijote_dir = None - def __init__(self, srcdir=None, postdir=None, quijote_dir=None): + def __init__(self, srcdir=None, postdir=None, borg_dir=None, + quijote_dir=None): self.srcdir = srcdir self.postdir = postdir + self.borg_dir = borg_dir self.quijote_dir = quijote_dir @staticmethod @@ -68,6 +73,26 @@ class Paths: self._check_directory(path) self._srcdir = path + @property + def borg_dir(self): + """ + Path to the folder where BORG MCMC chains are stored. + + Returns + ------- + path : str + """ + if self._borg_dir is None: + raise ValueError("`borg_dir` is not set!") + return self._borg_dir + + @borg_dir.setter + def borg_dir(self, path): + if path is None: + return + self._check_directory(path) + self._borg_dir = path + @property def quijote_dir(self): """ @@ -146,6 +171,21 @@ class Paths: return nsim return f"{str(nobs).zfill(2)}{str(nsim).zfill(3)}" + def borg_mcmc(self, nsim): + """ + Path to the BORG MCMC chain file. + + Parameters + ---------- + nsim : int + IC realisation index. + + Returns + ------- + path : str + """ + return join(self.borg_dir, "mcmc", f"mcmc_{nsim}.h5") + def mmain(self, nsnap, nsim): """ Path to the `mmain` CSiBORG files of summed substructure. @@ -373,7 +413,7 @@ class Paths: IC realisation index. in_rsp : bool Whether the calculation is performed in redshift space. - smooth_scale : float + smooth_scale : float, optional Smoothing scale in :math:`\mathrm{Mpc}/h` Returns diff --git a/scripts/field_prop.py b/scripts/field_prop.py index 41ded47..003161d 100644 --- a/scripts/field_prop.py +++ b/scripts/field_prop.py @@ -61,54 +61,30 @@ def density_field(nsim, parser_args, to_save=True): 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) + + if parser_args.kind == "density": + field = gen(parts, parser_args.grid, in_rsp=False, + verbose=parser_args.verbose) + if parser_args.in_rsp: + field = csiborgtools.field.field2rsp(*field, parts=parts, box=box, + verbose=parser_args.verbose) + else: + field = gen(parts, parser_args.grid, in_rsp=parser_args.in_rsp, + verbose=parser_args.verbose) + + if parser_args.smooth_scale > 0: + field = csiborgtools.field.smoothen_field( + field, parser_args.smooth_scale, box.boxsize * box.h, threads=1) if to_save: - fout = paths.field("density", parser_args.MAS, parser_args.grid, - nsim, parser_args.in_rsp) + fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid, + nsim, parser_args.in_rsp, parser_args.smooth_scale) print(f"{datetime.now()}: saving output to `{fout}`.") numpy.save(fout, field) return field -def density_field_smoothed(nsim, parser_args, to_save=True): - """ - Calculate the smoothed density field in the CSiBORG simulation. The - unsmoothed density field must already be precomputed. - - Parameters - ---------- - nsim : int - Simulation index. - parser_args : argparse.Namespace - Parsed arguments. - to_save : bool, optional - Whether to save the output to disk. - - Returns - ------- - smoothed_density : 3-dimensional array - """ - 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 - rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid, - nsim, in_rsp=False)) - rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale, - box.boxsize, threads=1) - if to_save: - fout = paths.field("density", parser_args.MAS, parser_args.grid, - nsim, parser_args.in_rsp, parser_args.smooth_scale) - print(f"{datetime.now()}: saving output to `{fout}`.") - numpy.save(fout, rho) - return rho - - ############################################################################### # Velocity field # ############################################################################### @@ -185,7 +161,7 @@ def potential_field(nsim, parser_args, to_save=True): nsim, in_rsp=False)) if parser_args.smooth_scale > 0: rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale, - box.boxsize, threads=1) + box.boxsize * box.h, threads=1) rho = density_gen.overdensity_field(rho) # Calculate the real space potentiel field gen = csiborgtools.field.PotentialField(box, parser_args.MAS) @@ -281,7 +257,7 @@ def environment_field(nsim, parser_args, to_save=True): nsim, in_rsp=False)) if parser_args.smooth_scale > 0: rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale, - box.boxsize, threads=1) + box.boxsize * box.h, threads=1) rho = density_gen.overdensity_field(rho) # Calculate the real space tidal tensor field, delete overdensity. if parser_args.verbose: @@ -339,28 +315,28 @@ if __name__ == "__main__": 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", - "environment"], + choices=["density", "rspdensity", "velocity", "radvel", + "potential", "environment"], 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("--smooth_scale", type=float, default=0) + parser.add_argument("--smooth_scale", type=float, default=0, + help="Smoothing scale in Mpc/h.") parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)), help="Verbosity flag for reading in particles.") + parser.add_argument("--simname", type=str, default="csiborg", + 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": - if parser_args.smooth_scale > 0: - density_field_smoothed(nsim, parser_args) - else: - density_field(nsim, parser_args) + if parser_args.kind == "density" or parser_args.kind == "rspdensity": + density_field(nsim, parser_args) elif parser_args.kind == "velocity": velocity_field(nsim, parser_args) elif parser_args.kind == "radvel": diff --git a/scripts_plots/plot_data.py b/scripts_plots/plot_data.py index 4bf2208..43ab86a 100644 --- a/scripts_plots/plot_data.py +++ b/scripts_plots/plot_data.py @@ -19,6 +19,7 @@ from argparse import ArgumentParser import matplotlib as mpl import matplotlib.pyplot as plt import numpy +from h5py import File import healpy import scienceplots # noqa @@ -244,7 +245,8 @@ def load_field(kind, nsim, grid, MAS, in_rsp=False, smooth_scale=None): def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS", - highres_only=True, slice_find=None, pdf=False): + vel_component=0, highres_only=True, slice_find=None, + pdf=False): r""" Plot the mean projected field, however can also plot a single slice. @@ -262,6 +264,8 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS", Smoothing scale in :math:`\mathrm{Mpc} / h`. MAS : str, optional Mass assignment scheme. + vel_component : int, optional + Which velocity field component to plot. highres_only : bool, optional Whether to only plot the high-resolution region. slice_find : float, optional @@ -283,12 +287,16 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS", smooth_scale=smooth_scale) density_gen = csiborgtools.field.DensityField(box, MAS) field = density_gen.overdensity_field(field) + 1 + elif kind == "borg_density": + field = File(paths.borg_mcmc(nsim), 'r') + field = field["scalars"]["BORG_final_density"][...] else: field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp, smooth_scale=smooth_scale) if kind == "velocity": - field = field[0, ...] + field = field[vel_component, ...] + field = box.box2vel(field) if highres_only: csiborgtools.field.fill_outside(field, numpy.nan, rmax=155.5, @@ -552,22 +560,23 @@ if __name__ == "__main__": if False: plot_hmf(pdf=False) - if True: + if False: kind = "overdensity" grid = 1024 plot_sky_distribution(kind, 7444, grid, nside=64, plot_groups=False, dmin=45, dmax=60, plot_halos=5e13, volume_weight=True) - if False: - kind = "density" + if True: + kind = "overdensity" grid = 256 smooth_scale = 0 # plot_projected_field("overdensity", 7444, grid, in_rsp=True, # highres_only=False) plot_projected_field(kind, 7444, grid, in_rsp=False, smooth_scale=smooth_scale, slice_find=0.5, - highres_only=False) + MAS="PCS", + highres_only=True) if False: paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)