Density field calculation (#58)

* Add iterative density field generation

* Edit particle halfwidth selection

* Update import

* Remove old file

* Add position wrapping

* Add RSD support

* Add density field calculation

* Edit paths to the density field

* Flip argument order

* Edit documentation

* Edit docs

* update nb

csiborgtools/field/density.py

@@ -16,16 +16,14 @@
 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
+from ..read.utils import radec_to_cartesian, real2redshift
 
 
 class BaseField(ABC):
@@ -189,11 +187,6 @@ 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
@@ -205,52 +198,11 @@ class DensityField(BaseField):
     ----------
     [1] https://pylians3.readthedocs.io/
     """
-    _pos = None
-    _mass = None
 
-    def __init__(self, pos, mass, box, MAS):
+    def __init__(self, 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.
@@ -294,17 +246,28 @@ class DensityField(BaseField):
         delta -= 1
         return delta
 
-    def __call__(self, grid, smooth_scale=None, verbose=True):
+    def __call__(self, parts, grid, in_rsp, flip_xz=True, nbatch=30,
+                 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.
-        smooth_scale : float, optional
+        in_rsp : bool
-            Gaussian kernal scale to smoothen the density field, in box units.
+            Whether to calculate the density field in redshift space.
-        verbose : bool
+        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
             Verbosity flag.
 
         Returns
@@ -318,12 +281,32 @@ 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)
-        MASL.MA(self.pos, rho, self.boxsize, self.MAS, W=self.mass,
+
-                verbose=verbose)
+        nparts = parts.shape[0]
-        if smooth_scale is not None:
+        batch_size = nparts // nbatch
-            rho = self.smoothen(rho, smooth_scale)
+        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
         return rho
 
 

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_select,  # noqa
+from .readsim import (MmainReader, ParticleReader, halfwidth_mask,  # 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

csiborgtools/read/paths.py

@@ -329,16 +329,18 @@ class CSiBORGPaths:
         fname = f"parts_{str(nsim).zfill(5)}.h5"
         return join(fdir, fname)
 
-    def density_field_path(self, mas, nsim):
+    def density_field_path(self, MAS, nsim, in_rsp):
         """
         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
         -------
@@ -348,7 +350,9 @@ 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"
+        fname = f"density_{MAS}_{str(nsim).zfill(5)}.npy"
+        if in_rsp:
+            fname = fname.replace("density", "density_rsp")
         return join(fdir, fname)
 
     def knnauto_path(self, run, nsim=None):

csiborgtools/read/readsim.py

@@ -534,34 +534,23 @@ def read_initcm(nsim, srcdir, fname="clump_{}_cm.npy"):
         return None
 
 
-def halfwidth_select(hw, particles):
+def halfwidth_mask(pos, hw):
     """
-    Select particles that in a cube of size `2 hw`, centered at the origin.
+    Mask of particles in a region of width `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 halfwidth.
+        Central region half-width.
-    particles : structured array
-        Particle array with keys `x`, `y`, `z`.
 
     Returns
     -------
-    particles : structured array
+    mask : 1-dimensional boolean array of shape `(nparticles, )`
-        The modified particle array.
     """
     assert 0 < hw < 0.5
-    mask = ((0.5 - hw < particles['x']) & (particles['x'] < 0.5 + hw)
+    return numpy.all((0.5 - hw < pos) & (pos < 0.5 + hw), axis=1)
-            & (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):

csiborgtools/read/utils.py

@@ -80,7 +80,8 @@ def radec_to_cartesian(X, isdeg=True):
     return dist * numpy.vstack([x, y, z]).T
 
 
-def real2redshift(pos, vel, origin, box, in_box_units, make_copy=True):
+def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True,
+                  make_copy=True):
     r"""
     Convert real-space position to redshift space position.
 
@@ -99,6 +100,9 @@ def real2redshift(pos, vel, origin, box, in_box_units, make_copy=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.
 
@@ -122,6 +126,12 @@ def real2redshift(pos, vel, origin, box, in_box_units, make_copy=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
 
 

scripts/field_density.py

@@ -0,0 +1,69 @@
+# 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 the density fields on CSiBORG simulations in the 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("--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)

scripts/field_prop.py

@@ -1,127 +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.
-"""
-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)