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
+    Evaluate a scalar field at Cartesian coordinates.
-    interpolation.
 
     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))]
+
+    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.,
+def observer_vobs(velocity_field):
-              verbose=True):
+    """
+    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

scripts/field_prop.py

@@ -61,54 +61,30 @@ def density_field(nsim, parser_args, to_save=True):
     nsnap = max(paths.get_snapshots(nsim))
     box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
     parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
-
     gen = csiborgtools.field.DensityField(box, parser_args.MAS)
+
+    if parser_args.kind == "density":
+        field = gen(parts, parser_args.grid, in_rsp=False,
+                    verbose=parser_args.verbose)
+        if parser_args.in_rsp:
+            field = csiborgtools.field.field2rsp(*field, parts=parts, box=box,
+                                                 verbose=parser_args.verbose)
+    else:
         field = gen(parts, parser_args.grid, in_rsp=parser_args.in_rsp,
                     verbose=parser_args.verbose)
 
+    if parser_args.smooth_scale > 0:
+        field = csiborgtools.field.smoothen_field(
+            field, parser_args.smooth_scale, box.boxsize * box.h, threads=1)
+
     if to_save:
-        fout = paths.field("density", parser_args.MAS, parser_args.grid,
+        fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
-                           nsim, parser_args.in_rsp)
+                           nsim, parser_args.in_rsp, parser_args.smooth_scale)
         print(f"{datetime.now()}: saving output to `{fout}`.")
         numpy.save(fout, field)
     return field
 
 
-def density_field_smoothed(nsim, parser_args, to_save=True):
-    """
-    Calculate the smoothed density field in the CSiBORG simulation. The
-    unsmoothed density field must already be precomputed.
-
-    Parameters
-    ----------
-    nsim : int
-        Simulation index.
-    parser_args : argparse.Namespace
-        Parsed arguments.
-    to_save : bool, optional
-        Whether to save the output to disk.
-
-    Returns
-    -------
-    smoothed_density : 3-dimensional array
-    """
-    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsnap = max(paths.get_snapshots(nsim))
-    box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
-
-    # Load the real space overdensity field
-    rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
-                                 nsim, in_rsp=False))
-    rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
-                                            box.boxsize, threads=1)
-    if to_save:
-        fout = paths.field("density", parser_args.MAS, parser_args.grid,
-                           nsim, parser_args.in_rsp, parser_args.smooth_scale)
-        print(f"{datetime.now()}: saving output to `{fout}`.")
-        numpy.save(fout, rho)
-    return rho
-
-
 ###############################################################################
 #                            Velocity field                                   #
 ###############################################################################
@@ -185,7 +161,7 @@ def potential_field(nsim, parser_args, to_save=True):
                                  nsim, in_rsp=False))
     if parser_args.smooth_scale > 0:
         rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
-                                                box.boxsize, threads=1)
+                                                box.boxsize * box.h, threads=1)
     rho = density_gen.overdensity_field(rho)
     # Calculate the real space potentiel field
     gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
@@ -281,7 +257,7 @@ def environment_field(nsim, parser_args, to_save=True):
                                  nsim, in_rsp=False))
     if parser_args.smooth_scale > 0:
         rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
-                                                box.boxsize, threads=1)
+                                                box.boxsize * box.h, threads=1)
     rho = density_gen.overdensity_field(rho)
     # Calculate the real space tidal tensor field, delete overdensity.
     if parser_args.verbose:
@@ -339,27 +315,27 @@ if __name__ == "__main__":
     parser.add_argument("--nsims", type=int, nargs="+", default=None,
                         help="IC realisations. `-1` for all simulations.")
     parser.add_argument("--kind", type=str,
-                        choices=["density", "velocity", "radvel", "potential",
+                        choices=["density", "rspdensity", "velocity", "radvel",
-                                 "environment"],
+                                 "potential", "environment"],
                         help="What derived field to calculate?")
     parser.add_argument("--MAS", type=str,
                         choices=["NGP", "CIC", "TSC", "PCS"])
     parser.add_argument("--grid", type=int, help="Grid resolution.")
     parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
                         help="Calculate in RSP?")
-    parser.add_argument("--smooth_scale", type=float, default=0)
+    parser.add_argument("--smooth_scale", type=float, default=0,
+                        help="Smoothing scale in Mpc/h.")
     parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
                         help="Verbosity flag for reading in particles.")
+    parser.add_argument("--simname", type=str, default="csiborg",
+                        help="Verbosity flag for reading in particles.")
     parser_args = parser.parse_args()
     comm = MPI.COMM_WORLD
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     nsims = get_nsims(parser_args, paths)
 
     def main(nsim):
-        if parser_args.kind == "density":
+        if parser_args.kind == "density" or parser_args.kind == "rspdensity":
-            if parser_args.smooth_scale > 0:
-                density_field_smoothed(nsim, parser_args)
-            else:
             density_field(nsim, parser_args)
         elif parser_args.kind == "velocity":
             velocity_field(nsim, parser_args)

scripts_plots/plot_data.py

@@ -19,6 +19,7 @@ from argparse import ArgumentParser
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy
+from h5py import File
 import healpy
 
 import scienceplots  # noqa
@@ -244,7 +245,8 @@ def load_field(kind, nsim, grid, MAS, in_rsp=False, smooth_scale=None):
 
 
 def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
-                         highres_only=True, slice_find=None, pdf=False):
+                         vel_component=0, highres_only=True, slice_find=None,
+                         pdf=False):
     r"""
     Plot the mean projected field, however can also plot a single slice.
 
@@ -262,6 +264,8 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
         Smoothing scale in :math:`\mathrm{Mpc} / h`.
     MAS : str, optional
         Mass assignment scheme.
+    vel_component : int, optional
+        Which velocity field component to plot.
     highres_only : bool, optional
         Whether to only plot the high-resolution region.
     slice_find : float, optional
@@ -283,12 +287,16 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
                            smooth_scale=smooth_scale)
         density_gen = csiborgtools.field.DensityField(box, MAS)
         field = density_gen.overdensity_field(field) + 1
+    elif kind == "borg_density":
+        field = File(paths.borg_mcmc(nsim), 'r')
+        field = field["scalars"]["BORG_final_density"][...]
     else:
         field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp,
                            smooth_scale=smooth_scale)
 
     if kind == "velocity":
-        field = field[0, ...]
+        field = field[vel_component, ...]
+        field = box.box2vel(field)
 
     if highres_only:
         csiborgtools.field.fill_outside(field, numpy.nan, rmax=155.5,
@@ -552,22 +560,23 @@ if __name__ == "__main__":
     if False:
         plot_hmf(pdf=False)
 
-    if True:
+    if False:
         kind = "overdensity"
         grid = 1024
         plot_sky_distribution(kind, 7444, grid, nside=64,
                               plot_groups=False, dmin=45, dmax=60,
                               plot_halos=5e13, volume_weight=True)
 
-    if False:
+    if True:
-        kind = "density"
+        kind = "overdensity"
         grid = 256
         smooth_scale = 0
         # plot_projected_field("overdensity", 7444, grid, in_rsp=True,
         #                      highres_only=False)
         plot_projected_field(kind, 7444, grid, in_rsp=False,
                              smooth_scale=smooth_scale, slice_find=0.5,
-                             highres_only=False)
+                             MAS="PCS",
+                             highres_only=True)
 
     if False:
         paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)