Update units to be consistent. (#78)

* Fix Quijote units

* Updates to units

* Fix how things are loaded

* Updating definitions & conventions

* Clear up how fiducial observers in quijote work

* Refactorize array manip

* Move function definition

* More code refactoring

* Remove unused argument

* Remove `convert_from_box`

* Make a note

* Converting particle units

* Add notes about units

* Remove box constants

* Add rho_crit0

* Fix spherical overdensity mass units

* Refactor more code

* Edit catalogue kwargs

* Edit the docstring

* Edit bounds

* Add new checks for empty array

* Remove unused import

* Remove old code

* Remove old function

* Update real 2 redshift

* Clear up the RSP conv

* Add comments

* Add some units

csiborgtools/field/density.py

@@ -181,6 +181,8 @@ class DensityField(BaseField):
                 vel[:, [0, 2]] = vel[:, [2, 0]]
 
             if in_rsp:
+                raise NotImplementedError("Redshift space needs to be fixed.")
+                # TODO change how called + units.
                 pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], self.box,
                                     in_box_units=True, periodic_wrap=True,
                                     make_copy=False)

csiborgtools/field/interp.py

@@ -255,6 +255,8 @@ def field2rsp(*fields, parts, vobs, box, nbatch=30, flip_partsxz=True,
             pos[:, [0, 2]] = pos[:, [2, 0]]
             vel[:, [0, 2]] = vel[:, [2, 0]]
         # Then move the particles to redshift space.
+        # TODO here the function is now called differently and pos assumes
+        # different units.
         rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
                                 in_box_units=True, periodic_wrap=True,
                                 make_copy=True)

csiborgtools/fits/__init__.py

@@ -14,4 +14,3 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 from .halo import (Halo, delta2ncells, center_of_mass,  # noqa
                    periodic_distance, shift_to_center_of_box, number_counts)  # noqa
-from .utils import split_jobs  # noqa

csiborgtools/fits/halo.py

@@ -20,6 +20,11 @@ from numba import jit
 from scipy.optimize import minimize
 
 
+GRAV = 6.6743e-11               # m^3 kg^-1 s^-2
+MSUN = 1.988409870698051e+30    # kg
+MPC2M = 3.0856775814671916e+22  # 1 Mpc is this many meters
+
+
 class BaseStructure(ABC):
     """
     Basic structure object for handling operations on its particles.
@@ -85,7 +90,7 @@ class BaseStructure(ABC):
         """
         return numpy.vstack([self[p] for p in ("vx", "vy", "vz")]).T
 
