New RSP density distinction (#72)

* Edit docs

* Fix kernel units

* Add NGP interpolation

* Add Vobs but not implemented

* Add BORG density option

* Add BORG mcmc path

* Organise imports

* Add new density field distinction

csiborgtools/__init__.py

@@ -17,6 +17,8 @@ from csiborgtools import clustering, field, fits, match, read  # noqa
 # Arguments to csiborgtools.read.Paths.
 paths_glamdring = {"srcdir": "/mnt/extraspace/hdesmond/",
                    "postdir": "/mnt/extraspace/rstiskalek/CSiBORG/",
+                   "BORG_dir": "/mnt/extraspace/rstiskalek/BORG/",
+                   "BORG_final_density": "/users/hdesmond/BORG_final/",
                    "quijote_dir": "/mnt/extraspace/rstiskalek/Quijote",
                    }
 

csiborgtools/field/__init__.py

@@ -17,9 +17,10 @@ from warnings import warn
 try:
     import MAS_library as MASL  # noqa
 
-    from .density import DensityField, PotentialField, VelocityField, TidalTensorField  # noqa
+    from .density import (DensityField, PotentialField,  # noqa
+                          TidalTensorField, VelocityField)
     from .interp import (evaluate_cartesian, evaluate_sky, field2rsp,  # noqa
-                         make_sky, fill_outside)  # noqa
+                         fill_outside, make_sky, observer_vobs)
     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

@@ -17,13 +17,13 @@ Density field and cross-correlation calculations.
 """
 from abc import ABC
 
-import numpy
-
 import MAS_library as MASL
+import numpy
 from numba import jit
 from tqdm import trange
 
 from ..read.utils import real2redshift
+from .interp import divide_nonzero
 from .utils import force_single_precision
 
 
@@ -249,7 +249,7 @@ class VelocityField(BaseField):
                                        / numpy.sqrt(px**2 + py**2 + pz**2))
         return radvel
 
-    def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30,
+    def __call__(self, parts, grid, flip_xz=True, nbatch=30,
                  verbose=True):
         """
         Calculate the velocity field using a Pylians routine [1, 2].
