diff --git a/csiborgtools/field/density.py b/csiborgtools/field/density.py index f0d782e..4f7fef7 100644 --- a/csiborgtools/field/density.py +++ b/csiborgtools/field/density.py @@ -16,14 +16,16 @@ Density field and cross-correlation calculations. """ from abc import ABC +from warnings import warn import MAS_library as MASL import numpy +import Pk_library as PKL import smoothing_library as SL from tqdm import trange from .utils import force_single_precision -from ..read.utils import radec_to_cartesian, real2redshift +from ..read.utils import radec_to_cartesian class BaseField(ABC): @@ -187,6 +189,11 @@ class DensityField(BaseField): Parameters ---------- + pos : 2-dimensional array of shape `(N, 3)` + Particle position array. Columns must be ordered as `['x', 'y', 'z']`. + The positions are assumed to be in box units, i.e. :math:`\in [0, 1 ]`. + mass : 1-dimensional array of shape `(N,)` + Particle mass array. Assumed to be in box units. box : :py:class:`csiborgtools.read.BoxUnits` The simulation box information and transformations. MAS : str @@ -198,11 +205,52 @@ class DensityField(BaseField): ---------- [1] https://pylians3.readthedocs.io/ """ + _pos = None + _mass = None - def __init__(self, box, MAS): + def __init__(self, pos, mass, box, MAS): + self.pos = pos + self.mass = mass self.box = box self.MAS = MAS + @property + def pos(self): + """ + Particle position array. + + Returns + ------- + particles : 2-dimensional array + """ + return self._particles + + @pos.setter + def pos(self, pos): + assert pos.ndim == 2 + warn("Flipping the `x` and `z` coordinates of the particle positions.", + UserWarning, stacklevel=1) + pos[:, [0, 2]] = pos[:, [2, 0]] + pos = force_single_precision(pos, "particle_position") + self._pos = pos + + @property + def mass(self): + """ + Particle mass array. + + Returns + ------- + mass : 1-dimensional array + """ + return self._mass + + @mass.setter + def mass(self, mass): + assert mass.ndim == 1 + mass = force_single_precision(mass, "particle_mass") + self._mass = mass + def smoothen(self, field, smooth_scale, threads=1): """ Smooth a field with a Gaussian filter. @@ -246,28 +294,17 @@ class DensityField(BaseField): delta -= 1 return delta - def __call__(self, parts, grid, in_rsp, flip_xz=True, nbatch=30, - verbose=True): + def __call__(self, grid, smooth_scale=None, verbose=True): """ Calculate the density field using a Pylians routine [1, 2]. - Iteratively loads the particles into memory, flips their `x` and `z` - coordinates. Particles are assumed to be in box units, with positions - in [0, 1] and observer in the centre of the box if RSP is applied. Parameters ---------- - parts : 2-dimensional array of shape `(n_parts, 7)` - Particle positions, velocities and masses. - Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`. grid : int Grid size. - in_rsp : bool - Whether to calculate the density field in redshift space. - flip_xz : bool, optional - Whether to flip the `x` and `z` coordinates. - nbatch : int, optional - Number of batches to split the particle loading into. - verbose : bool, optional + smooth_scale : float, optional + Gaussian kernal scale to smoothen the density field, in box units. + verbose : bool Verbosity flag. Returns @@ -281,32 +318,12 @@ class DensityField(BaseField): [2] https://github.com/franciscovillaescusa/Pylians3/blob/master /library/MAS_library/MAS_library.pyx """ + # Pre-allocate and do calculations rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32) - - nparts = parts.shape[0] - batch_size = nparts // nbatch - start = 0 - for __ in trange(nbatch + 1) if verbose else range(nbatch + 1): - end = min(start + batch_size, nparts) - pos = parts[start:end] - pos, vel, mass = pos[:, :3], pos[:, 3:6], pos[:, 6] - - pos = force_single_precision(pos, "particle_position") - vel = force_single_precision(vel, "particle_velocity") - mass = force_single_precision(mass, "particle_mass") - if flip_xz: - pos[:, [0, 2]] = pos[:, [2, 0]] - vel[:, [0, 2]] = vel[:, [2, 0]] - - if in_rsp: - pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], self.box, - in_box_units=True, periodic_wrap=True, - make_copy=False) - - MASL.MA(pos, rho, self.boxsize, self.MAS, W=mass, verbose=False) - if end == nparts: - break - start = end + MASL.MA(self.pos, rho, self.boxsize, self.MAS, W=self.mass, + verbose=verbose) + if smooth_scale is not None: + rho = self.smoothen(rho, smooth_scale) return rho diff --git a/csiborgtools/read/__init__.py b/csiborgtools/read/__init__.py index a18f6dd..a0c6d67 100644 --- a/csiborgtools/read/__init__.py +++ b/csiborgtools/read/__init__.py @@ -21,7 +21,7 @@ from .overlap_summary import (NPairsOverlap, PairOverlap, # noqa binned_resample_mean) from .paths import CSiBORGPaths # noqa from .pk_summary import PKReader # noqa -from .readsim import (MmainReader, ParticleReader, halfwidth_mask, # noqa +from .readsim import (MmainReader, ParticleReader, halfwidth_select, # noqa load_clump_particles, load_parent_particles, read_initcm) from .tpcf_summary import TPCFReader # noqa from .utils import (cartesian_to_radec, cols_to_structured, # noqa diff --git a/csiborgtools/read/paths.py b/csiborgtools/read/paths.py index dcb5656..6c5bf64 100644 --- a/csiborgtools/read/paths.py +++ b/csiborgtools/read/paths.py @@ -329,18 +329,16 @@ class CSiBORGPaths: fname = f"parts_{str(nsim).zfill(5)}.h5" return join(fdir, fname) - def density_field_path(self, MAS, nsim, in_rsp): + def density_field_path(self, mas, nsim): """ Path to the files containing the calculated density fields. Parameters ---------- - MAS : str + mas : str Mass-assignment scheme. Currently only SPH is supported. nsim : int IC realisation index. - in_rsp : bool - Whether the density field is calculated in redshift space. Returns ------- @@ -350,9 +348,7 @@ class CSiBORGPaths: if not isdir(fdir): makedirs(fdir) warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1) - fname = f"density_{MAS}_{str(nsim).zfill(5)}.npy" - if in_rsp: - fname = fname.replace("density", "density_rsp") + fname = f"density_{mas}_{str(nsim).zfill(5)}.npy" return join(fdir, fname) def knnauto_path(self, run, nsim=None): diff --git a/csiborgtools/read/readsim.py b/csiborgtools/read/readsim.py index b6942a5..32cb319 100644 --- a/csiborgtools/read/readsim.py +++ b/csiborgtools/read/readsim.py @@ -534,23 +534,34 @@ def read_initcm(nsim, srcdir, fname="clump_{}_cm.npy"): return None -def halfwidth_mask(pos, hw): +def halfwidth_select(hw, particles): """ - Mask of particles in a region of width `2 hw, centered at the origin. + Select particles that in a cube of size `2 hw`, centered at the origin. + Note that this directly modifies the original array and throws away + particles outside the central region. Parameters ---------- - pos : 2-dimensional array of shape `(nparticles, 3)` - Particle positions, in box units. hw : float - Central region half-width. + Central region halfwidth. + particles : structured array + Particle array with keys `x`, `y`, `z`. Returns ------- - mask : 1-dimensional boolean array of shape `(nparticles, )` + particles : structured array + The modified particle array. """ assert 0 < hw < 0.5 - return numpy.all((0.5 - hw < pos) & (pos < 0.5 + hw), axis=1) + mask = ((0.5 - hw < particles['x']) & (particles['x'] < 0.5 + hw) + & (0.5 - hw < particles['y']) & (particles['y'] < 0.5 + hw) + & (0.5 - hw < particles['z']) & (particles['z'] < 0.5 + hw)) + # Subselect the particles + particles = particles[mask] + # Rescale to range [0, 1] + for p in ('x', 'y', 'z'): + particles[p] = (particles[p] - 0.5 + hw) / (2 * hw) + return particles def load_clump_particles(clid, particles, clump_map, clid2map): diff --git a/csiborgtools/read/utils.py b/csiborgtools/read/utils.py index 01c8517..ab2537b 100644 --- a/csiborgtools/read/utils.py +++ b/csiborgtools/read/utils.py @@ -80,8 +80,7 @@ def radec_to_cartesian(X, isdeg=True): return dist * numpy.vstack([x, y, z]).T -def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True, - make_copy=True): +def real2redshift(pos, vel, origin, box, in_box_units, make_copy=True): r""" Convert real-space position to redshift space position. @@ -100,9 +99,6 @@ def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True, to be in :math:`\mathrm{Mpc}`, velocity in :math:`\mathrm{km} \mathrm{s}^{-1}` and math:`h=0.705`, or otherwise matching the box. - periodic_wrap : bool, optional - Whether to wrap around the box, particles may be outside the default - bounds once RSD is applied. make_copy : bool, optional Whether to make a copy of `pos` before modifying it. @@ -126,12 +122,6 @@ def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True, for i in range(3): pos[:, i] += origin[i] - - if periodic_wrap: - boxsize = 1. if in_box_units else box.box2mpc(1.) - # Wrap around the box: x > 1 -> x - 1, x < 0 -> x + 1 - pos[pos > boxsize] -= boxsize - pos[pos < 0] += boxsize return pos diff --git a/scripts/field_density.py b/scripts/field_density.py deleted file mode 100644 index 22f9ff2..0000000 --- a/scripts/field_density.py +++ /dev/null @@ -1,70 +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 evaluate field properties at the galaxy positions. - -NOTE THAT ONLY MAX SNAP -""" -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("--grid", type=int, help="Grid resolution.") -parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)), - help="Calculate the density field in redshift space?") -parser.add_argument("--MAS", type=str, choices=["NGP", "CIC", "TSC", "PCS"]) -args = parser.parse_args() -paths = csiborgtools.read.CSiBORGPaths(**csiborgtools.paths_glamdring) - -if args.ics is None or args.ics[0] == -1: - ics = paths.get_ics(tonew=False) -else: - ics = args.ics - - -for i in csiborgtools.fits.split_jobs(len(ics), nproc)[rank]: - nsim = ics[i] - print(f"{datetime.now()}: rank {rank} working on simulation {nsim}.", - flush=True) - - nsnap = max(paths.get_snapshots(nsim)) - box = csiborgtools.read.BoxUnits(nsnap, nsim, paths) - parts = csiborgtools.read.read_h5(paths.particles_path(nsim))["particles"] - density_generator = csiborgtools.field.DensityField(box, args.MAS) - - rho = density_generator(parts, args.grid, in_rsp=args.in_rsp, - verbose=verbose) - - fout = paths.density_field_path(args.MAS, nsim, args.in_rsp) - print(f"{datetime.now()}: rank {rank} saving output to `{fout}`.") - numpy.save(fout, rho) diff --git a/scripts/field_prop.py b/scripts/field_prop.py new file mode 100644 index 0000000..4bd0db6 --- /dev/null +++ b/scripts/field_prop.py @@ -0,0 +1,127 @@ +# 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 evaluate field properties at the galaxy positions. +""" +from argparse import ArgumentParser +from datetime import datetime +from os import remove +from os.path import join + +import numpy +from mpi4py import MPI + +try: + import csiborgtools +except ModuleNotFoundError: + import sys + sys.path.append("../") + import csiborgtools + +import utils + +dumpdir = "/mnt/extraspace/rstiskalek/csiborg/" +parser = ArgumentParser() +parser.add_argument("--survey", type=str, choices=["SDSS"]) +parser.add_argument("--grid", type=int) +parser.add_argument("--MAS", type=str, choices=["NGP", "CIC", "TSC", "PCS"]) +parser.add_argument("--halfwidth", type=float) +parser.add_argument("--smooth_scale", type=float, default=None) +args = parser.parse_args() +# Smooth scale of 0 means no smoothing. Note that this is in Mpc/h +args.smooth_scale = None if args.smooth_scale == 0 else args.smooth_scale + +# Get MPI things +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +nproc = comm.Get_size() + +# Galaxy positions +survey = utils.surveys[args.survey]()() +pos = numpy.vstack([survey[p] for p in ("DIST", "RA", "DEC")]).T +pos = pos.astype(numpy.float32) + +# File paths +fname = "out_{}_{}_{}_{}_{}".format( + survey.name, args.grid, args.MAS, args.halfwidth, args.smooth_scale) +ftemp = join(dumpdir, "temp_fields", fname + "_{}.npy") +fperm = join(dumpdir, "fields", fname + ".npy") + +# Edit depending on what is calculated +dtype = {"names": ["delta", "phi"], "formats": [numpy.float32] * 2} + +# CSiBORG simulation paths +paths = csiborgtools.read.CSiBORGPaths(**csiborgtools.paths_glamdring) +ics = paths.get_ics(tonew=False) +nsims = len(ics) + +for n in csiborgtools.utils.split_jobs(nsims, nproc)[rank]: + print("Rank {}@{}: working on {}th IC.".format(rank, datetime.now(), n), + flush=True) + nsim = ics[n] + nsnap = max(paths.get_snapshots(nsim)) + reader = csiborgtools.read.ParticleReader(paths) + box = csiborgtools.read.BoxUnits(nsnap, nsim, paths) + + # Read particles and select a subset of them + particles = reader.read_particle(nsnap, nsim, ["x", "y", "z", "M"], + verbose=False) + if args.halfwidth < 0.5: + particles = csiborgtools.read.halfwidth_select( + args.halfwidth, particles) + length = box.box2mpc(2 * args.halfwidth) * box.h # Mpc/h + else: + length = box.box2mpc(1) * box.h # Mpc/h + + # Initialise the field object and output array + field = csiborgtools.field.DensityField(particles, length, box, args.MAS) + out = numpy.full(pos.shape[0], numpy.nan, dtype=dtype) + + # Calculate the overdensity field and interpolate at galaxy positions + feval = field.overdensity_field(args.grid, args.smooth_scale, + verbose=False) + out["delta"] = field.evaluate_sky(feval, pos=pos, isdeg=True)[0] + + # Potential + feval = field.potential_field(args.grid, args.smooth_scale, verbose=False) + out["phi"] = field.evaluate_sky(feval, pos=pos, isdeg=True)[0] + + # Calculate the remaining fields + # ... + # ... + + # Dump the results + with open(ftemp.format(nsim), "wb") as f: + numpy.save(f, out) + +# Wait for all ranks to finish +comm.Barrier() +if rank == 0: + print("Collecting files...", flush=True) + + out = numpy.full((nsims, pos.shape[0]), numpy.nan, dtype=dtype) + + for n in range(nsims): + nsim = ics[n] + with open(ftemp.format(nsim), "rb") as f: + fin = numpy.load(f, allow_pickle=True) + for name in dtype["names"]: + out[name][n, ...] = fin[name] + # Remove the temporary file + remove(ftemp.format(nsim)) + + print("Saving results to `{}`.".format(fperm), flush=True) + with open(fperm, "wb") as f: + numpy.save(f, out)