-    def spherical_overdensity_mass(self, delta_mult, kind="crit", tol=1e-8,
+    def spherical_overdensity_mass(self, delta_mult, kind="crit", rtol=1e-8,
                                    maxiter=100, npart_min=10):
         r"""
         Calculate spherical overdensity mass and radius via the iterative
@@ -97,7 +102,7 @@ class BaseStructure(ABC):
             Overdensity multiple.
         kind : str, optional
             Either `crit` or `matter`, for critical or matter overdensity
-        tol : float, optional
+        rtol : float, optional
             Tolerance for the change in the center of mass or radius.
         maxiter : int, optional
             Maximum number of iterations.
@@ -107,33 +112,34 @@ class BaseStructure(ABC):
         Returns
         -------
         mass :  float
-            The requested spherical overdensity mass.
+            The requested spherical overdensity mass in :math:`M_\odot / h`.
         rad : float
-            The radius of the sphere enclosing the requested overdensity.
+            The radius of the sphere enclosing the requested overdensity in box
+            units.
         cm : 1-dimensional array of shape `(3, )`
             The center of mass of the sphere enclosing the requested
-            overdensity.
+            overdensity in box units.
         """
         assert kind in ["crit", "matter"]
-        rho = delta_mult * self.box.box_rhoc
+
+        # Calculate density based on the provided kind
+        rho = delta_mult * self.box.rho_crit0
         if kind == "matter":
             rho *= self.box.Om
-        pos = self.pos
-        mass = self["M"]
 
-        # Initial guesses
+        pos, mass = self.pos, self["M"]
+
+        # Initial estimates for center of mass and radius
         init_cm = center_of_mass(pos, mass, boxsize=1)
-        init_rad = mass_to_radius(numpy.sum(mass), rho) * 1.5
+        init_rad = self.box.mpc2box(mass_to_radius(numpy.sum(mass), rho) * 1.5)
 
-        rad = init_rad
+        rad, cm = init_rad, numpy.copy(init_cm)
-        cm = numpy.copy(init_cm)
 
-        success = False
+        for _ in range(maxiter):
-        for __ in range(maxiter):
-            # Calculate the distance of each particle from the current guess.
             dist = periodic_distance(pos, cm, boxsize=1)
             within_rad = dist <= rad
-            # Heuristic reset if there are too few enclosed particles.
+
+            # Heuristic reset if too few enclosed particles
             if numpy.sum(within_rad) < npart_min:
                 js = numpy.random.choice(len(self), len(self), replace=True)
                 cm = center_of_mass(pos[js], mass[js], boxsize=1)
@@ -141,41 +147,40 @@ class BaseStructure(ABC):
                 dist = periodic_distance(pos, cm, boxsize=1)
                 within_rad = dist <= rad
 
-            # Calculate the enclosed mass for the current CM and radius.
+                # If there are still too few particles, then skip this
-            enclosed_mass = numpy.sum(mass[within_rad])
+                # iteration.
+                if numpy.sum(within_rad) < npart_min:
+                    continue
 
-            # Calculate the new CM and radius from this mass.
+            enclosed_mass = numpy.sum(mass[within_rad])
-            new_rad = mass_to_radius(enclosed_mass, rho)
+            new_rad = self.box.mpc2box(mass_to_radius(enclosed_mass, rho))
             new_cm = center_of_mass(pos[within_rad], mass[within_rad],
                                     boxsize=1)
 
-            # Update the CM and radius
+            # Check convergence based on center of mass and radius
-            prev_cm, cm = cm, new_cm
+            cm_conv = numpy.linalg.norm(cm - new_cm) < rtol
-            prev_rad, rad = rad, new_rad
+            rad_conv = abs(rad - new_rad) < rtol
-
-            # Check if the change in CM and radius is small enough.
-            dcm = numpy.linalg.norm(cm - prev_cm)
-            drad = abs(rad - prev_rad)
-            if dcm < tol or drad < tol:
-                success = True
-                break
-
-        if not success:
-            return numpy.nan, numpy.nan, numpy.full(3, numpy.nan)
 
+            if cm_conv or rad_conv:
                 return enclosed_mass, rad, cm
 
+            cm, rad = new_cm, new_rad
+
+        # Return NaN values if no convergence after max iterations
+        return numpy.nan, numpy.nan, numpy.full(3, numpy.nan, numpy.float32)
+
     def angular_momentum(self, ref, rad, npart_min=10):
-        """
+        r"""
         Calculate angular momentum around a reference point using all particles
-        within a radius. The angular momentum is returned in box units.
+        within a radius. Units are
+        :math:`(M_\odot / h) (\mathrm{Mpc} / h) \mathrm{km} / \mathrm{s}`.
 
         Parameters
         ----------
         ref : 1-dimensional array of shape `(3, )`
-            Reference point.
+            Reference point in box units.
         rad : float
-            Radius around the reference point.
+            Radius around the reference point in box units.
         npart_min : int, optional
             Minimum number of enclosed particles for a radius to be
             considered trustworthy.
@@ -184,16 +189,21 @@ class BaseStructure(ABC):
         -------
         angmom : 1-dimensional array or shape `(3, )`
         """
-        pos = self.pos
+        # Calculate the distance of each particle from the reference point.
-        mask = periodic_distance(pos, ref, boxsize=1) < rad
+        distances = periodic_distance(self.pos, ref, boxsize=1)
-        if numpy.sum(mask) < npart_min:
-            return numpy.full(3, numpy.nan)
 
-        mass = self["M"][mask]
+        # Filter particles within the provided radius.
-        pos = pos[mask]
+        mask = distances < rad
-        vel = self.vel[mask]
+        if numpy.sum(mask) < npart_min:
-        # Velocitities in the object CM frame
+            return numpy.full(3, numpy.nan, numpy.float32)
+
+        mass, pos, vel = self["M"][mask], self.pos[mask], self.vel[mask]
+
+        # Convert positions to Mpc / h and center around the reference point.
+        pos = self.box.box2mpc(pos) - ref
+        # Adjust velocities to be in the CM frame.
         vel -= numpy.average(vel, axis=0, weights=mass)
+        # Calculate angular momentum.
         return numpy.sum(mass[:, numpy.newaxis] * numpy.cross(pos, vel),
                          axis=0)
 
@@ -205,9 +215,9 @@ class BaseStructure(ABC):
         Parameters
         ----------
         ref : 1-dimensional array of shape `(3, )`
-            Reference point.
+            Reference point in box units.
         rad : float
-            Radius around the reference point.
+            Radius around the reference point in box units.
 
         Returns
         -------
@@ -219,12 +229,18 @@ class BaseStructure(ABC):
         Bullock, J. S.; Dekel, A.;  Kolatt, T. S.; Kravtsov, A. V.;
         Klypin, A. A.; Porciani, C.; Primack, J. R.
         """
-        pos = self.pos
+        # Filter particles within the provided radius
-        mask = periodic_distance(pos, ref, boxsize=1) < rad
+        mask = periodic_distance(self.pos, ref, boxsize=1) < rad
-        mass = numpy.sum(self["M"][mask])
+        # Calculate the total mass of the enclosed particles
-        circvel = numpy.sqrt(self.box.box_G * mass / rad)
+        enclosed_mass = numpy.sum(self["M"][mask])
-        angmom_norm = numpy.linalg.norm(self.angular_momentum(ref, rad))
+        # Convert the radius from box units to Mpc/h
-        return angmom_norm / (numpy.sqrt(2) * mass * circvel * rad)
+        rad_mpc = self.box.box2mpc(rad)
+        # Circular velocity in km/s
+        circvel = (GRAV * enclosed_mass * MSUN / (rad_mpc * MPC2M))**0.5 * 1e-3
+        # Magnitude of the angular momentum
+        l_norm = numpy.linalg.norm(self.angular_momentum(ref, rad))
+        # Compute and return the Bullock spin parameter
+        return l_norm / (numpy.sqrt(2) * enclosed_mass * circvel * rad_mpc)
 
     def nfw_concentration(self, ref, rad, conc_min=1e-3, npart_min=10):
         """
@@ -234,9 +250,9 @@ class BaseStructure(ABC):
         Parameters
         ----------
         ref : 1-dimensional array of shape `(3, )`
-            Reference point.
+            Reference point in box units.
         rad : float
-            Radius around the reference point.
+            Radius around the reference point in box units.
         conc_min : float
             Minimum concentration limit.
         npart_min : int, optional
@@ -247,42 +263,43 @@ class BaseStructure(ABC):
         -------
         conc : float
         """
-        pos = self.pos
+        dist = periodic_distance(self.pos, ref, boxsize=1)
-        dist = periodic_distance(pos, ref, boxsize=1)
         mask = dist < rad
+
         if numpy.sum(mask) < npart_min:
             return numpy.nan
 
-        dist = dist[mask]
+        dist, weight = dist[mask], self["M"][mask]
-        weight = self["M"][mask]
         weight /= numpy.mean(weight)
 
-        # We do the minimization in log space
+        # Objective function for minimization
-        def negll_nfw_concentration(log_c, xs, weight):
+        def negll_nfw_concentration(log_c, xs, w):
             c = 10**log_c
             ll = xs / (1 + c * xs)**2 * c**2
             ll *= (1 + c) / ((1 + c) * numpy.log(1 + c) - c)
-            ll = numpy.sum(numpy.log(weight * ll))
+            ll = numpy.sum(numpy.log(w * ll))
             return -ll
 
-        res = minimize(negll_nfw_concentration, x0=1.5,
+        initial_guess = 1.5
+        res = minimize(negll_nfw_concentration, x0=initial_guess,
                        args=(dist / rad, weight, ), method='Nelder-Mead',
                        bounds=((numpy.log10(conc_min), 5),))
 
         if not res.success:
             return numpy.nan
 
-        res = 10**res["x"][0]
+        conc_value = 10**res["x"][0]
-        if res < conc_min or numpy.isclose(res, conc_min):
+        if conc_value < conc_min or numpy.isclose(conc_value, conc_min):
             return numpy.nan
 
-        return res
+        return conc_value
 
     def __getitem__(self, key):
-        keys = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']
+        key_to_index = {'x': 0, 'y': 1, 'z': 2,
-        if key not in keys:
+                        'vx': 3, 'vy': 4, 'vz': 5, 'M': 6}
+        if key not in key_to_index:
             raise RuntimeError(f"Invalid key `{key}`!")
-        return self.particles[:, keys.index(key)]
+        return self.particles[:, key_to_index[key]]
 
     def __len__(self):
         return self.particles.shape[0]
@@ -304,7 +321,6 @@ class Halo(BaseStructure):
 
     def __init__(self, particles, box):
         self.particles = particles
-        # self.info = info
         self.box = box
 
 

csiborgtools/fits/utils.py

@@ -1,43 +0,0 @@
-# Copyright (C) 2022 Richard Stiskalek, Deaglan Bartlett
-# This program is free software; you can redistribute it and/or modify it
-# under the terms of the GNU General Public License as published by the
-# Free Software Foundation; either version 3 of the License, or (at your
-# option) any later version.
-#
-# This program is distributed in the hope that it will be useful, but
-# WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
-# Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""Fitting utility functions."""
-
-import numpy
-
-
-def split_jobs(njobs, ncpu):
-    """
-    Split `njobs` amongst `ncpu`.
-
-    Parameters
-    ----------
-    njobs : int
-        Number of jobs.
-    ncpu : int
-        Number of CPUs.
-
-    Returns
-    -------
-    jobs : list of lists of integers
-        Outer list of each CPU and inner lists for CPU's jobs.
-    """
-    njobs_per_cpu, njobs_remainder = divmod(njobs, ncpu)
-    jobs = numpy.arange(njobs_per_cpu * ncpu).reshape((njobs_per_cpu, ncpu)).T
-
-    jobs = jobs.tolist()
-    for i in range(njobs_remainder):
-        jobs[i].append(njobs_per_cpu * ncpu + i)
-
-    return jobs

csiborgtools/match/__init__.py

@@ -14,6 +14,4 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 from .match import (ParticleOverlap, RealisationsMatcher,  # noqa
                     calculate_overlap, calculate_overlap_indxs,
-                    cosine_similarity)
+                    cosine_similarity, find_neighbour)
-from .nearest_neighbour import find_neighbour  # noqa
-from .utils import concatenate_parts  # noqa

csiborgtools/match/match.py

@@ -1008,3 +1008,43 @@ def radius_neighbours(knn, X, radiusX, radiusKNN, nmult=1.0,
             indxs[i] = indxs[i].astype(numpy.int32)
 
     return numpy.asarray(indxs, dtype=object)
+
+
+def find_neighbour(nsim0, cats):
+    """
+    Find the nearest neighbour of halos from a reference catalogue indexed
+    `nsim0` in the remaining simulations.
+
+    Parameters
+    ----------
+    nsim0 : int
+        Index of the reference simulation.
+    cats : dict
+        Dictionary of halo catalogues. Keys must be the simulation indices.
+
+    Returns
+    -------
+    dists : 2-dimensional array of shape `(nhalos, len(cats) - 1)`
+        Distances to the nearest neighbour.
+    cross_hindxs : 2-dimensional array of shape `(nhalos, len(cats) - 1)`
+        Halo indices of the nearest neighbour.
+    """
+    cat0 = cats[nsim0]
+    X = cat0.position(in_initial=False, subtract_observer=True)
+
+    nhalos = X.shape[0]
+    num_cats = len(cats) - 1
+
+    dists = numpy.full((nhalos, num_cats), numpy.nan, dtype=numpy.float32)
+    cross_hindxs = numpy.full((nhalos, num_cats), numpy.nan, dtype=numpy.int32)
+
+    # Filter out the reference simulation from the dictionary
+    filtered_cats = {k: v for k, v in cats.items() if k != nsim0}
+
+    for i, catx in enumerate(filtered_cats):
+        dist, ind = catx.nearest_neighbours(X, radius=1, in_initial=False,
+                                            knearest=True)
+        dists[:, i] = numpy.ravel(dist)
+        cross_hindxs[:, i] = catx["index"][numpy.ravel(ind)]
+
+    return dists, cross_hindxs

csiborgtools/match/nearest_neighbour.py

@@ -1,56 +0,0 @@
-# Copyright (C) 2022 Richard Stiskalek
-# This program is free software; you can redistribute it and/or modify it
-# under the terms of the GNU General Public License as published by the
-# Free Software Foundation; either version 3 of the License, or (at your
-# option) any later version.
-#
-# This program is distributed in the hope that it will be useful, but
-# WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
-# Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""
-Tools for finding the nearest neighbours of reference simulation haloes from
-cross simulations.
-"""
-import numpy
-
-
-def find_neighbour(nsim0, cats):
-    """
-    Find the nearest neighbour of halos in `cat0` in `catx`.
-
-    Parameters
-    ----------
-    nsim0 : int
-        Index of the reference simulation.
-    cats : dict
-        Dictionary of halo catalogues. Keys must be the simulation indices.
-
-    Returns
-    -------
-    dists : 2-dimensional array of shape `(nhalos, len(cats) - 1)`
-        Distances to the nearest neighbour.
-    cross_hindxs : 2-dimensional array of shape `(nhalos, len(cats) - 1)`
-        Halo indices of the nearest neighbour.
-    """
-    cat0 = cats[nsim0]
-    X = cat0.position(in_initial=False)
-    shape = (X.shape[0], len(cats) - 1)
-    dists = numpy.full(shape, numpy.nan, dtype=numpy.float32)
-    cross_hindxs = numpy.full(shape, numpy.nan, dtype=numpy.int32)
-
-    i = 0
-    for nsimx, catx in cats.items():
-        if nsimx == nsim0:
-            continue
-        dist, ind = catx.nearest_neighbours(X, radius=1, in_initial=False,
-                                            knearest=True)
-        dists[:, i] = dist.reshape(-1,)
-        cross_hindxs[:, i] = catx["index"][ind.reshape(-1,)]
-        i += 1
-
-    return dists, cross_hindxs

csiborgtools/match/utils.py

@@ -1,67 +0,0 @@
-# Copyright (C) 2022 Richard Stiskalek
-# This program is free software; you can redistribute it and/or modify it
-# under the terms of the GNU General Public License as published by the
-# Free Software Foundation; either version 3 of the License, or (at your
-# option) any later version.
-#
-# This program is distributed in the hope that it will be useful, but
-# WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
-# Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""Useful functions."""
-import numpy
-
-
-def concatenate_parts(list_parts, include_velocities=False):
-    """
-    Concatenate a list of particle arrays into a single array.
-
-    Parameters
-    ----------
-    list_parts : list of structured arrays
-        List of particle arrays.
-    include_velocities : bool, optional
-        Whether to include velocities in the output array.
-
-    Returns
-    -------
-    parts_out : structured array
-    """
-    # Count how large array will be needed
-    N = 0
-    for part in list_parts:
-        N += part.size
-    # Infer dtype of positions
-    if list_parts[0]["x"].dtype.char in numpy.typecodes["AllInteger"]:
-        posdtype = numpy.int32
-    else:
-        posdtype = numpy.float32
-
-    # We pre-allocate an empty array. By default, we include just particle
-    # positions, which may be specified by cell IDs if integers, and masses.
-    # Additionally also outputs velocities.
-    if include_velocities:
-        dtype = {
-            "names": ["x", "y", "z", "vx", "vy", "vz", "M"],
-            "formats": [posdtype] * 3 + [numpy.float32] * 4,
-        }
-    else:
-        dtype = {
-            "names": ["x", "y", "z", "M"],
-            "formats": [posdtype] * 3 + [numpy.float32],
-        }
-    parts_out = numpy.full(N, numpy.nan, dtype)
-
-    # Fill it one clump by another
-    start = 0
-    for parts in list_parts:
-        end = start + parts.size
-        for p in dtype["names"]:
-            parts_out[p][start:end] = parts[p]
-        start = end
-
-    return parts_out

csiborgtools/read/__init__.py

@@ -26,6 +26,5 @@ from .pk_summary import PKReader  # noqa
 from .readsim import (MmainReader, CSiBORGReader, QuijoteReader, halfwidth_mask,  # noqa
                       load_halo_particles)  # noqa
 from .tpcf_summary import TPCFReader  # noqa
-from .utils import (M200_to_R200, cartesian_to_radec,  # noqa
+from .utils import (cartesian_to_radec, cols_to_structured, radec_to_cartesian,  # noqa
-                    cols_to_structured, radec_to_cartesian, read_h5,
+                    read_h5, real2redshift)  # noqa
-                    real2redshift)  # noqa

csiborgtools/read/box_units.py

@@ -17,29 +17,12 @@ Simulation box unit transformations.
 """
 from abc import ABC, abstractmethod, abstractproperty
 
-import numpy
 from astropy import constants, units
 from astropy.cosmology import LambdaCDM
 
 from .readsim import CSiBORGReader, QuijoteReader
 
 
-CSIBORG_CONV_NAME = {
-    "length": ["x", "y", "z", "peak_x", "peak_y", "peak_z", "Rs", "rmin",
-               "rmax", "r200c", "r500c", "r200m", "r500m", "x0", "y0", "z0",
-               "lagpatch_size"],
-    "velocity": ["vx", "vy", "vz"],
-    "mass": ["mass_cl", "totpartmass", "m200c", "m500c", "mass_mmain", "M",
-             "m200m", "m500m"],
-    "density": ["rho0"]
-    }
-
-QUIJOTE_CONV_NAME = {
-    "length": ["x", "y", "z", "x0", "y0", "z0", "Rs", "r200c", "r500c",
-               "r200m", "r500m", "lagpatch_size"],
-    "mass": ["group_mass", "totpartmass", "m200c", "m500c", "m200m", "m500m"],
-    }
-
 ###############################################################################
 #                              Base box                                       #
 ###############################################################################
@@ -77,6 +60,18 @@ class BaseBox(ABC):
         """
         return self.cosmo.H0.value
 
+    @property
+    def rho_crit0(self):
+        r"""
+        Present-day critical density in :math:`M_\odot h^2 / \mathrm{cMpc}^3`.
+
+        Returns
+        -------
+        rho_crit0 : float
+        """
+        rho_crit0 = self.cosmo.critical_density0
+        return rho_crit0.to_value(units.solMass / units.Mpc**3)
+
     @property
     def h(self):
         r"""
@@ -86,7 +81,7 @@ class BaseBox(ABC):
         -------
         h : float
         """
-        return self.H0 / 100
+        return self._h
 
     @property
     def Om0(self):
@@ -111,31 +106,70 @@ class BaseBox(ABC):
         pass
 
     @abstractmethod
-    def convert_from_box(self, data, names):
+    def mpc2box(self, length):
         r"""
-        Convert columns named `names` in array `data` from box units to
+        Convert length from :math:`\mathrm{cMpc} / h` to box units.
-        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
-        array is passed by reference and directly modified, even though it is
-        also explicitly returned. Additionally centres the box coordinates on
-        the observer, if they are being transformed.
 
         Parameters
         ----------
-        data : structured array
+        length : float
-            Input array.
+            Length in :math:`\mathrm{cMpc}`
-        names : list of str
-            Columns to be converted.
 
         Returns
         -------
-        data : structured array
+        length : float
-            Input array with converted columns.
+            Length in box units.
+        """
+        pass
+
+    @abstractmethod
+    def box2mpc(self, length):
+        r"""
+        Convert length from box units to :math:`\mathrm{cMpc} / h`.
+
+        Parameters
+        ----------
+        length : float
+            Length in box units.
+
+        Returns
+        -------
+        length : float
+            Length in :math:`\mathrm{cMpc} / h`
+        """
+        pass
+
+    @abstractmethod
+    def solarmass2box(self, mass):
+        r"""
+        Convert mass from :math:`M_\odot / h` to box units.
+
+        Parameters
+        ----------
+        mass : float
+            Mass in :math:`M_\odot / h`.
+
+        Returns
+        -------
+        mass : float
+            Mass in box units.
+        """
+        pass
+
+    @abstractmethod
+    def box2solarmass(self, mass):
+        r"""
+        Convert mass from box units to :math:`M_\odot / h`.
+
+        Parameters
+        ----------
+        mass : float
+            Mass in box units.
+
+        Returns
+        -------
+        mass : float
+            Mass in :math:`M_\odot / h`.
         """
         pass
 
@@ -169,124 +203,31 @@ class CSiBORGBox(BaseBox):
                 "omega_k", "omega_b", "unit_l", "unit_d", "unit_t"]
         for par in pars:
             setattr(self, "_" + par, info[par])
-
+        self._h = self._H0 / 100
-        self._cosmo = LambdaCDM(H0=self._H0, Om0=self._omega_m,
+        self._cosmo = LambdaCDM(H0=100, Om0=self._omega_m,
                                 Ode0=self._omega_l, Tcmb0=2.725 * units.K,
                                 Ob0=self._omega_b)
         self._Msuncgs = constants.M_sun.cgs.value  # Solar mass in grams
 
-    @property
-    def box_G(self):
-        """
-        Gravitational constant :math:`G` in box units. Given everything else
-        it looks like `self.unit_t` is in seconds.
-
-        Returns
-        -------
-        G : float
-        """
-        return constants.G.cgs.value * (self._unit_d * self._unit_t**2)
-
-    @property
-    def box_H0(self):
-        """
-        Present time Hubble constant :math:`H_0` in box units.
-
-        Returns
-        -------
-        H0 : float
-        """
-        return self.H0 * 1e5 / units.Mpc.to(units.cm) * self._unit_t
-
-    @property
-    def box_c(self):
-        """
-        Speed of light in box units.
-
-        Returns
-        -------
-        c : float
-        """
-        return constants.c.cgs.value * self._unit_t / self._unit_l
-
-    @property
-    def box_rhoc(self):
-        """
-        Critical density in box units.
-
-        Returns
-        -------
-        rhoc : float
-        """
-        return 3 * self.box_H0**2 / (8 * numpy.pi * self.box_G)
-
-    def box2kpc(self, length):
-        r"""
-        Convert length from box units to :math:`\mathrm{ckpc}` (with
-        :math:`h=0.705`).
-
-        Parameters
-        ----------
-        length : float
-            Length in box units.
-
-        Returns
-        -------
-        length : float
-            Length in :math:`\mathrm{ckpc}`
-        """
-        return length * (self._unit_l / units.kpc.to(units.cm) / self._aexp)
-
-    def kpc2box(self, length):
-        r"""
-        Convert length from :math:`\mathrm{ckpc}` (with :math:`h=0.705`) to
-        box units.
-
-        Parameters
-        ----------
-        length : float
-            Length in :math:`\mathrm{ckpc}`
-
-        Returns
-        -------
-        length : float
-            Length in box units.
-        """
-        return length / (self._unit_l / units.kpc.to(units.cm) / self._aexp)
-
     def mpc2box(self, length):
-        r"""
+        conv = (self._unit_l / units.kpc.to(units.cm) / self._aexp) * 1e-3
-        Convert length from :math:`\mathrm{cMpc}` (with :math:`h=0.705`) to
+        conv *= self._h
-        box units.
+        return length / conv
-
-        Parameters
-        ----------
-        length : float
-            Length in :math:`\mathrm{cMpc}`
-
-        Returns
-        -------
-        length : float
-            Length in box units.
-        """
-        return self.kpc2box(length * 1e3)
 
     def box2mpc(self, length):
-        r"""
+        conv = (self._unit_l / units.kpc.to(units.cm) / self._aexp) * 1e-3
-        Convert length from box units to :math:`\mathrm{cMpc}` (with
+        conv *= self._h
-        :math:`h=0.705`).
+        return length * conv
 
-        Parameters
+    def solarmass2box(self, mass):
-        ----------
+        conv = (self._unit_d * self._unit_l**3) / self._Msuncgs
-        length : float
+        conv *= self.h
-            Length in box units.
+        return mass / conv
 
-        Returns
+    def box2solarmass(self, mass):
-        -------
+        conv = (self._unit_d * self._unit_l**3) / self._Msuncgs
-        length : float
+        conv *= self.h
-            Length in :math:`\mathrm{ckpc}`
+        return mass * conv
-        """
-        return self.box2kpc(length) * 1e-3
 
     def box2vel(self, vel):
         r"""
@@ -304,105 +245,6 @@ class CSiBORGBox(BaseBox):
         """
         return vel * (1e-2 * self._unit_l / self._unit_t / self._aexp) * 1e-3
 
-    def solarmass2box(self, mass):
-        r"""
-        Convert mass from :math:`M_\odot` (with :math:`h=0.705`) to box units.
-
-        Parameters
-        ----------
-        mass : float
-            Mass in :math:`M_\odot`.
-
-        Returns
-        -------
-        mass : float
-            Mass in box units.
-        """
-        return mass / (self._unit_d * self._unit_l**3) * self._Msuncgs
-
-    def box2solarmass(self, mass):
-        r"""
-        Convert mass from box units to :math:`M_\odot` (with :math:`h=0.705`).
-        It appears that `self.unit_d` is density in units of
-        :math:`\mathrm{g}/\mathrm{cm}^3`.
-
-        Parameters
-        ----------
-        mass : float
-            Mass in box units.
-
-        Returns
-        -------
-        mass : float
-            Mass in :math:`M_\odot`.
-        """
-        return mass * (self._unit_d * self._unit_l**3) / self._Msuncgs
-
-    def box2dens(self, density):
-        r"""
-        Convert density from box units to :math:`M_\odot / \mathrm{Mpc}^3`
-        (with :math:`h=0.705`).
-
-        Parameters
-        ----------
-        density : float
-            Density in box units.
-
-        Returns
-        -------
-        density : float
-            Density in :math:`M_\odot / \mathrm{pc}^3`.
-        """
-        return (density * self._unit_d
-                / self._Msuncgs * (units.Mpc.to(units.cm)) ** 3)
-
-    def dens2box(self, density):
-        r"""
-        Convert density from :math:`M_\odot / \mathrm{Mpc}^3`
-        (with :math:`h=0.705`) to box units.
-
-        Parameters
-        ----------
-        density : float
-            Density in :math:`M_\odot / \mathrm{pc}^3`.
-
-        Returns
-        -------
-        density : float
-            Density in box units.
-        """
-        return (density / self._unit_d * self._Msuncgs
-                / (units.Mpc.to(units.cm)) ** 3)
-
-    def convert_from_box(self, data, names):
-        names = [names] if isinstance(names, str) else names
-        transforms = {"length": self.box2mpc,
-                      "mass": self.box2solarmass,
-                      "velocity": self.box2vel,
-                      "density": self.box2dens}
-
-        for name in names:
-            if name not in data.dtype.names:
-                continue
-
-            # Convert
-            found = False
-            for unittype, suppnames in CSIBORG_CONV_NAME.items():
-                if name in suppnames:
-                    data[name] = transforms[unittype](data[name])
-                    found = True
-                    continue
-            # If nothing found
-            if not found:
-                raise NotImplementedError(
-                    f"Conversion of `{name}` is not defined.")
-
-            # Center at the observer
-            if name in ["x0", "y0", "z0"]:
-                data[name] -= transforms["length"](0.5)
-
-        return data
-
     @property
     def boxsize(self):
         return self.box2mpc(1.)
@@ -428,10 +270,10 @@ class QuijoteBox(BaseBox):
     def __init__(self, nsnap, nsim, paths):
         zdict = {4: 0.0, 3: 0.5, 2: 1.0, 1: 2.0, 0: 3.0}
         assert nsnap in zdict.keys(), f"`nsnap` must be in {zdict.keys()}."
-        self._aexp = 1 / (1 + zdict[nsnap])
-
         info = QuijoteReader(paths).read_info(nsnap, nsim)
-        self._cosmo = LambdaCDM(H0=info["Hubble"], Om0=info["Omega_m"],
+        self._aexp = 1 / (1 + zdict[nsnap])
+        self._h = info["h"]
+        self._cosmo = LambdaCDM(H0=100, Om0=info["Omega_m"],
                                 Ode0=info["Omega_l"], Tcmb0=2.725 * units.K)
         self._info = info
 
@@ -440,35 +282,9 @@ class QuijoteBox(BaseBox):
         return self._info["BoxSize"]
 
     def box2mpc(self, length):
-        r"""
-        Convert length from box units to :math:`\mathrm{cMpc} / h`.
-
-        Parameters
-        ----------
-        length : float
-            Length in box units.
-
-        Returns
-        -------
-        length : float
-            Length in :math:`\mathrm{cMpc} / h`
-        """
         return length * self.boxsize
 
     def mpc2box(self, length):
-        r"""
-        Convert length from :math:`\mathrm{cMpc} / h` to box units.
-
-        Parameters
-        ----------
-        length : float
-            Length in :math:`\mathrm{cMpc} / h`.
-
-        Returns
-        -------
-        length : float
-            Length in box units.
-        """
         return length / self.boxsize
 
     def solarmass2box(self, mass):
@@ -502,33 +318,3 @@ class QuijoteBox(BaseBox):
             Mass in :math:`M_\odot / h`.
         """
         return mass * self._info["TotMass"]
-
-    def convert_from_box(self, data, names):
-        names = [names] if isinstance(names, str) else names
-        transforms = {"length": self.box2mpc,
-                      "mass": self.box2solarmass,
-                      # "velocity": self.box2vel,
-                      # "density": self.box2dens,
-                      }
-
-        for name in names:
-            if name not in data.dtype.names:
-                continue
-
-            # Convert
-            found = False
-            for unittype, suppnames in QUIJOTE_CONV_NAME.items():
-                if name in suppnames:
-                    data[name] = transforms[unittype](data[name])
-                    found = True
-                    continue
-            # If nothing found
-            if not found:
-                raise NotImplementedError(
-                    f"Conversion of `{name}` is not defined.")
-
-            # # Center at the observer
-            # if name in ["x0", "y0", "z0"]:
-            #     data[name] -= transforms["length"](0.5)
-
-        return data

csiborgtools/read/halo_cat.py

@@ -24,6 +24,7 @@ from itertools import product
 from math import floor
 
 import numpy
+
 from readfof import FoF_catalog
 from sklearn.neighbors import NearestNeighbors
 
@@ -57,6 +58,7 @@ class BaseCatalogue(ABC):
 
     @nsim.setter
     def nsim(self, nsim):
+        assert isinstance(nsim, int)
         self._nsim = nsim
 
     @abstractproperty
@@ -98,19 +100,9 @@ class BaseCatalogue(ABC):
         data : structured array
         """
         if self._data is None:
-            raise RuntimeError("Catalogue data not loaded!")
+            raise RuntimeError("`data` is not set!")
         return self._data
 
-    def apply_bounds(self, bounds):
-        for key, (xmin, xmax) in bounds.items():
-            xmin = -numpy.inf if xmin is None else xmin
-            xmax = numpy.inf if xmax is None else xmax
-            if key == "dist":
-                x = self.radial_distance(in_initial=False)
-            else:
-                x = self[key]
-            self._data = self._data[(x > xmin) & (x <= xmax)]
-
     @abstractproperty
     def box(self):
         """
@@ -122,72 +114,156 @@ class BaseCatalogue(ABC):
         """
         pass
 
-    def position(self, in_initial=False, cartesian=True):
+    def load_initial(self, data, paths, simname):
+        """
+        Load initial snapshot fits from the script `fit_init.py`.
+
+        Parameters
+        ----------
+        data : structured array
+            The catalogue to which append the new data.
+        paths : :py:class:`csiborgtools.read.Paths`
+            Paths manager.
+        simname : str
+            Simulation name.
+
+        Returns
+        -------
+        data : structured array
+        """
+        fits = numpy.load(paths.initmatch(self.nsim, simname, "fit"))
+        X, cols = [], []
+
+        for col in fits.dtype.names:
+            if col == "index":
+                continue
+            cols.append(col + "0" if col in ['x', 'y', 'z'] else col)
+            X.append(fits[col])
+
+        data = add_columns(data, X, cols)
+        for p in ('x0', 'y0', 'z0', 'lagpatch_size'):
+            data[p] = self.box.box2mpc(data[p])
+
+        return data
+
+    def load_fitted(self, data, paths, simname):
+        """
+        Load halo fits from the script `fit_halos.py`.
+
+        Parameters
+        ----------
+        data : structured array
+            The catalogue to which append the new data.
+        paths : :py:class:`csiborgtools.read.Paths`
+            Paths manager.
+        simname : str
+            Simulation name.
+
+        Returns
+        -------
+        data : structured array
+        """
+        fits = numpy.load(paths.structfit(self.nsnap, self.nsim, simname))
+
+        cols = [col for col in fits.dtype.names if col != "index"]
+        X = [fits[col] for col in cols]
+        data = add_columns(data, X, cols)
+        box = self.box
+
+        data["r200c"] = box.box2mpc(data["r200c"])
+
+        return data
+
+    def filter_data(self, data, bounds):
+        """
+        Filters data based on specified bounds for each key.
+
+        Parameters
+        ----------
+        data : structured array
+            The data to be filtered.
+        bounds : dict
+            A dictionary with keys corresponding to data columns or `dist` and
+            values as a tuple of `(xmin, xmax)`. If `xmin` or `xmax` is `None`,
+            it defaults to negative infinity and positive infinity,
+            respectively.
+
+        Returns
+        -------
+        data : structured array
+            The filtered data based on the provided bounds.
+        """
+        for key, (xmin, xmax) in bounds.items():
+            if key == "dist":
+                pos = numpy.vstack([data[p] - self.observer_location[i]
+                                    for i, p in enumerate("xyz")]).T
+                values_to_filter = numpy.linalg.norm(pos, axis=1)
+            else:
+                values_to_filter = data[key]
+
+            min_bound = xmin if xmin is not None else -numpy.inf
+            max_bound = xmax if xmax is not None else numpy.inf
+
+            data = data[(values_to_filter > min_bound)
+                        & (values_to_filter <= max_bound)]
+
+        return data
+
+    @property
+    def observer_location(self):
         r"""
-        Position components. If not Cartesian, then RA is in :math:`[0, 360)`
+        Location of the observer in units :math:`\mathrm{Mpc} / h`.
-        degrees and DEC is in :math:`[-90, 90]` degrees.
+
+        Returns
+        -------
+            obs_pos : 1-dimensional array of shape `(3,)`
+        """
+        if self._observer_location is None:
+            raise RuntimeError("`observer_location` is not set!")
+        return self._observer_location
+
+    @observer_location.setter
+    def observer_location(self, obs_pos):
+        assert isinstance(obs_pos, (list, tuple, numpy.ndarray))
+        obs_pos = numpy.asanyarray(obs_pos)
+        assert obs_pos.shape == (3,)
+        self._observer_location = obs_pos
+
+    def position(self, in_initial=False, cartesian=True,
+                 subtract_observer=False):
+        r"""
+        Return position components (Cartesian or RA/DEC).
 
         Parameters
         ----------
         in_initial : bool, optional
-            Whether to return the initial snapshot positions.
+            If True, return positions from the initial snapshot, otherwise the
+            final snapshot.
         cartesian : bool, optional
-            Whether to return the Cartesian or spherical position components.
+            If True, return Cartesian positions. Otherwise, return dist/RA/DEC
-            By default Cartesian.
+            centered at the observer.
+        subtract_observer : bool, optional
+            If True, subtract the observer's location from the returned
+            positions. This is only relevant if `cartesian` is True.
 
         Returns
         -------
-        pos : 2-dimensional array of shape `(nobjects, 3)`
+        pos : ndarray, shape `(nobjects, 3)`
+            Position components.
         """
-        if in_initial:
+        suffix = '0' if in_initial else ''
-            ps = ["x0", "y0", "z0"]
+        component_keys = [f"{comp}{suffix}" for comp in ('x', 'y', 'z')]
-        else:
-            ps = ["x", "y", "z"]
-        pos = numpy.vstack([self[p] for p in ps]).T
-        if not cartesian:
-            pos = cartesian_to_radec(pos)
-        return pos
 
-    def velocity(self):
+        pos = numpy.vstack([self[key] for key in component_keys]).T
-        r"""
-        Cartesian velocity components in :math:`\mathrm{km} / \mathrm{s}`.
 
-        Returns
+        if subtract_observer or not cartesian:
-        -------
+            pos -= self.observer_location
-        vel : 2-dimensional array of shape `(nobjects, 3)`
-        """
-        return numpy.vstack([self["v{}".format(p)] for p in ("x", "y", "z")]).T
 
-    def redshift_space_position(self, cartesian=True):
+        return cartesian_to_radec(pos) if not cartesian else pos
-        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 radial_distance(self, in_initial=False):
         r"""
-        Distance of haloes from the origin.
+        Distance of haloes from the observer in :math:`\mathrm{cMpc}`.
 
         Parameters
         ----------
@@ -198,9 +274,41 @@ class BaseCatalogue(ABC):
         -------
         radial_distance : 1-dimensional array of shape `(nobjects,)`
         """
-        pos = self.position(in_initial=in_initial, cartesian=True)
+        pos = self.position(in_initial=in_initial, cartesian=True,
+                            subtract_observer=True)
         return numpy.linalg.norm(pos, axis=1)
 
+    def velocity(self):
+        r"""
+        Return Cartesian velocity in :math:`\mathrm{km} / \mathrm{s}`.
+
+        Returns
+        -------
+        vel : 2-dimensional array of shape `(nobjects, 3)`
+        """
+        return numpy.vstack([self["v{}".format(p)] for p in ("x", "y", "z")]).T
+
+    def redshift_space_position(self, cartesian=True):
+        """
+        Calculates the position of objects in redshift space. Positions can be
+        returned  in either Cartesian coordinates (default) or spherical
+        coordinates (dist/RA/dec).
+
+        Parameters
+        ----------
+        cartesian : bool, optional
+            Returns position in Cartesian coordinates if True, else in
+            spherical coordinates.
+
+        Returns
+        -------
+        pos : 2-dimensional array of shape `(nobjects, 3)`
+            Position of objects in the desired coordinate system.
+        """
+        rsp = real2redshift(self.position(cartesian=True), self.velocity(),
+                            self.observer_location, self.box, make_copy=False)
+        return rsp if cartesian else cartesian_to_radec(rsp)
+
     def angmomentum(self):
         """
         Cartesian angular momentum components of halos in the box coordinate
@@ -214,8 +322,9 @@ class BaseCatalogue(ABC):
 
     @lru_cache(maxsize=2)
     def knn(self, in_initial):
-        """
+        r"""
-        kNN object fitted on all catalogue objects. Caches the kNN object.
+        kNN object for catalogue objects with caching. Positions are centered
+        on the observer.
 
         Parameters
         ----------
@@ -225,51 +334,49 @@ class BaseCatalogue(ABC):
         Returns
         -------
         knn : :py:class:`sklearn.neighbors.NearestNeighbors`
+            kNN object fitted with object positions.
         """
-        knn = NearestNeighbors()
+        pos = self.position(in_initial=in_initial)
-        return knn.fit(self.position(in_initial=in_initial))
+        return NearestNeighbors().fit(pos)
 
     def nearest_neighbours(self, X, radius, in_initial, knearest=False,
-                           return_mass=False, masss_key=None):
+                           return_mass=False, mass_key=None):
         r"""
-        Sorted nearest neigbours within `radius` of `X` in the initial or final
+        Return nearest neighbours within `radius` of `X` in a given snapshot.
-        snapshot. However, if `knearest` is `True` then the `radius` is assumed
-        to be the integer number of nearest neighbours to return.
 
         Parameters
         ----------
-        X : 2-dimensional array of shape `(n_queries, 3)`
+        X : 2D array, shape `(n_queries, 3)`
-            Cartesian query position components in :math:`\mathrm{cMpc}`.
+            Query positions in :math:`\mathrm{cMpc} / h`. Expected to be
+            centered on the observer.
         radius : float or int
-            Limiting neighbour distance. If `knearest` is `True` then this is
+            Limiting distance or number of neighbours, depending on `knearest`.
-            the number of nearest neighbours to return.
         in_initial : bool
-            Whether to define the kNN on the initial or final snapshot.
+            Use the initial or final snapshot for kNN.
         knearest : bool, optional
-            Whether `radius` is the number of nearest neighbours to return.
+            If True, `radius` is the number of neighbours to return.
         return_mass : bool, optional
-            Whether to return the masses of the nearest neighbours.
+            Return masses of the nearest neighbours.
-        masss_key : str, optional
+        mass_key : str, optional
-            Key of the mass column in the catalogue. Must be provided if
+            Mass column key. Required if `return_mass` is True.
-            `return_mass` is `True`.
 
         Returns
         -------
-        dist : list of 1-dimensional arrays
+        dist : list of arrays
-            List of length `n_queries` whose elements are arrays of distances
+            Distances to the nearest neighbours for each query.
-            to the nearest neighbours.
+        indxs : list of arrays
-        knns : list of 1-dimensional arrays
+            Indices of nearest neighbours for each query.
-            List of length `n_queries` whose elements are arrays of indices of
+        mass (optional): list of arrays
-            nearest neighbours in this catalogue.
+            Masses of the nearest neighbours for each query.
         """
-        if not (X.ndim == 2 and X.shape[1] == 3):
+        if X.shape != (len(X), 3):
-            raise TypeError("`X` must be an array of shape `(n_samples, 3)`.")
+            raise ValueError("`X` must be of shape `(n_samples, 3)`.")
-        if knearest:
+        if knearest and not isinstance(radius, int):
-            assert isinstance(radius, int)
+            raise ValueError("`radius` must be an integer if `knearest`.")
-        if return_mass:
+        if return_mass and not mass_key:
-            assert masss_key is not None
+            raise ValueError("`mass_key` must be provided if `return_mass`.")
-        knn = self.knn(in_initial)
 
+        knn = self.knn(in_initial)
         if knearest:
             dist, indxs = knn.kneighbors(X, radius)
         else:
@@ -278,35 +385,26 @@ class BaseCatalogue(ABC):
         if not return_mass:
             return dist, indxs
 
-        if knearest:
+        mass = [self[mass_key][indx] for indx in indxs]
-            mass = numpy.copy(dist)
-            for i in range(dist.shape[0]):
-                mass[i, :] = self[masss_key][indxs[i]]
-        else:
-            mass = deepcopy(dist)
-            for i in range(dist.size):
-                mass[i] = self[masss_key][indxs[i]]
-
         return dist, indxs, mass
 
     def angular_neighbours(self, X, ang_radius, in_rsp, rad_tolerance=None):
         r"""
-        Find nearest neighbours within `ang_radius` of query points `X`.
+        Find nearest neighbours within `ang_radius` of query points `X` in the
-        Optionally applies radial tolerance, which is expected to be in
+        final snaphot. Optionally applies radial distance tolerance, which is
-        :math:`\mathrm{cMpc}`.
+        expected to be in :math:`\mathrm{cMpc} / h`.
 
         Parameters
         ----------
         X : 2-dimensional array of shape `(n_queries, 2)` or `(n_queries, 3)`
-            Query positions. If 2-dimensional, then RA and DEC in degrees.
+            Query positions. Either RA/dec in degrees or dist/RA/dec with
-            If 3-dimensional, then radial distance in :math:`\mathrm{cMpc}`,
+            distance in :math:`\mathrm{cMpc} / h`.
-            RA and DEC in degrees.
         in_rsp : bool
-            Whether to use redshift space positions of haloes.
+            If True, use redshift space positions of haloes.
         ang_radius : float
             Angular radius in degrees.
         rad_tolerance : float, optional
-            Radial tolerance in :math:`\mathrm{cMpc}`.
+            Radial distance tolerance in :math:`\mathrm{cMpc} / h`.
 
         Returns
         -------
@@ -316,46 +414,50 @@ class BaseCatalogue(ABC):
             Indices of each neighbour in this catalogue.
         """
         assert X.ndim == 2
-        # We first get positions of haloes in this catalogue, store their
+
-        # radial distance and normalise them to unit vectors.
+        # Get positions of haloes in this catalogue
         if in_rsp:
+            # TODO what to do with subtracting the observer here?
             pos = self.redshift_space_position(cartesian=True)
         else:
-            pos = self.position(in_initial=False, cartesian=True)
+            pos = self.position(in_initial=False, cartesian=True,
+                                subtract_observer=True)
+
+        # Convert halo positions to unit vectors.
         raddist = numpy.linalg.norm(pos, axis=1)
         pos /= raddist.reshape(-1, 1)
-        # We convert RAdec query positions to unit vectors. If no radial
+
-        # distance provided add it.
+        # Convert RA/dec query positions to unit vectors. If no radial
+        # distance is provided artificially add it.
         if X.shape[1] == 2:
             X = numpy.vstack([numpy.ones_like(X[:, 0]), X[:, 0], X[:, 1]]).T
             radquery = None
         else:
             radquery = X[:, 0]
-
         X = radec_to_cartesian(X)
+
+        # Find neighbours
         knn = NearestNeighbors(metric="cosine")
         knn.fit(pos)
-        # Convert angular radius to cosine difference.
         metric_maxdist = 1 - numpy.cos(numpy.deg2rad(ang_radius))
         dist, ind = knn.radius_neighbors(X, radius=metric_maxdist,
                                          sort_results=True)
-        # And the cosine difference to angular distance.
+
+        # Convert cosine difference to angular distance
         for i in range(X.shape[0]):
             dist[i] = numpy.rad2deg(numpy.arccos(1 - dist[i]))
 
-        # Apply the radial tolerance
+        # Apply radial tolerance
-        if rad_tolerance is not None:
+        if rad_tolerance and radquery:
-            assert radquery is not None
             for i in range(X.shape[0]):
-                mask = numpy.abs(raddist[ind[i]] - radquery) < rad_tolerance
+                mask = numpy.abs(raddist[ind[i]] - radquery[i]) < rad_tolerance
-                dist[i] = dist[i][mask]
+                dist[i], ind[i] = dist[i][mask], ind[i][mask]
-                ind[i] = ind[i][mask]
+
         return dist, ind
 
-    @property
     def keys(self):
         """
-        Catalogue keys.
+        Return catalogue keys.
 
         Returns
         -------
@@ -364,11 +466,12 @@ class BaseCatalogue(ABC):
         return self.data.dtype.names
 
     def __getitem__(self, key):
+        # If key is an integer, return the corresponding row.
         if isinstance(key, (int, numpy.integer)):
             assert key >= 0
-            return self.data[key]
+        elif key not in self.keys():
-        if key not in self.keys:
             raise KeyError(f"Key '{key}' not in catalogue.")
+
         return self.data[key]
 
     def __len__(self):
@@ -382,7 +485,10 @@ class BaseCatalogue(ABC):
 
 class CSiBORGHaloCatalogue(BaseCatalogue):
     r"""
-    CSiBORG FoF halo catalogue.
+    CSiBORG FoF halo catalogue with units:
+        - Length: :math:`cMpc / h`
+        - Velocity: :math:`km / s`
+        - Mass: :math:`M_\odot / h`
 
     Parameters
     ----------
@@ -390,74 +496,52 @@ class CSiBORGHaloCatalogue(BaseCatalogue):
         IC realisation index.
     paths : py:class`csiborgtools.read.Paths`
         Paths object.
+    observer_location : array, optional
+        Observer's location in :math:`\mathrm{Mpc} / h`.
     bounds : dict
-        Parameter bounds to apply to the catalogue. The keys are the parameter
+        Parameter bounds; keys as names, values as (min, max) tuples. Use
-        names and the items are a len-2 tuple of (min, max) values. In case of
+        `dist` for radial distance, `None` for no bound.
-        no minimum or maximum, use `None`. For radial distance from the origin
-        use `dist`.
     load_fitted : bool, optional
-        Whether to load fitted quantities.
+        Load fitted quantities.
     load_initial : bool, optional
-        Whether to load initial positions.
+        Load initial positions.
     with_lagpatch : bool, optional
-        Whether to only load halos with a resolved Lagrangian patch.
+        Load halos with a resolved Lagrangian patch.
-    rawdata : bool, optional
-        Whether to return the raw data. In this case applies no cuts and
-        transformations.
     """
 
-    def __init__(self, nsim, paths, bounds={"dist": (0, 155.5 / 0.705)},
+    def __init__(self, nsim, paths, observer_location=[338.85, 338.85, 338.85],
-                 load_fitted=True, load_initial=True, with_lagpatch=True,
+                 bounds={"dist": (0, 155.5)},
-                 rawdata=False):
+                 load_fitted=True, load_initial=True, with_lagpatch=False):
         self.nsim = nsim
         self.paths = paths
+        self.observer_location = observer_location
         reader = CSiBORGReader(paths)
-        self._data = reader.read_fof_halos(self.nsim)
+        data = reader.read_fof_halos(self.nsim)
+        box = self.box
 
-        if load_fitted:
+        # We want coordinates to be [0, 677.7] in Mpc / h
-            fits = numpy.load(paths.structfit(self.nsnap, nsim, "csiborg"))
-            cols = [col for col in fits.dtype.names if col != "index"]
-            X = [fits[col] for col in cols]
-            self._data = add_columns(self._data, X, cols)
-
-        if load_initial:
-            fits = numpy.load(paths.initmatch(nsim, "csiborg", "fit"))
-            X, cols = [], []
-            for col in fits.dtype.names:
-                if col == "index":
-                    continue
-                if col in ['x', 'y', 'z']:
-                    cols.append(col + "0")
-                else:
-                    cols.append(col)
-                X.append(fits[col])
-
-            self._data = add_columns(self._data, X, cols)
-
-        if rawdata:
         for p in ('x', 'y', 'z'):
-                self._data[p] = self.box.mpc2box(self._data[p]) + 0.5
+            data[p] = data[p] * box.h + box.box2mpc(1) / 2
-        else:
+        # Similarly mass in units of Msun / h
-            if with_lagpatch:
+        data["fof_totpartmass"] *= box.h
-                self._data = self._data[numpy.isfinite(self["lagpatch_size"])]
+        data["fof_m200c"] *= box.h
-            # Flip positions and convert from code units to cMpc. Convert M too
+        # Because of a RAMSES bug, we must flip the x and z coordinates
-            flip_cols(self._data, "x", "z")
+        flip_cols(data, 'x', 'z')
-            if load_fitted:
-                flip_cols(self._data, "vx", "vz")
-                names = ["totpartmass", "rho0", "r200c",
-                         "r500c", "m200c", "m500c", "r200m", "m200m",
-                         "r500m", "m500m", "vx", "vy", "vz"]
-                self._data = self.box.convert_from_box(self._data, names)
 
         if load_initial:
-                flip_cols(self._data, "x0", "z0")
+            data = self.load_initial(data, paths, "csiborg")
-                for p in ("x0", "y0", "z0"):
+            flip_cols(data, "x0", "z0")
-                    self._data[p] -= 0.5
+        if load_fitted:
-                names = ["x0", "y0", "z0", "lagpatch_size"]
+            data = self.load_fitted(data, paths, "csiborg")
-                self._data = self.box.convert_from_box(self._data, names)
+            flip_cols(data, "vx", "vz")
+
+        if load_initial and with_lagpatch:
+            data = data[numpy.isfinite(data["lagpatch_size"])]
 
         if bounds is not None:
-                self.apply_bounds(bounds)
+            data = self.filter_data(data, bounds)
+
+        self._data = data
 
     @property
     def nsnap(self):
@@ -481,8 +565,11 @@ class CSiBORGHaloCatalogue(BaseCatalogue):
 
 
 class QuijoteHaloCatalogue(BaseCatalogue):
-    """
+    r"""
-    Quijote FoF halo catalogue.
+    Quijote FoF halo catalogue with units:
+        - Length: :math:`cMpc / h`
+        - Velocity: :math:`km / s`
+        - Mass: :math:`M_\odot / h`
 
     Parameters
     ----------
@@ -492,34 +579,30 @@ class QuijoteHaloCatalogue(BaseCatalogue):
         Paths object.
     nsnap : int
         Snapshot index.
-    origin : len-3 tuple, optional
+    observer_location : array, optional
-        Where to place the origin of the box. In units of :math:`cMpc / h`.
+        Observer's location in :math:`\mathrm{Mpc} / h`.
     bounds : dict
-        Parameter bounds to apply to the catalogue. The keys are the parameter
+        Parameter bounds; keys as parameter names, values as (min, max)
-        names and the items are a len-2 tuple of (min, max) values. In case of
+        tuples. Use `dist` for radial distance, `None` for no bound.
-        no minimum or maximum, use `None`. For radial distance from the origin
+    load_fitted : bool, optional
-        use `dist`.
+        Load fitted quantities from `fit_halos.py`.
     load_initial : bool, optional
-        Whether to load initial positions.
+        Load initial positions from `fit_init.py`.
     with_lagpatch : bool, optional
-        Whether to only load halos with a resolved Lagrangian patch.
+        Load halos with a resolved Lagrangian patch.
-    rawdata : bool, optional
-        Whether to return the raw data. In this case applies no cuts and
-        transformations.
-    **kwargs : dict
-        Keyword arguments for backward compatibility.
     """
     _nsnap = None
     _origin = None
 
-    def __init__(self, nsim, paths, nsnap, origin=[0., 0., 0.],
+    def __init__(self, nsim, paths, nsnap,
-                 bounds=None, load_initial=True, with_lagpatch=True,
+                 observer_location=[500., 500., 500.],
-                 rawdata=False, **kwargs):
+                 bounds=None, load_fitted=True, load_initial=True,
+                 with_lagpatch=False):
+        self.nsim = nsim
         self.paths = paths
         self.nsnap = nsnap
-        self.origin = origin
+        self.observer_location = observer_location
         self._box = QuijoteBox(nsnap, nsim, paths)
-        self._boxwidth = self.box.boxsize
 
         fpath = self.paths.fof_cat(nsim, "quijote")
         fof = FoF_catalog(fpath, self.nsnap, long_ids=False, swap=False,
@@ -532,44 +615,29 @@ class QuijoteHaloCatalogue(BaseCatalogue):
                 ("index", numpy.int32)]
         data = cols_to_structured(fof.GroupLen.size, cols)
 
-        pos = self.box.mpc2box(fof.GroupPos / 1e3)
+        pos = fof.GroupPos / 1e3
         vel = fof.GroupVel * (1 + self.redshift)
         for i, p in enumerate(["x", "y", "z"]):
-            data[p] = pos[:, i] - self.origin[i]
+            data[p] = pos[:, i]
             data["v" + p] = vel[:, i]
-        data["group_mass"] = self.box.solarmass2box(fof.GroupMass * 1e10)
+        data["group_mass"] = fof.GroupMass * 1e10
         data["npart"] = fof.GroupLen
         # We want to start indexing from 1. Index 0 is reserved for
         # particles unassigned to any FoF group.
         data["index"] = 1 + numpy.arange(data.size, dtype=numpy.int32)
 
         if load_initial:
-            fits = numpy.load(paths.initmatch(nsim, "quijote", "fit"))
+            data = self.load_initial(data, paths, "quijote")
-            X, cols = [], []
+        if load_fitted:
-            for col in fits.dtype.names:
+            assert nsnap == 4
-                if col == "index":
-                    continue
-                if col in ['x', 'y', 'z']:
-                    cols.append(col + "0")
-                else:
-                    cols.append(col)
-                X.append(fits[col])
-            data = add_columns(data, X, cols)
 
-        self._data = data
+        if load_initial and with_lagpatch:
-        if not rawdata:
+            data = data[numpy.isfinite(data["lagpatch_size"])]
-            if with_lagpatch:
-                mask = numpy.isfinite(self._data["lagpatch_size"])
-                self._data = self._data[mask]
-
-                names = ["x", "y", "z", "group_mass"]
-                self._data = self.box.convert_from_box(self._data, names)
-                if load_initial:
-                    names = ["x0", "y0", "z0", "lagpatch_size"]
-                    self._data = self.box.convert_from_box(self._data, names)
 
         if bounds is not None:
-                self.apply_bounds(bounds)
+            data = self.filter_data(data, bounds)
+
+        self._data = data
 
     @property
     def nsnap(self):
@@ -596,40 +664,12 @@ class QuijoteHaloCatalogue(BaseCatalogue):
         -------
         redshift : float
         """
-        z_dict = {4: 0.0, 3: 0.5, 2: 1.0, 1: 2.0, 0: 3.0}
+        return {4: 0.0, 3: 0.5, 2: 1.0, 1: 2.0, 0: 3.0}[self.nsnap]
-        return z_dict[self.nsnap]
 
     @property
     def box(self):
-        """
-        Quijote box object.
-
-        Returns
-        -------
-        box : instance of :py:class:`csiborgtools.units.BaseBox`
-        """
         return self._box
 
-    @property
-    def origin(self):
-        """
-        Origin of the box with respect to the initial box units.
-
-        Returns
-        -------
-        origin : len-3 tuple
-        """
-        if self._origin is None:
-            raise ValueError("`origin` is not set.")
-        return self._origin
-
-    @origin.setter
-    def origin(self, origin):
-        if isinstance(origin, (list, tuple)):
-            origin = numpy.asanyarray(origin)
-        assert origin.ndim == 1 and origin.size == 3
-        self._origin = origin
-
     def pick_fiducial_observer(self, n, rmax):
         r"""
         Return a copy of itself, storing only halos within `rmax` of the new
@@ -640,22 +680,15 @@ class QuijoteHaloCatalogue(BaseCatalogue):
         n : int
             Fiducial observer index.
         rmax : float
-            Maximum distance from the fiducial observer in :math:`cMpc`.
+            Max. distance from the fiducial obs. in :math:`\mathrm{cMpc} / h`.
 
         Returns
         -------
         cat : instance of csiborgtools.read.QuijoteHaloCatalogue
         """
-        new_origin = fiducial_observers(self.box.boxsize, rmax)[n]
-        # We make a copy of the catalogue to avoid modifying the original.
-        # Then, we shift coordinates back to the original box frame and then to
-        # the new origin.
         cat = deepcopy(self)
-        for i, p in enumerate(('x', 'y', 'z')):
+        cat.observer_location = fiducial_observers(self.box.boxsize, rmax)[n]
-            cat._data[p] += self.origin[i]
+        cat._data = cat.filter_data(cat._data, {"dist": (0, rmax)})
-            cat._data[p] -= new_origin[i]
-
-        cat.apply_bounds({"dist": (0, rmax)})
         return cat
 
 
@@ -666,26 +699,20 @@ class QuijoteHaloCatalogue(BaseCatalogue):
 
 def fiducial_observers(boxwidth, radius):
     """
-    Positions of fiducial observers in a box, such that that the box is
+    Compute observer positions in a box, subdivided into spherical regions.
-    subdivided among them into spherical regions.
 
     Parameters
     ----------
     boxwidth : float
-        Box width.
+        Width of the box.
     radius : float
         Radius of the spherical regions.
 
     Returns
     -------
-    origins : list of len-3 lists
+    origins : list of lists
-        Positions of the observers.
+        Positions of the observers, with each position as a len-3 list.
     """
-    nobs = floor(boxwidth / (2 * radius))  # Number of observers per dimension
+    nobs = floor(boxwidth / (2 * radius))
-
+    return [[val * radius for val in position]
-    origins = list(product([1, 3, 5], repeat=nobs))
+            for position in product([1, 3, 5], repeat=nobs)]
-    for i in range(len(origins)):
-        origins[i] = list(origins[i])
-        for j in range(nobs):
-            origins[i][j] *= radius
-    return origins

csiborgtools/read/readsim.py

@@ -531,6 +531,8 @@ class MmainReader:
         the position of the parent, the summed mass and the fraction of mass in
         substructure. Corresponds to the PHEW Halo finder.
 
+        NOTE: this code is no longer used and the units may be inconsistent.
+
         Parameters
         ----------
         nsim : int
@@ -642,8 +644,8 @@ class QuijoteReader:
 
         if verbose:
             print(f"{datetime.now()}: reading particle velocities.")
-        # NOTE convert to box units.
+        # Unlike the positions, we keep velocities in km/s
-        vel = readgadget.read_block(snapshot, "VEL ", ptype)  # km/s
+        vel = readgadget.read_block(snapshot, "VEL ", ptype)
         vel *= (1 + info["redshift"])
 
         for i, v in enumerate(['vx', 'vy', 'vz']):
@@ -657,9 +659,9 @@ class QuijoteReader:
         if verbose:
             print(f"{datetime.now()}: reading particle masses.")
         if return_structured:
-            out["M"] = info["PartMass"] / info["TotMass"]
+            out["M"] = info["PartMass"]
         else:
-            out[:, 6] = info["PartMass"] / info["TotMass"]
+            out[:, 6] = info["PartMass"]
 
         return out, pids
 

csiborgtools/read/utils.py

@@ -18,7 +18,6 @@ Various coordinate transformations.
 from os.path import isfile
 
 import numpy
-from astropy import units
 from h5py import File
 
 ###############################################################################
@@ -81,7 +80,7 @@ def radec_to_cartesian(X, isdeg=True):
     return numpy.vstack([x, y, z]).T
 
 
-def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True,
+def real2redshift(pos, vel, observer_location, box, periodic_wrap=True,
                   make_copy=True):
     r"""
     Convert real-space position to redshift space position.
@@ -89,18 +88,13 @@ def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True,
     Parameters
     ----------
     pos : 2-dimensional array `(nsamples, 3)`
-        Real-space Cartesian position components.
+        Real-space Cartesian components in :math:`\mathrm{cMpc} / h`.
     vel : 2-dimensional array `(nsamples, 3)`
-        Cartesian velocity components.
+        Cartesian velocity in :math:`\mathrm{km} \mathrm{s}^{-1}`.
-    origin : 1-dimensional array `(3,)`
+    observer_location: 1-dimensional array `(3,)`
-        Origin of the coordinate system in the `pos` reference frame.
+        Observer location in :math:`\mathrm{cMpc} / h`.
     box : py:class:`csiborg.read.CSiBORGBox`
         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.
     periodic_wrap : bool, optional
         Whether to wrap around the box, particles may be outside the default
         bounds once RSD is applied.
@@ -110,55 +104,30 @@ def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True,
     Returns
     -------
     pos : 2-dimensional array `(nsamples, 3)`
-        Redshift-space Cartesian position components, with an observer assumed
+        Redshift-space Cartesian position in :math:`\mathrm{cMpc} / h`.
-        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]
 
+    # Place the observer at the origin
+    pos -= observer_location
+    # Dot product of position vector and velocity
+    vr_dot = numpy.sum(pos * vel, axis=1)
+    # Compute the norm squared of the displacement
     norm2 = numpy.sum(pos**2, axis=1)
-    dot = numpy.einsum("ij,ij->i", pos, vel)
+    pos *= (1 + box._aexp / box.H0 * vr_dot / norm2).reshape(-1, 1)
-    pos *= (1 + a / H0 * dot / norm2).reshape(-1, 1)
+    # Place the observer back at the original location
-
+    pos += observer_location
-    for i in range(3):
-        pos[:, i] += origin[i]
 
     if periodic_wrap:
-        boxsize = 1. if in_box_units else box.box2mpc(1.)
+        boxsize = box.box2mpc(1.)
-        # Wrap around the box: x > 1 -> x - 1, x < 0 -> x + 1
+        # Wrap around the box.
-        pos[pos > boxsize] -= boxsize
+        pos = numpy.where(pos > boxsize, pos - boxsize, pos)
-        pos[pos < 0] += boxsize
+        pos = numpy.where(pos < 0, pos + boxsize, pos)
+
     return pos
 
 
-def M200_to_R200(M200, cosmo):
-    r"""
-    Convert :math:M_{200} to :math:`R_{200}`.
-
-    Parameters
-    ----------
-    M200 : float
-        :math:`M_{200}` in :math:`M_{\odot}`.
-    cosmo : astropy cosmology object
-        Cosmology.
-
-    Returns
-    -------
-    R200 : float
-        :math:`R_{200}` in :math:`\mathrm{Mpc}`.
-    """
-    Msun = 1.98847e30
-    M200 = 1e14 * Msun * units.kg
-    rhoc = cosmo.critical_density0
-    R200 = (M200 / (4 * numpy.pi / 3 * 200 * rhoc))**(1. / 3)
-    return R200.to(units.Mpc).value
-
-
 ###############################################################################
 #                          Array manipulation                                 #
 ###############################################################################
@@ -173,215 +142,121 @@ def cols_to_structured(N, cols):
     N : int
         Structured array size.
     cols: list of tuples
-        Column names and dtypes. Each tuple must written as `(name, dtype)`.
+        Column names and dtypes. Each tuple must be written as `(name, dtype)`.
 
     Returns
     -------
     out : structured array
-        Initialised structured array.
+        Initialized structured array.
     """
-    if not isinstance(cols, list) and all(isinstance(c, tuple) for c in cols):
+    if not (isinstance(cols, list)
-        raise TypeError("`cols` must be a list of tuples.")
+            and all(isinstance(c, tuple) and len(c) == 2 for c in cols)):
+        raise TypeError("`cols` must be a list of (name, dtype) tuples.")
+
+    names, formats = zip(*cols)
+    dtype = {"names": names, "formats": formats}
 
-    dtype = {"names": [col[0] for col in cols],
-             "formats": [col[1] for col in cols]}
     return numpy.full(N, numpy.nan, dtype=dtype)
 
 
 def add_columns(arr, X, cols):
     """
-    Add new columns to a record array `arr`. Creates a new array.
+    Add new columns `X` to a record array `arr`. Creates a new array.
 
     Parameters
     ----------
-    arr : record array
+    arr : structured array
-        Record array to add columns to.
+        Structured array to add columns to.
-    X : (list of) 1-dimensional array(s) or 2-dimensional array
+    X : (list of) 1-dimensional array(s)
         Columns to be added.
     cols : str or list of str
         Column names to be added.
 
     Returns
     -------
-    out : record array
+    out : structured array
     """
-    # Make sure cols is a list of str and X a 2D array
     cols = [cols] if isinstance(cols, str) else cols
+
+    # Convert X to a list of 1D arrays for consistency
     if isinstance(X, numpy.ndarray) and X.ndim == 1:
-        X = X.reshape(-1, 1)
+        X = [X]
-    if isinstance(X, list) and all(x.ndim == 1 for x in X):
+    elif isinstance(X, numpy.ndarray):
-        X = numpy.vstack([X]).T
+        raise ValueError("`X` should be a 1D array or a list of 1D arrays.")
-    if len(cols) != X.shape[1]:
-        raise ValueError("Number of columns of `X` does not match `cols`.")
-    if arr.size != X.shape[0]:
-        raise ValueError("Number of rows of `X` does not match size of `arr`.")
 
-    # Get the new data types
+    if len(X) != len(cols):
-    dtype = arr.dtype.descr
+        raise ValueError("Mismatch between `X` and `cols` lengths.")
-    for i, col in enumerate(cols):
-        dtype.append((col, X[i, :].dtype.descr[0][1]))
 
-    # Fill in the old array
+    if not all(isinstance(x, numpy.ndarray) and x.ndim == 1 for x in X):
-    out = numpy.full(arr.size, numpy.nan, dtype=dtype)
+        raise ValueError("All elements of `X` should be 1D arrays.")
+
+    if not all(x.size == arr.size for x in X):
+        raise ValueError("All arrays in `X` must have the same size as `arr`.")
+
+    # Define new dtype
+    dtype = list(arr.dtype.descr) + [(col, x.dtype) for col, x in zip(cols, X)]
+
+    # Create a new array and fill in values
+    out = numpy.empty(arr.size, dtype=dtype)
     for col in arr.dtype.names:
         out[col] = arr[col]
-    for i, col in enumerate(cols):
+    for col, x in zip(cols, X):
-        out[col] = X[:, i]
+        out[col] = x
 
     return out
 
 
 def rm_columns(arr, cols):
     """
-    Remove columns `cols` from a record array `arr`. Creates a new array.
+    Remove columns `cols` from a structured array `arr`. Allocates a new array.
 
     Parameters
     ----------
-    arr : record array
+    arr : structured array
-        Record array to remove columns from.
+        Structured array to remove columns from.
     cols : str or list of str
         Column names to be removed.
 
     Returns
     -------
-    out : record array
+    out : structured array
     """
-    # Check columns we wish to delete are in the array
+    # Ensure cols is a list
     cols = [cols] if isinstance(cols, str) else cols
-    for col in cols:
-        if col not in arr.dtype.names:
-            raise ValueError("Column `{}` not in `arr`.".format(col))
 
-    # Get a new dtype without the cols to be deleted
+    # Check columns we wish to delete are in the array
-    new_dtype = []
+    missing_cols = [col for col in cols if col not in arr.dtype.names]
-    for dtype, name in zip(arr.dtype.descr, arr.dtype.names, strict=True):
+    if missing_cols:
-        if name not in cols:
+        raise ValueError(f"Columns `{missing_cols}` not in `arr`.")
-            new_dtype.append(dtype)
 
-    # Allocate a new array and fill it in.
+    # Define new dtype without the cols to be deleted
-    out = numpy.full(arr.size, numpy.nan, new_dtype)
+    new_dtype = [(n, dt) for n, dt in arr.dtype.descr if n not in cols]
+
+    # Allocate a new array and fill in values
+    out = numpy.empty(arr.size, dtype=new_dtype)
     for name in out.dtype.names:
         out[name] = arr[name]
 
     return out
 
 
-def list_to_ndarray(arrs, cols):
-    """
-    Convert a list of structured arrays of CSiBORG simulation catalogues to
-    an 3-dimensional array.
-
-    Parameters
-    ----------
-    arrs : list of structured arrays
-        List of CSiBORG catalogues.
-    cols : str or list of str
-        Columns to be extracted from the CSiBORG catalogues.
-
-    Returns
-    -------
-    out : 3-dimensional array
-        Catalogue array of shape `(n_realisations, n_samples, n_cols)`, where
-        `n_samples` is the maximum number of samples over the CSiBORG
-        catalogues.
-    """
-    if not isinstance(arrs, list):
-        raise TypeError("`arrs` must be a list of structured arrays.")
-    cols = [cols] if isinstance(cols, str) else cols
-
-    Narr = len(arrs)
-    Nobj_max = max([arr.size for arr in arrs])
-    Ncol = len(cols)
-    # Preallocate the array and fill it
-    out = numpy.full((Narr, Nobj_max, Ncol), numpy.nan)
-    for i in range(Narr):
-        Nobj = arrs[i].size
-        for j in range(Ncol):
-            out[i, :Nobj, j] = arrs[i][cols[j]]
-    return out
-
-
-def array_to_structured(arr, cols):
-    """
-    Create a structured array from a 2-dimensional array.
-
-    Parameters
-    ----------
-    arr : 2-dimensional array
-        Original array of shape `(n_samples, n_cols)`.
-    cols : list of str
-        Columns of the structured array
-
-    Returns
-    -------
-    out : structured array
-        Output structured array.
-    """
-    cols = [cols] if isinstance(cols, str) else cols
-    if arr.ndim != 2 and arr.shape[1] != len(cols):
-        raise TypeError("`arr` must be a 2D array `(n_samples, n_cols)`.")
-
-    dtype = {"names": cols, "formats": [arr.dtype] * len(cols)}
-    out = numpy.full(arr.shape[0], numpy.nan, dtype=dtype)
-    for i, col in enumerate(cols):
-        out[col] = arr[:, i]
-
-    return out
-
-
 def flip_cols(arr, col1, col2):
     """
-    Flip values in columns `col1` and `col2`. `arr` is passed by reference and
+    Flip values in columns `col1` and `col2`. `arr` is modified in place.
-    is not explicitly returned back.
 
     Parameters
     ----------
     arr : structured array
-        Array whose columns are to be converted.
+        Array whose columns are to be flipped.
     col1 : str
         First column name.
     col2 : str
         Second column name.
-
-    Returns
-    -------
-    None
     """
-    dum = numpy.copy(arr[col1])
+    if col1 not in arr.dtype.names or col2 not in arr.dtype.names:
-    arr[col1] = arr[col2]
+        raise ValueError(f"Both `{col1}` and `{col2}` must exist in `arr`.")
-    arr[col2] = dum
 
-
+    arr[col1], arr[col2] = numpy.copy(arr[col2]), numpy.copy(arr[col1])
-def extract_from_structured(arr, cols):
-    """
-    Extract columns `cols` from a structured array. The  array dtype is set
-    to be that of the first column in `cols`.
-
-    Parameters
-    ----------
-    arr : structured array
-        Array from which to extract columns.
-    cols : list of str or str
-        Column to extract.
-
-    Returns
-    -------
-    out : 2- or 1-dimensional array
-        Array with shape `(n_particles, len(cols))`. If `len(cols)` is 1
-        flattens the array.
-    """
-    cols = [cols] if isinstance(cols, str) else cols
-    for col in cols:
-        if col not in arr.dtype.names:
-            raise ValueError(f"Invalid column `{col}`!")
-    # Preallocate an array and populate it
-    out = numpy.zeros((arr.size, len(cols)), dtype=arr[cols[0]].dtype)
-    for i, col in enumerate(cols):
-        out[:, i] = arr[col]
-    # Optionally flatten
-    if len(cols) == 1:
-        return out.reshape(-1, )
-    return out
 
 
 ###############################################################################

scripts/fit_halos.py

@@ -99,12 +99,12 @@ def _main(nsim, simname, verbose):
     if simname == "csiborg":
         box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
         cat = csiborgtools.read.CSiBORGHaloCatalogue(
-            nsim, paths, with_lagpatch=False, load_initial=False, rawdata=True,
+            nsim, paths, bounds=None, load_fitted=False, load_initial=False)
-            load_fitted=False)
     else:
         box = csiborgtools.read.QuijoteBox(nsnap, nsim, paths)
         cat = csiborgtools.read.QuijoteHaloCatalogue(
-            nsim, paths, nsnap, load_initial=False, rawdata=True)
+            nsim, paths, nsnap, bounds=None, load_fitted=False,
+            load_initial=False)
 
     # Particle archive
     f = csiborgtools.read.read_h5(paths.particles(nsim, simname))
@@ -116,6 +116,7 @@ def _main(nsim, simname, verbose):
     for i in trange(len(cat)) if verbose else range(len(cat)):
         hid = cat["index"][i]
         out["index"][i] = hid
+        # print("i = ", i)
         part = csiborgtools.read.load_halo_particles(hid, particles, halo_map,
                                                      hid2map)
         # Skip if no particles.

scripts/fit_hmf.py

@@ -59,12 +59,12 @@ def get_counts(nsim, bins, paths, parser_args):
 
     if simname == "csiborg":
         cat = csiborgtools.read.CSiBORGHaloCatalogue(
-            nsim, paths, bounds=bounds, with_lagpatch=False,
+            nsim, paths, bounds=bounds, load_initial=False)
-            load_initial=False)
         logmass = numpy.log10(cat["totpartmass"])
         counts = csiborgtools.fits.number_counts(logmass, bins)
     elif simname == "quijote":
-        cat0 = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4)
+        cat0 = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4,
+                                                      load_initial=False)
         nmax = int(cat0.box.boxsize // (2 * parser_args.Rmax))**3
         counts = numpy.full((nmax, len(bins) - 1), numpy.nan,
                             dtype=numpy.float32)
@@ -74,7 +74,8 @@ def get_counts(nsim, bins, paths, parser_args):
             logmass = numpy.log10(cat["group_mass"])
             counts[nobs, :] = csiborgtools.fits.number_counts(logmass, bins)
     elif simname == "quijote_full":
-        cat = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4)
+        cat = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4,
+                                                     load_initial=False)
         logmass = numpy.log10(cat["group_mass"])
         counts = csiborgtools.fits.number_counts(logmass, bins)
     else:

scripts/fit_init.py

@@ -68,7 +68,8 @@ def _main(nsim, simname, verbose):
         cat = csiborgtools.read.CSiBORGHaloCatalogue(
             nsim, paths, rawdata=True, load_fitted=False, load_initial=False)
     else:
-        cat = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4)
+        cat = csiborgtools.read.QuijoteHaloCatalogue(
+            nsim, paths, nsnap=4, load_fitted=False, load_initial=False)
     hid2map = {hid: i for i, hid in enumerate(halo_map[:, 0])}
 
     # Initialise the overlapper.

scripts/pre_dumppart.py

@@ -11,11 +11,16 @@
 # You should have received a copy of the GNU General Public License along
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""
+r"""
 Script to load in the simulation particles, sort them by their FoF halo ID and
 dump into a HDF5 file. Stores the first and last index of each halo in the
 particle array. This can be used for fast slicing of the array to acces
 particles of a single clump.
+
+Ensures the following units:
+    - Positions in box units.
+    - Velocities in :math:`\mathrm{km} / \mathrm{s}`.
+    - Masses in :math:`M_\odot / h`.
 """
 from argparse import ArgumentParser
 from datetime import datetime
@@ -118,6 +123,14 @@ def main(nsim, simname, verbose):
         pars_extract = None
     parts, pids = partreader.read_particle(
         nsnap, nsim, pars_extract, return_structured=False, verbose=verbose)
+
+    # In case of CSiBORG, we need to convert the mass and velocities from
+    # box units.
+    if simname == "csiborg":
+        box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
+        parts[:, [3, 4, 5]] = box.box2vel(parts[:, [3, 4, 5]])
+        parts[:, 6] = box.box2solarmass(parts[:, 6])
+
     # Now we in two steps save the particles and particle IDs.
     if verbose:
         print(f"{datetime.now()}: dumping particles from {nsim}.", flush=True)

scripts/pre_sortinit.py

@@ -12,9 +12,13 @@
 # You should have received a copy of the GNU General Public License along
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""
+r"""
 Script to sort the initial snapshot particles according to their final
 snapshot ordering, which is sorted by the halo IDs.
+
+Ensures the following units:
+    - Positions in box units.
+    - Masses in :math:`M_\odot / h`.
 """
 from argparse import ArgumentParser
 from datetime import datetime
@@ -75,6 +79,13 @@ def _main(nsim, simname, verbose):
         nsnap = -1
     part0, pid0 = partreader.read_particle(
         nsnap, nsim, pars_extract, return_structured=False, verbose=verbose)
+
+    # In CSiBORG we need to convert particle masses from box units.
+    if simname == "csiborg":
+        box = csiborgtools.read.CSiBORGBox(
+            max(paths.get_snapshots(nsim, simname)), nsim, paths)
+        part0[:, 3] = box.box2solarmass(part0[:, 3])
+
     # Quijote's initial snapshot information also contains velocities but we
     # don't need those.
     if simname == "quijote":

scripts/utils.py

@@ -89,9 +89,13 @@ def read_single_catalogue(args, config, nsim, run, rmax, paths, nobs=None):
         raise KeyError(f"No configuration for run {run}.")
     # We first read the full catalogue without applying any bounds.
     if args.simname == "csiborg":
-        cat = csiborgtools.read.CSiBORGHaloCatalogue(nsim, paths)
+        cat = csiborgtools.read.CSiBORGHaloCatalogue(
+            nsim, paths, load_fitted=True, load_inital=True,
+            with_lagpatch=False)
     else:
-        cat = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4)
+        cat = csiborgtools.read.QuijoteHaloCatalogue(
+            nsim, paths, nsnap=4, load_fitted=True, load_initial=True,
+            with_lagpatch=False)
         if nobs is not None:
             # We may optionally already here pick a fiducial observer.
             cat = cat.pick_fiducial_observer(nobs, args.Rmax)

scripts_plots/plot_data.py

@@ -49,8 +49,10 @@ def open_csiborg(nsim):
     cat : csiborgtools.read.CSiBORGHaloCatalogue
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    bounds = {"totpartmass": (None, None), "dist": (0, 155/0.705)}
+    bounds = {"totpartmass": (None, None), "dist": (0, 155)}
-    return csiborgtools.read.CSiBORGHaloCatalogue(nsim, paths, bounds=bounds)
+    return csiborgtools.read.CSiBORGHaloCatalogue(
+        nsim, paths, bounds=bounds, load_fitted=True, load_initial=True,
+        with_lagpatch=False)
 
 
 def open_quijote(nsim, nobs=None):
@@ -69,9 +71,11 @@ def open_quijote(nsim, nobs=None):
     cat : csiborgtools.read.QuijoteHaloCatalogue
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    cat = csiborgtools.read.QuijoteHaloCatalogue(nsim, paths, nsnap=4)
+    cat = csiborgtools.read.QuijoteHaloCatalogue(
+        nsim, paths, nsnap=4, load_fitted=True, load_initial=True,
+        with_lagpatch=False)
     if nobs is not None:
-        cat = cat.pick_fiducial_observer(nobs, rmax=155.5 / 0.705)
+        cat = cat.pick_fiducial_observer(nobs, rmax=155.5)
     return cat
 
 
@@ -101,7 +105,7 @@ def plot_mass_vs_ncells(nsim, pdf=False):
         plt.yscale("log")
         for n in [1, 10, 100]:
             plt.axvline(n * 512 * mpart, c="black", ls="--", zorder=0, lw=0.8)
-        plt.xlabel(r"$M_{\rm tot} / M_\odot$")
+        plt.xlabel(r"$M_{\rm tot} ~ [M_\odot$ / h]")
         plt.ylabel(r"$N_{\rm cells}$")
 
         for ext in ["png"] if pdf is False else ["png", "pdf"]:
@@ -198,7 +202,7 @@ def plot_hmf(pdf=False):
         ax[1].axhline(1, color="k", ls=plt.rcParams["lines.linestyle"],
                       lw=0.5 * plt.rcParams["lines.linewidth"], zorder=0)
         ax[0].set_ylabel(r"$\frac{\mathrm{d} n}{\mathrm{d}\log M_{\rm h}}~\mathrm{dex}^{-1}$")  # noqa
-        ax[1].set_xlabel(r"$M_{\rm h}$ [$M_\odot$]")
+        ax[1].set_xlabel(r"$M_{\rm h}~[M_\odot / h]$")
         ax[1].set_ylabel(r"$\mathrm{CSiBORG} / \mathrm{Quijote}$")
 
         ax[0].set_xscale("log")
@@ -268,7 +272,7 @@ def plot_hmf_quijote_full(pdf=False):
                       lw=0.5 * plt.rcParams["lines.linewidth"], zorder=0)
         ax[0].set_ylabel(r"$\frac{\mathrm{d}^2 n}{\mathrm{d}\log M_{\rm h} \mathrm{d} V}~[\mathrm{dex}^{-1} (\mathrm{Mpc / h})^{-3}]$",  # noqa
                          fontsize="small")
-        ax[1].set_xlabel(r"$M_{\rm h}$ [$M_\odot$]")
+        ax[1].set_xlabel(r"$M_{\rm h}~[$M_\odot / h]$", fontsize="small")
         ax[1].set_ylabel(r"$\mathrm{HMF} / \langle \mathrm{HMF} \rangle$",
                          fontsize="small")
 

scripts_plots/plot_knn.py

@@ -99,7 +99,7 @@ def plot_knn(runname):
 #                              color=cols[k % len(cols)], zorder=0)
 
         plt.legend()
-        plt.xlabel(r"$r~[\mathrm{Mpc}]$")
+        plt.xlabel(r"$r~[\mathrm{Mpc} / h]$")
         plt.ylabel(r"$P(k | V = 4 \pi r^3 / 3)$")
 
         for ext in ["png"]:

scripts_plots/plot_match.py

@@ -54,7 +54,7 @@ def open_cat(nsim):
     cat : csiborgtools.read.CSiBORGHaloCatalogue
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    bounds = {"totpartmass": (1e12, None)}
+    bounds = {"dist": (0, 155), "totpartmass": (1e12, None)}
     return csiborgtools.read.CSiBORGHaloCatalogue(nsim, paths, bounds=bounds)
 
 
@@ -86,7 +86,7 @@ def plot_mass_vs_pairoverlap(nsim0, nsimx):
         plt.hexbin(x, y, mincnt=1, bins="log",
                    gridsize=50)
         plt.colorbar(label="Counts in bins")
-        plt.xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        plt.xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         plt.ylabel("Pair overlap")
         plt.ylim(0., 1.)
 
@@ -130,7 +130,7 @@ def plot_mass_vs_maxpairoverlap(nsim0, nsimx):
         plt.hexbin(x, y, mincnt=1, bins="log",
                    gridsize=50)
         plt.colorbar(label="Counts in bins")
-        plt.xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        plt.xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         plt.ylabel("Maximum pair overlap")
         plt.ylim(0., 1.)
 
@@ -214,9 +214,9 @@ def plot_mass_vsmedmaxoverlap(nsim0):
                             numpy.nanstd(max_overlap, axis=1), gridsize=30,
                             C=x, reduce_C_function=numpy.nanmean)
 
-        axs[0].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[0].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[0].set_ylabel(r"Mean max. pair overlap")
-        axs[1].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[1].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[1].set_ylabel(r"Uncertainty of max. pair overlap")
         axs[2].set_xlabel(r"Mean max. pair overlap")
         axs[2].set_ylabel(r"Uncertainty of max. pair overlap")
@@ -287,14 +287,15 @@ def plot_summed_overlap_vs_mass(nsim0):
         axs[2].plot(t, t, color="red", linestyle="--")
         axs[0].set_ylim(0.)
         axs[1].set_ylim(0.)
-        axs[0].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[0].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[0].set_ylabel("Mean summed overlap")
-        axs[1].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[1].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[1].set_ylabel("Uncertainty of summed overlap")
         axs[2].set_xlabel(r"$1 - $ mean summed overlap")
         axs[2].set_ylabel("Mean prob. of no match")
 
-        label = ["Bin counts", "Bin counts", r"$\log M_{\rm tot} / M_\odot$"]
+        label = ["Bin counts", "Bin counts",
+                 r"$\log M_{\rm tot} ~ [M_\odot / h]$"]
         ims = [im1, im2, im3]
         for i in range(3):
             axins = inset_axes(axs[i], width="100%", height="5%",
@@ -338,7 +339,7 @@ def plot_mass_vs_separation(nsim0, nsimx, plot_std=False, min_overlap=0.0):
     catx = open_cat(nsimx)
 
     reader = csiborgtools.read.PairOverlap(cat0, catx, paths,
-                                           maxdist=155 / 0.705)
+                                           maxdist=155)
     mass = numpy.log10(reader.cat0("totpartmass"))
     dist = reader.dist(in_initial=False, norm_kind="r200c")
     overlap = reader.overlap(True)
@@ -373,7 +374,7 @@ def plot_mass_vs_separation(nsim0, nsimx, plot_std=False, min_overlap=0.0):
         ax.plot(xrange, numpy.polyval(p, xrange), color="red",
                 linestyle="--")
         fig.colorbar(cx, label="Bin counts")
-        ax.set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        ax.set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         ax.set_ylabel(r"$\log \langle \Delta R / R_{\rm 200c}\rangle$")
 
         fig.tight_layout()
@@ -460,10 +461,10 @@ def plot_maxoverlap_mass(nsim0):
         axs[0].plot(t, t + 0.2, color="red", linestyle="--", alpha=0.5)
         axs[0].plot(t, t - 0.2, color="red", linestyle="--", alpha=0.5)
 
-        axs[0].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[0].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[1].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[1].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[0].set_ylabel(r"Max. overlap mean of $\log M_{\rm tot} / M_\odot$")
+        axs[0].set_ylabel(r"Max. overlap mean of $\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[1].set_ylabel(r"Max. overlap std. of $\log M_{\rm tot} / M_\odot$")
+        axs[1].set_ylabel(r"Max. overlap std. of $\log M_{\rm tot} ~ [M_\odot / h]$")
 
         ims = [im0, im1]
         for i in range(2):
@@ -518,9 +519,9 @@ def plot_maxoverlapstat(nsim0, key):
         m = numpy.isfinite(key_val) & numpy.isfinite(mu)
         print("True to expectation corr: ", kendalltau(key_val[m], mu[m]))
 
-        axs[0].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[0].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[0].set_ylabel(r"Max. overlap mean of ${}$".format(key_label))
-        axs[1].set_xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        axs[1].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[1].set_ylabel(r"Max. overlap std. of ${}$".format(key_label))
         axs[2].set_xlabel(r"${}$".format(key_label))
         axs[2].set_ylabel(r"Max. overlap mean of ${}$".format(key_label))
@@ -622,10 +623,9 @@ def plot_mass_vs_expected_mass(nsim0, min_overlap=0, max_prob_nomatch=1):
                             gridsize=50, C=mass[mask],
                             reduce_C_function=numpy.nanmedian)
         axs[2].axhline(0, color="red", linestyle="--", alpha=0.5)
-
+        axs[0].set_xlabel(r"True $\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[0].set_xlabel(r"True $\log M_{\rm tot} / M_\odot$")
+        axs[0].set_ylabel(r"Expected $\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[0].set_ylabel(r"Expected $\log M_{\rm tot} / M_\odot$")
+        axs[1].set_xlabel(r"True $\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[1].set_xlabel(r"True $\log M_{\rm tot} / M_\odot$")
         axs[1].set_ylabel(r"Std. of $\sigma_{\log M_{\rm tot}}$")
         axs[2].set_xlabel(r"1 - median prob. of no match")
         axs[2].set_ylabel(r"$\log M_{\rm tot} - \log M_{\rm tot, exp}$")
@@ -636,7 +636,8 @@ def plot_mass_vs_expected_mass(nsim0, min_overlap=0, max_prob_nomatch=1):
         axs[0].plot(t, t - 0.2, color="red", linestyle="--", alpha=0.5)
 
         ims = [im0, im1, im2]
-        labels = ["Bin counts", "Bin counts", r"$\log M_{\rm tot}$"]
+        labels = ["Bin counts", "Bin counts",
+                  r"$\log M_{\rm tot} ~ [M_\odot / h]$"]
         for i in range(3):
             axins = inset_axes(axs[i], width="100%", height="5%",
                                loc='upper center', borderpad=-0.75)
@@ -758,10 +759,10 @@ def plot_dist(run, kind, kwargs, runs_to_mass, pulled_cdf=False, r200=None):
 
         fig, ax = plt.subplots()
         if run != "mass009":
-            ax.set_title(r"${} \leq \log M_{{\rm tot}} / M_\odot < {}$"
+            ax.set_title(r"${} \leq \log M_{{\rm tot}} / (M_\odot  h) < {}$"
                          .format(*runs_to_mass[run]), fontsize="small")
         else:
-            ax.set_title(r"$\log M_{{\rm tot}} / M_\odot \geq {}$"
+            ax.set_title(r"$\log M_{{\rm tot}} / (M_\odot h) \geq {}$"
                          .format(runs_to_mass[run][0]), fontsize="small")
         # Plot data
         nrad = y_csiborg.shape[0]
@@ -778,12 +779,12 @@ def plot_dist(run, kind, kwargs, runs_to_mass, pulled_cdf=False, r200=None):
             ax.plot(x2, y2, c="gray", ls="--",
                     label="Quijote" if i == 0 else None)
 
-        fig.colorbar(cmap, ax=ax, label=r"$R_{\rm dist}~[\mathrm{Mpc}]$")
+        fig.colorbar(cmap, ax=ax, label=r"$R_{\rm dist}~[\mathrm{Mpc} / h]$")
         ax.grid(alpha=0.5, lw=0.4)
         # Plot labels
         if pulled_cdf:
             if r200 is None:
-                ax.set_xlabel(r"$\tilde{r}_{1\mathrm{NN}}~[\mathrm{Mpc}]$")
+                ax.set_xlabel(r"$\tilde{r}_{1\mathrm{NN}}~[\mathrm{Mpc} / h]$")
                 if kind == "pdf":
                     ax.set_ylabel(r"$p(\tilde{r}_{1\mathrm{NN}})$")
                 else:
@@ -796,7 +797,7 @@ def plot_dist(run, kind, kwargs, runs_to_mass, pulled_cdf=False, r200=None):
                     ax.set_ylabel(r"$\mathrm{CDF}(\tilde{r}_{1\mathrm{NN}} / R_{200c})$")  # noqa
         else:
             if r200 is None:
-                ax.set_xlabel(r"$r_{1\mathrm{NN}}~[\mathrm{Mpc}]$")
+                ax.set_xlabel(r"$r_{1\mathrm{NN}}~[\mathrm{Mpc} / h]$")
                 if kind == "pdf":
                     ax.set_ylabel(r"$p(r_{1\mathrm{NN}})$")
                 else:
@@ -901,7 +902,7 @@ def plot_cdf_diff(runs, kwargs, pulled_cdf, runs_to_mass):
             ax.fill_between(r, *numpy.percentile(dy, [16, 84], axis=0),
                             alpha=0.5, color=cmap.to_rgba(runs_to_mass[i]))
         fig.colorbar(cmap, ax=ax, ticks=runs_to_mass,
-                     label=r"$\log M_{\rm tot} / M_\odot$")
+                     label=r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         ax.set_xlim(0.0, 55)
         ax.set_ylim(0)
 
@@ -909,17 +910,17 @@ def plot_cdf_diff(runs, kwargs, pulled_cdf, runs_to_mass):
 
         # Plot labels
         if pulled_cdf:
-            ax.set_xlabel(r"$\tilde{r}_{1\mathrm{NN}}~[\mathrm{Mpc}]$")
+            ax.set_xlabel(r"$\tilde{r}_{1\mathrm{NN}}~[\mathrm{Mpc} / h]$")
         else:
-            ax.set_xlabel(r"$r_{1\mathrm{NN}}~[\mathrm{Mpc}]$")
+            ax.set_xlabel(r"$r_{1\mathrm{NN}}~[\mathrm{Mpc} / h]$")
         ax.set_ylabel(r"$\Delta \mathrm{CDF}(r_{1\mathrm{NN}})$")
 
         # Plot labels
         if pulled_cdf:
-            ax.set_xlabel(r"$\tilde{r}_{1\mathrm{NN}}~[\mathrm{Mpc}]$")
+            ax.set_xlabel(r"$\tilde{r}_{1\mathrm{NN}}~[\mathrm{Mpc} / h]$")
             ax.set_ylabel(r"$\Delta \mathrm{CDF}(\tilde{r}_{1\mathrm{NN}})$")
         else:
-            ax.set_xlabel(r"$r_{1\mathrm{NN}}~[\mathrm{Mpc}]$")
+            ax.set_xlabel(r"$r_{1\mathrm{NN}}~[\mathrm{Mpc} / h]$")
             ax.set_ylabel(r"$\Delta \mathrm{CDF}(r_{1\mathrm{NN}})$")
 
         fig.tight_layout()
@@ -1104,7 +1105,7 @@ def plot_significance(simname, runs, nsim, nobs, kind, kwargs, runs_to_mass):
 
         cbar_ax = fig.add_axes([1.0, 0.125, 0.035, 0.85])
         fig.colorbar(cmap, cax=cbar_ax, ticks=runs_to_mass,
-                     label=r"$\log M_{\rm tot} / M_\odot$")
+                     label=r"$\log M_{\rm tot} ~ [M_\odot / h]$")
 
         ax[0].set_xlim(z[0], z[-1])
         ax[0].set_ylim(1e-5, 1.)
@@ -1216,7 +1217,7 @@ def plot_significance_vs_mass(simname, runs, nsim, nobs, kind, kwargs,
         corr = plt_utils.latex_float(*kendalltau(xs[mask], ys[mask]))
         plt.title(r"$\tau = {}, p = {}$".format(*corr), fontsize="small")
 
-        plt.xlabel(r"$\log M_{\rm tot} / M_\odot$")
+        plt.xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         if kind == "ks":
             plt.ylabel(r"$\log p$-value of $r_{1\mathrm{NN}}$ distribution")
             plt.ylim(top=0)