Forward model fields to RSP (#66)

* Add radial velocity field

* Add overdensity plot

* Flip velocities too

* Add field calculations

* Add RSP mapping

* Add potential in RSP

* Add projected field plotting

csiborgtools/field/__init__.py

@@ -18,7 +18,8 @@ try:
     import MAS_library as MASL  # noqa
 
     from .density import DensityField, PotentialField, VelocityField  # noqa
-    from .interp import evaluate_cartesian, evaluate_sky, make_sky  # noqa
+    from .interp import (evaluate_cartesian, evaluate_sky, field2rsp,  # noqa
+                         make_sky)
     from .utils import nside2radec, smoothen_field  # noqa
 except ImportError:
     warn("MAS_library not found, `DensityField` will not be available", UserWarning)  # noqa

csiborgtools/field/density.py

@@ -20,8 +20,10 @@ TODO:
 """
 from abc import ABC
 
-import MAS_library as MASL
 import numpy
+
+import MAS_library as MASL
+from numba import jit
 from tqdm import trange
 
 from ..read.utils import real2redshift
@@ -220,6 +222,36 @@ class VelocityField(BaseField):
         self.box = box
         self.MAS = MAS
 
+    @staticmethod
+    @jit(nopython=True)
+    def radial_velocity(rho_vel):
+        """
+        Calculate the radial velocity field around the observer in the centre
+        of the box.
+
+        Parameters
+        ----------
+        rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
+            Velocity field along each axis.
+
+        Returns
+        -------
+        radvel : 3-dimensional array of shape `(grid, grid, grid)`.
+            Radial velocity field.
+        """
+        grid = rho_vel.shape[1]
+        radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
+        for i in range(grid):
+            px = i - 0.5 * (grid - 1)
+            for j in range(grid):
+                py = j - 0.5 * (grid - 1)
+                for k in range(grid):
+                    pz = k - 0.5 * (grid - 1)
+                    vx, vy, vz = rho_vel[:, i, j, k]
+                    radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
+                                       / numpy.sqrt(px**2 + py**2 + pz**2))
+        return radvel
+
     def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30,
                  verbose=True):
         """
@@ -245,7 +277,7 @@ class VelocityField(BaseField):
 
         Returns
         -------
-        rho_vel : 3-dimensional array of shape `(3, grid, grid, grid)`.
+        rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
             Velocity field along each axis.
 
         References
@@ -272,6 +304,7 @@ class VelocityField(BaseField):
             mass = force_single_precision(mass, "particle_mass")
             if flip_xz:
                 pos[:, [0, 2]] = pos[:, [2, 0]]
+                vel[:, [0, 2]] = vel[:, [2, 0]]
             vel *= mass.reshape(-1, 1) / mpart
 
             for i in range(3):

csiborgtools/field/interp.py

@@ -17,9 +17,11 @@ Tools for interpolating 3D fields at arbitrary positions.
 """
 import MAS_library as MASL
 import numpy
+from numba import jit
+from scipy.ndimage import gaussian_filter
 from tqdm import trange
 
-from ..read.utils import radec_to_cartesian
+from ..read.utils import radec_to_cartesian, real2redshift
 from .utils import force_single_precision
 
 
@@ -137,3 +139,99 @@ def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True):
                 evaluate_sky(field, pos=dir_loop, box=box, isdeg=True))
     out *= dx
     return out
+
+
+@jit(nopython=True)
+def divide_nonzero(field0, field1):
+    """
+    Divide two fields where the second one is not zero. If the second field
+    is zero, the first one is left unchanged. Operates in-place.
+
+    Parameters
+    ----------
+    field0 : 3-dimensional array of shape `(grid, grid, grid)`
+        Field to be divided.
+    field1 : 3-dimensional array of shape `(grid, grid, grid)`
+        Field to divide by.
+
+    Returns
+    -------
+    field0 : 3-dimensional array of shape `(grid, grid, grid)`
+        Field divided by the second one.
+    """
+    assert field0.shape == field1.shape
+
+    imax, jmax, kmax = field0.shape
+    for i in range(imax):
+        for j in range(jmax):
+            for k in range(kmax):
+                if field1[i, j, k] != 0:
+                    field0[i, j, k] /= field1[i, j, k]
+
+
+def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
+              verbose=True):
+    """
+    Forward model real space scalar field to redshift space. Attaches field
+    values to particles, those are then moved to redshift space and from their
+    positions reconstructs back the field on a regular grid by NGP
+    interpolation. This by definition produces a discontinuous field.
+
+    Parameters
+    ----------
+    field : 3-dimensional array of shape `(grid, grid, grid)`
+        Real space field to be evolved to redshift space.
+    parts_pos : 2-dimensional array of shape `(n_parts, 3)`
+        Particle positions in real space.
+    parts_vel : 2-dimensional array of shape `(n_parts, 3)`
+        Particle velocities in real space.
+    box : :py:class:`csiborgtools.read.CSiBORGBox`
+        The simulation box information and transformations.
+    nbatch : int, optional
+        Number of batches to use when moving particles to redshift space.
+        Particles are assumed to be lazily loaded to memory.
+    flip_partsxz : bool, optional
+        Whether to flip the x and z coordinates of the particles. This is
+        because of a BORG bug.
+    init_value : float, optional
+        Initial value of the RSP field on the grid.
+    verbose : bool, optional
+        Verbosity flag.
+
+    Returns
+    -------
+    rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
+    """
+    rsp_field = numpy.full(field.shape, init_value, dtype=numpy.float32)
+    cellcounts = numpy.zeros(rsp_field.shape, dtype=numpy.float32)
+    # We iterate over the fields and in the inner loop over the particles. This
+    # is slower than iterating over the particles and in the inner loop over
+    # the fields, but it is more memory efficient. Typically we will only have
+    # one field.
+    nparts = parts.shape[0]
+    batch_size = nparts // nbatch
+    start = 0
+    for k in trange(nbatch + 1) if verbose else range(nbatch + 1):
+        end = min(start + batch_size, nparts)
+        pos = parts[start:end]
+        pos, vel = pos[:, :3], pos[:, 3:6]
+        if flip_partsxz:
+            pos[:, [0, 2]] = pos[:, [2, 0]]
+            vel[:, [0, 2]] = vel[:, [2, 0]]
+        # Evaluate the field at the particle positions in real space.
+        values = evaluate_cartesian(field, pos=pos)
+        # Move particles to redshift space.
+        rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
+                                in_box_units=True, periodic_wrap=True,
+                                make_copy=True)
+        # Assign particles' values to the grid.
+        MASL.MA(rsp_pos, rsp_field, 1., "NGP", W=values)
+        # Count the number of particles in each grid cell.
+        MASL.MA(rsp_pos, cellcounts, 1., "NGP")
+        if end == nparts:
+            break
+        start = end
+
+    # Finally divide by the number of particles in each cell and smooth.
+    divide_nonzero(rsp_field, cellcounts)
+    return rsp_field

csiborgtools/read/paths.py

@@ -363,7 +363,8 @@ class Paths:
         Parameters
         ----------
         kind : str
-            Field type. Must be one of: `density`, `velocity`, `potential`.
+            Field type. Must be one of: `density`, `velocity`, `potential`,
+            `radvel`.
         MAS : str
            Mass-assignment scheme.
         grid : int
@@ -378,7 +379,7 @@ class Paths:
         path : str
         """
         fdir = join(self.postdir, "environment")
-        assert kind in ["density", "velocity", "potential"]
+        assert kind in ["density", "velocity", "potential", "radvel"]
         if not isdir(fdir):
             makedirs(fdir)
             warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)