Run density field estimator (#56)

* 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

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):
-        self.pos = pos
-        self.mass = mass
+    def __init__(self, box, MAS):
         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
-            Gaussian kernal scale to smoothen the density field, in box units.
-        verbose : bool
+        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
             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)
-        if smooth_scale is not None:
-            rho = self.smoothen(rho, smooth_scale)
+
+        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
         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.
-    Note that this directly modifies the original array and throws away
-    particles outside the central region.
+    Mask of particles in a region of width `2 hw, centered at the origin.
 
     Parameters
     ----------
+    pos : 2-dimensional array of shape `(nparticles, 3)`
+        Particle positions, in box units.
     hw : float
-        Central region halfwidth.
-    particles : structured array
-        Particle array with keys `x`, `y`, `z`.
+        Central region half-width.
 
     Returns
     -------
-    particles : structured array
-        The modified particle array.
+    mask : 1-dimensional boolean array of shape `(nparticles, )`
     """
     assert 0 < hw < 0.5
-    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
+    return numpy.all((0.5 - hw < pos) & (pos < 0.5 + hw), axis=1)
 
 
 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