@@ -263,8 +263,6 @@ class VelocityField(BaseField):
             Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`.
         grid : int
             Grid size.
-        mpart : float
-            Particle mass.
         flip_xz : bool, optional
             Whether to flip the `x` and `z` coordinates.
         nbatch : int, optional
@@ -287,6 +285,7 @@ class VelocityField(BaseField):
         rho_vely = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
         rho_velz = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
         rho_vel = [rho_velx, rho_vely, rho_velz]
+        cellcounts = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
 
         nparts = parts.shape[0]
         batch_size = nparts // nbatch
@@ -302,16 +301,22 @@ class VelocityField(BaseField):
             if flip_xz:
                 pos[:, [0, 2]] = pos[:, [2, 0]]
                 vel[:, [0, 2]] = vel[:, [2, 0]]
-            vel *= mass.reshape(-1, 1) / mpart
+            vel *= mass.reshape(-1, 1)
 
             for i in range(3):
                 MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[:, i],
                         verbose=False)
+
+            MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass,
+                    verbose=False)
             if end == nparts:
                 break
             start = end
 
-        return numpy.stack(rho_vel)
+        for i in range(3):
+            divide_nonzero(rho_vel[i], cellcounts)
+
+        return numpy.stack([rho_velx, rho_vely, rho_velz])
 
 
 ###############################################################################

csiborgtools/field/interp.py

@@ -24,10 +24,9 @@ from ..read.utils import radec_to_cartesian, real2redshift
 from .utils import force_single_precision
 
 
-def evaluate_cartesian(*fields, pos):
+def evaluate_cartesian(*fields, pos, interp="CIC"):
     """
-    Evaluate a scalar field at Cartesian coordinates using CIC
-    interpolation.
+    Evaluate a scalar field at Cartesian coordinates.
 
     Parameters
     ----------
@@ -36,19 +35,29 @@ def evaluate_cartesian(*fields, pos):
     pos : 2-dimensional array of shape `(n_samples, 3)`
         Positions to evaluate the density field. Assumed to be in box
         units.
+    interp : str, optional
+        Interpolation method. Can be either `CIC` or `NGP`.
 
     Returns
     -------
     interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
     """
+    assert interp in ["CIC", "NGP"]
     boxsize = 1.
     pos = force_single_precision(pos, "pos")
 
     nsamples = pos.shape[0]
     interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32)
                      for __ in range(len(fields))]
-    for i, field in enumerate(fields):
-        MASL.CIC_interp(field, boxsize, pos, interp_fields[i])
+
+    if interp == "CIC":
+        for i, field in enumerate(fields):
+            MASL.CIC_interp(field, boxsize, pos, interp_fields[i])
+    else:
+        pos = numpy.floor(pos * fields[0].shape[0]).astype(numpy.int32)
+        for i, field in enumerate(fields):
+            for j in range(nsamples):
+                interp_fields[i][j] = field[pos[j, 0], pos[j, 1], pos[j, 2]]
 
     if len(fields) == 1:
         return interp_fields[0]
@@ -168,8 +177,30 @@ def divide_nonzero(field0, field1):
                     field0[i, j, k] /= field1[i, j, k]
 
 
-def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
-              verbose=True):
+def observer_vobs(velocity_field):
+    """
+    Calculate the observer velocity from a velocity field. Assumes the observer
+    is in the centre of the box.
+
+    Parameters
+    ----------
+    velocity_field : 4-dimensional array of shape `(3, grid, grid, grid)`
+        Velocity field to calculate the observer velocity from.
+
+    Returns
+    -------
+    vobs : 1-dimensional array of shape `(3,)`
+        Observer velocity in units of `velocity_field`.
+    """
+    pos = numpy.asanyarray([0.5, 0.5, 0.5]).reshape(1, 3)
+    vobs = numpy.full(3, numpy.nan, dtype=numpy.float32)
+    for i in range(3):
+        vobs[i] = evaluate_cartesian(velocity_field[i, ...], pos=pos)[0]
+    return vobs
+
+
+def field2rsp(*fields, parts, vobs, box, nbatch=30, flip_partsxz=True,
+              init_value=0., verbose=True):
     """
     Forward model real space scalar fields to redshift space. Attaches field
     values to particles, those are then moved to redshift space and from their
@@ -182,6 +213,8 @@ def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
     parts : 2-dimensional array of shape `(n_parts, 6)`
         Particle positions and velocities in real space. Must be organised as
         `x, y, z, vx, vy, vz`.
+    vobs : 1-dimensional array of shape `(3,)`
+        Observer velocity in units matching `parts`.
     box : :py:class:`csiborgtools.read.CSiBORGBox`
         The simulation box information and transformations.
     nbatch : int, optional
@@ -199,6 +232,7 @@ def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
     -------
     rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
     """
+    raise NotImplementedError("Figure out what to do with Vobs.")
     nfields = len(fields)
     # Check that all fields have the same shape.
     if nfields > 1:

csiborgtools/read/paths.py

@@ -31,16 +31,21 @@ class Paths:
         Path to the folder where the RAMSES outputs are stored.
     postdir: str, optional
         Path to the folder where post-processed files are stored.
+    borg_dir : str, optional
+        Path to the folder where BORG MCMC chains are stored.
     quiote_dir : str, optional
         Path to the folder where Quijote simulations are stored.
     """
     _srcdir = None
     _postdir = None
+    _borg_dir = None
     _quijote_dir = None
 
-    def __init__(self, srcdir=None, postdir=None, quijote_dir=None):
+    def __init__(self, srcdir=None, postdir=None, borg_dir=None,
+                 quijote_dir=None):
         self.srcdir = srcdir
         self.postdir = postdir
+        self.borg_dir = borg_dir
         self.quijote_dir = quijote_dir
 
     @staticmethod
@@ -68,6 +73,26 @@ class Paths:
         self._check_directory(path)
         self._srcdir = path
 
+    @property
+    def borg_dir(self):
+        """
+        Path to the folder where BORG MCMC chains are stored.
+
+        Returns
+        -------
+        path : str
+        """
+        if self._borg_dir is None:
+            raise ValueError("`borg_dir` is not set!")
+        return self._borg_dir
+
+    @borg_dir.setter
+    def borg_dir(self, path):
+        if path is None:
+            return
+        self._check_directory(path)
+        self._borg_dir = path
+
     @property
     def quijote_dir(self):
         """
@@ -146,6 +171,21 @@ class Paths:
             return nsim
         return f"{str(nobs).zfill(2)}{str(nsim).zfill(3)}"
 
+    def borg_mcmc(self, nsim):
+        """
+        Path to the BORG MCMC chain file.
+
+        Parameters
+        ----------
+        nsim : int
+            IC realisation index.
+
+        Returns
+        -------
+        path : str
+        """
+        return join(self.borg_dir, "mcmc", f"mcmc_{nsim}.h5")
+
     def mmain(self, nsnap, nsim):
         """
         Path to the `mmain` CSiBORG files of summed substructure.
@@ -373,7 +413,7 @@ class Paths:
             IC realisation index.
         in_rsp : bool
             Whether the calculation is performed in redshift space.
-        smooth_scale : float
+        smooth_scale : float, optional
             Smoothing scale in :math:`\mathrm{Mpc}/h`
 
         Returns