Revert "Run density field estimator (#56)" (#57)

This reverts commit 98d0578fa7.

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,
+    def __call__(self, grid, smooth_scale=None, verbose=True):
-                 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
+        smooth_scale : float, optional
-            Whether to calculate the density field in redshift space.
+            Gaussian kernal scale to smoothen the density field, in box units.
-        flip_xz : bool, optional
+        verbose : bool
-            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
@@ -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)
-
+        MASL.MA(self.pos, rho, self.boxsize, self.MAS, W=self.mass,
-        nparts = parts.shape[0]
+                verbose=verbose)
-        batch_size = nparts // nbatch
+        if smooth_scale is not None:
-        start = 0
+            rho = self.smoothen(rho, smooth_scale)
-        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_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

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"
+        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,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):

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,
+def real2redshift(pos, vel, origin, box, in_box_units, make_copy=True):
-                  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
 
 

scripts/field_density.py

@@ -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)

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)