Add redshift space position (#55)

* Add docs

* Delete blank

* Add real to redshift space function

* RSD move back observer

* Convert velocoties to km/s

* Add units to RSP calculation

* Add RSP to catalogues

* Add RSD to angular neighbour search

csiborgtools/read/__init__.py

@@ -25,4 +25,4 @@ 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
-                    radec_to_cartesian, read_h5)
+                    radec_to_cartesian, read_h5, real2redshift)

csiborgtools/read/box_units.py

@@ -27,6 +27,7 @@ CONV_NAME = {
     "length": ["x", "y", "z", "peak_x", "peak_y", "peak_z", "Rs", "rmin",
                "rmax", "r200c", "r500c", "r200m", "x0", "y0", "z0",
                "lagpatch"],
+    "velocity": ["vx", "vy", "vz"],
     "mass": ["mass_cl", "totpartmass", "m200c", "m500c", "mass_mmain", "M",
              "m200m"],
     "density": ["rho0"]}
@@ -93,8 +94,8 @@ class BoxUnits:
     @property
     def H0(self):
         r"""
-        The Hubble parameter at the time of the snapshot
-        in :math:`\mathrm{Mpc} / \mathrm{km} / \mathrm{s}`.
+        The Hubble parameter at the time of the snapshot in units of
+        :math:`\mathrm{km} \mathrm{s}^{-1} \mathrm{Mpc}^{-1}`.
 
         Returns
         -------
@@ -229,7 +230,7 @@ class BoxUnits:
 
     def box2vel(self, vel):
         r"""
-        Convert velocity from box units to :math:`\mathrm{m} / \mathrm{s}^2`.
+        Convert velocity from box units to :math:`\mathrm{km} \mathrm{s}^{-1}`.
 
         Parameters
         ----------
@@ -239,9 +240,9 @@ class BoxUnits:
         Returns
         -------
         vel : float
-            Velocity in :math:`\mathrm{m} / \mathrm{s}^2`
+            Velocity in :math:`\mathrm{km} \mathrm{s}^{-1}`.
         """
-        return vel * (1e-2 * self._unit_l / self._unit_t / self._aexp)
+        return vel * (1e-2 * self._unit_l / self._unit_t / self._aexp) * 1e-3
 
     def box2cosmoredshift(self, dist):
         r"""
@@ -396,6 +397,7 @@ class BoxUnits:
         physical units, such that
             - length -> :math:`Mpc`,
             - mass -> :math:`M_\odot`,
+            - velocity -> :math:`\mathrm{km} / \mathrm{s}`,
             - density -> :math:`M_\odot / \mathrm{Mpc}^3`.
 
         Any other conversions are currently not implemented. Note that the
@@ -418,6 +420,7 @@ class BoxUnits:
         names = [names] if isinstance(names, str) else names
         transforms = {"length": self.box2mpc,
                       "mass": self.box2solarmass,
+                      "velocity": self.box2vel,
                       "density": self.box2dens,
                       }
 

csiborgtools/read/halo_cat.py

@@ -22,7 +22,7 @@ from .box_units import BoxUnits
 from .paths import CSiBORGPaths
 from .readsim import ParticleReader
 from .utils import (add_columns, cartesian_to_radec, flip_cols,
-                    radec_to_cartesian)
+                    radec_to_cartesian, real2redshift)
 
 
 class BaseCatalogue(ABC):
@@ -152,6 +152,33 @@ class BaseCatalogue(ABC):
         """
         return numpy.vstack([self["v{}".format(p)] for p in ("x", "y", "z")]).T
 
+    def redshift_space_position(self, cartesian=True):
+        r"""
+        Redshift space position components. If Cartesian, then in
+        :math:`\mathrm{cMpc}`. If spherical, then radius is in
+        :math:`\mathrm{cMpc}`, RA in :math:`[0, 360)` degrees and DEC in
+        :math:`[-90, 90]` degrees. Note that the position is defined as the
+        minimum of the gravitationl potential.
+
+        Parameters
+        ----------
+        cartesian : bool, optional
+            Whether to return the Cartesian or spherical position components.
+            By default Cartesian.
+
+        Returns
+        -------
+        pos : 2-dimensional array of shape `(nobjects, 3)`
+        """
+        pos = self.position(cartesian=True)
+        vel = self.velocity()
+        origin = [0., 0., 0.]
+        rsp = real2redshift(pos, vel, origin, self.box, in_box_units=False,
+                            make_copy=False)
+        if not cartesian:
+            rsp = cartesian_to_radec(rsp)
+        return rsp
+
     def angmomentum(self):
         """
         Cartesian angular momentum components of halos in the box coordinate
@@ -207,7 +234,7 @@ class BaseCatalogue(ABC):
         knn = self.knn(in_initial)
         return knn.radius_neighbors(X, radius, sort_results=True)
 
-    def angular_neighbours(self, X, ang_radius, rad_tolerance=None):
+    def angular_neighbours(self, X, ang_radius, in_rsp, rad_tolerance=None):
         r"""
         Find nearest neighbours within `ang_radius` of query points `X`.
         Optionally applies radial tolerance, which is expected to be in
@@ -219,6 +246,8 @@ class BaseCatalogue(ABC):
             Query positions. If 2-dimensional, then RA and DEC in degrees.
             If 3-dimensional, then radial distance in :math:`\mathrm{cMpc}`,
             RA and DEC in degrees.
+        in_rsp : bool
+            Whether to use redshift space positions of haloes.
         ang_radius : float
             Angular radius in degrees.
         rad_tolerance : float, optional
@@ -234,7 +263,10 @@ class BaseCatalogue(ABC):
         assert X.ndim == 2
         # We first get positions of haloes in this catalogue, store their
         # radial distance and normalise them to unit vectors.
-        pos = self.position(in_initial=False, cartesian=True)
+        if in_rsp:
+            pos = self.redshift_space_position(cartesian=True)
+        else:
+            pos = self.position(in_initial=False, cartesian=True)
         raddist = numpy.linalg.norm(pos, axis=1)
         pos /= raddist.reshape(-1, 1)
         # We convert RAdec query positions to unit vectors. If no radial
@@ -329,7 +361,8 @@ class ClumpsCatalogue(BaseCatalogue):
             for p in ("x", "y", "z"):
                 self._data[p] -= 0.5
             names = ["x", "y", "z", "mass_cl", "totpartmass", "rho0", "r200c",
-                     "r500c", "m200c", "m500c", "r200m", "m200m"]
+                     "r500c", "m200c", "m500c", "r200m", "m200m",
+                     "vx", "vy", "vz"]
             self._data = self.box.convert_from_boxunits(self._data, names)
             if maxdist is not None:
                 dist = numpy.sqrt(self._data["x"]**2 + self._data["y"]**2
@@ -420,7 +453,8 @@ class HaloCatalogue(BaseCatalogue):
             for p in ("x", "y", "z"):
                 self._data[p] -= 0.5
             names = ["x", "y", "z", "M", "totpartmass", "rho0", "r200c",
-                     "r500c", "m200c", "m500c", "r200m", "m200m"]
+                     "r500c", "m200c", "m500c", "r200m", "m200m",
+                     "vx", "vy", "vz"]
             self._data = self.box.convert_from_boxunits(self._data, names)
 
             if load_initial:

csiborgtools/read/overlap_summary.py

@@ -21,7 +21,6 @@ from os.path import isfile
 import numpy
 from tqdm import tqdm
 
-
 ###############################################################################
 #                         Overlap of two simulations                          #
 ###############################################################################
@@ -602,6 +601,11 @@ class NPairsOverlap:
         return len(self.pairs)
 
 
+###############################################################################
+#                       Various support functions.                            #
+###############################################################################
+
+
 def binned_resample_mean(x, y, prob, bins, nresample=50, seed=42):
     """
     Calculate binned average of `y` by MC resampling. Each point is kept with

csiborgtools/read/utils.py

@@ -80,6 +80,51 @@ 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):
+    r"""
+    Convert real-space position to redshift space position.
+
+    Parameters
+    ----------
+    pos : 2-dimensional array `(nsamples, 3)`
+        Real-space Cartesian position components.
+    vel : 2-dimensional array `(nsamples, 3)`
+        Cartesian velocity components.
+    origin : 1-dimensional array `(3,)`
+        Origin of the coordinate system in the `pos` reference frame.
+    box : py:class:`csiborg.read.BoxUnits`
+        Box units.
+    in_box_units: bool
+        Whether `pos` and `vel` are in box units. If not, position is assumed
+        to be in :math:`\mathrm{Mpc}`, velocity in
+        :math:`\mathrm{km} \mathrm{s}^{-1}` and math:`h=0.705`, or otherwise
+        matching the box.
+    make_copy : bool, optional
+        Whether to make a copy of `pos` before modifying it.
+
+    Returns
+    -------
+    pos : 2-dimensional array `(nsamples, 3)`
+        Redshift-space Cartesian position components, with an observer assumed
+        at the `origin`.
+    """
+    a = box._aexp
+    H0 = box.box_H0 if in_box_units else box.H0
+
+    if make_copy:
+        pos = numpy.copy(pos)
+    for i in range(3):
+        pos[:, i] -= origin[i]
+
+    norm2 = numpy.sum(pos**2, axis=1)
+    dot = numpy.einsum("ij,ij->i", pos, vel)
+    pos *= (1 + a / H0 * dot / norm2).reshape(-1, 1)
+
+    for i in range(3):
+        pos[:, i] += origin[i]
+    return pos
+
+
 ###############################################################################
 #                          Array manipulation                                 #
 ###############################################################################