Merge branch 'master' of https://github.com/Richard-Sti/galomatch

README.md

@@ -1,11 +1,16 @@
-# CSiBORG analysis tools :dart:
+# CSiBORG tools
 
-## TODO :scroll:
+## :scroll: Short-term TODO
-- [ ] Calculate $M_{\rm vir}, R_{\rm vir}, c$ from $R_s, \rho_0, \ldots$
+- [x] Calculate $M_{\rm vir}, R_{\rm vir}, c$ from $R_s, \rho_0, \ldots$
-- [ ] Calculate $M_{\rm 500c}$ by sphere shrinking
+- [x] In `NFWPosterior` correct for the radius in which particles are fitted.
-- [ ] Calculate the cross-correlation in CSiBORG. Should see the scale of the constraints?
+- [x] Calculate $M_{\rm 500c}$ by sphere shrinking
+- [x] Change to log10 of the scale factor
+- [ ] Calculate uncertainty on $R_{\rm s}$, switch to `JAX` and get gradients.
+
+
+## :hourglass: Long-term TODO
 - [ ] Improve file naming system
+- [ ] Calculate the cross-correlation in CSiBORG. Should see the scale of the constraints?
 
 
-## Open questions :bulb:
+## :bulb: Open questions
-- Get uncertainty on the fitted $R_{\rm s}$? If so get this directly from JAX.

csiborgtools/fits/halo.py

@@ -18,6 +18,7 @@ Tools for splitting the particles and a clump object.
 
 
 import numpy
+from scipy.optimize import minimize_scalar
 from os import remove
 from warnings import warn
 from os.path import join
@@ -235,7 +236,7 @@ def pick_single_clump(n, particles, particle_clumps, clumps):
 
 
 class Clump:
-    """
+    r"""
     A clump (halo) object to handle the particles and their clump's data.
 
     Parameters
@@ -254,28 +255,40 @@ class Clump:
         Clump center coordinate along the y-axis.
     z0 : float
         Clump center coordinate along the z-axis.
-    clump_mass : float
+    clump_mass : float, optional
-        Mass of the clump.
+        Mass of the clump. By default not set.
-    vx : 1-dimensional array
+    vx : 1-dimensional array, optional
-        Particle velocity along the x-axis.
+        Particle velocity along the x-axis. By default not set.
-    vy : 1-dimensional array
+    vy : 1-dimensional array, optional
-        Particle velocity along the y-axis.
+        Particle velocity along the y-axis. By default not set.
-    vz : 1-dimensional array
+    vz : 1-dimensional array, optional
-        Particle velocity along the z-axis.
+        Particle velocity along the z-axis. By default not set.
+    index : int, optional
+        The halo finder index of this clump. By default not set.
+    rhoc : float, optional
+        The critical density :math:`\rho_c` at this snapshot in box units. By
+        default not set.
     """
     _r = None
+    _rmin = None
+    _rmax = None
     _pos = None
     _clump_pos = None
     _clump_mass = None
     _vel = None
+    _Npart = None
+    _index = None
+    _rhoc = None
 
     def __init__(self, x, y, z, m, x0, y0, z0, clump_mass=None,
-                 vx=None, vy=None, vz=None):
+                 vx=None, vy=None, vz=None, index=None, rhoc=None):
         self.pos = (x, y, z, x0, y0, z0)
         self.clump_pos = (x0, y0, z0)
         self.clump_mass = clump_mass
         self.vel = (vx, vy, vz)
         self.m = m
+        self.index = index
+        self.rhoc = rhoc
 
     @property
     def pos(self):
@@ -294,7 +307,46 @@ class Clump:
         """Sets `pos` and calculates radial distance."""
         x, y, z, x0, y0, z0 = X
         self._pos = numpy.vstack([x - x0, y - y0, z - z0]).T
-        self.r = numpy.sum(self.pos**2, axis=1)**0.5
+        self._r = numpy.sum(self.pos**2, axis=1)**0.5
+        self._rmin = numpy.min(self._r)
+        self._rmax = numpy.max(self._r)
+        self._Npart = self._r.size
+
+    @property
+    def r(self):
+        """
+        Radial distance of the particles from the clump peak.
+
+        Returns
+        -------
+        r : 1-dimensional array
+            Array of shape `(n_particles, )`.
+        """
+        return self._r
+
+    @property
+    def rmin(self):
+        """
+        The minimum radial distance of a particle.
+
+        Returns
+        -------
+        rmin : float
+            The minimum distance.
+        """
+        return self._rmin
+
+    @property
+    def rmax(self):
+        """
+        The maximum radial distance of a particle.
+
+        Returns
+        -------
+        rmin : float
+            The maximum distance.
+        """
+        return self._rmax
 
     @property
     def Npart(self):
@@ -306,7 +358,7 @@ class Clump:
         Npart : int
             Number of particles.
         """
-        return self.r.size
+        return self._Npart
 
     @property
     def clump_pos(self):
@@ -338,12 +390,14 @@ class Clump:
         mass : float
             Clump mass.
         """
+        if self._clump_mass is None:
+            raise ValueError("Clump mass `clump_mass` has not been set.")
         return self._clump_mass
 
     @clump_mass.setter
     def clump_mass(self, mass):
         """Sets `clump_mass`, making sure it is a float."""
-        if not isinstance(mass, float):
+        if mass is not None and not isinstance(mass, float):
             raise ValueError("`clump_mass` must be a float.")
         self._clump_mass = mass
 
@@ -391,25 +445,46 @@ class Clump:
         self._m = m
 
     @property
-    def r(self):
+    def index(self):
         """
-        Radial distance of particles from the clump peak.
+        The halo finder clump index.
 
         Returns
         -------
-        r : 1-dimensional array
+        hindex : int
-            Array of shape `(n_particles, )`
+            The index.
         """
-        return self._r
+        if self._index is None:
+            raise ValueError("Halo index `hindex` has not been set.")
+        return self._index
 
-    @r.setter
+    @index.setter
-    def r(self, r):
+    def index(self, n):
-        """Sets `r`. Again checks the shape."""
+        """Sets the halo index, making sure it is an integer."""
-        if not isinstance(r, numpy.ndarray) and r.ndim == 1:
+        if n is not None and not (isinstance(n, (int, numpy.int64)) and n > 0):
-            raise TypeError("`r` must be a 1-dimensional array.")
+            raise ValueError("Halo index `index` must be an integer > 0.")
-        if not numpy.all(r >= 0):
+        self._index = n
-            raise ValueError("`r` larger than zero.")
+
-        self._r = r
+    @property
+    def rhoc(self):
+        """
+        The critical density :math:`\rho_c` at this snapshot in box units.
+
+        Returns
+        -------
+        rhoc : float
+            The critical density.
+        """
+        if self._rhoc is None:
+            raise ValueError("The critical density `rhoc` has not been set.")
+        return self._rhoc
+
+    @rhoc.setter
+    def rhoc(self, rhoc):
+        """Sets the critical density. Makes sure it is > 0."""
+        if rhoc is not None and not rhoc > 0:
+            raise ValueError("Critical density `rho_c` must be > 0.")
+        self._rhoc = rhoc
 
     @property
     def total_particle_mass(self):
@@ -435,8 +510,165 @@ class Clump:
         """
         return numpy.mean(self.pos + self.clump_pos, axis=0)
 
+    def enclosed_spherical_mass(self, rmax, rmin=None):
+        """
+        The enclosed spherical mass between two radii. All quantities remain
+        in the box units.
+
+        Parameters
+        ----------
+        rmax : float
+            The maximum radial distance.
+        rmin : float, optional
+            The minimum radial distance. By default the radial distance of the
+            closest particle.
+
+        Returns
+        -------
+        M_enclosed : float
+            The enclosed mass.
+        """
+        rmin = self.rmin if rmin is None else rmin
+        return numpy.sum(self.m[(self.r >= rmin) & (self.r <= rmax)])
+
+    def enclosed_spherical_density(self, rmax, rmin=None):
+        """
+        The enclosed spherical density between two radii. All quantities
+        remain in box units.
+
+        Parameters
+        ----------
+        rmax : float
+            The maximum radial distance.
+        rmin : float, optional
+            The minimum radial distance. By default the radial distance of the
+            closest particle.
+
+        Returns
+        -------
+        rho_enclosed : float
+            The enclosed density.
+        """
+        rmin = self.rmin if rmin is None else rmin
+        M = self.enclosed_spherical_mass(rmax, rmin)
+        V = 4 * numpy.pi / 3 * (rmax**3 - rmin**3)
+        return M / V
+
+    def radius_enclosed_overdensity(self, delta):
+        r"""
+        Radius of where the mean enclosed spherical density reaches a multiple
+        of the critical radius at a given redshift `self.rho_c`. Returns
+        `numpy.nan` if the fit does not converge. Note that `rhoc` must be in
+        box units!
+
+        Parameters
+        ----------
+        delta : int or float
+            The :math:`\delta_{\rm x}` parameters where :math:`\mathrm{x}` is
+            the overdensity multiple.
+
+        Returns
+        -------
+        rx : float
+            The radius where the enclosed density reaches required value.
+        """
+        # Loss function to minimise
+        def loss(r):
+            return abs(self.enclosed_spherical_density(r, self.rmin)
+                       - delta * self.rhoc)
+
+        res = minimize_scalar(loss, bounds=(self.rmin, self.rmax),
+                              method='bounded')
+        return res.x if res.success else numpy.nan
+
+    @property
+    def r200(self):
+        r"""
+        The radius at which the mean spherical density reaches 200 times
+        the critical density, :math:`R_{200c}`. Returns `numpy.nan` if the
+        estimate fails.
+
+        Returns
+        -------
+        r200 : float
+            The R200c radius
+        """
+        return self.radius_enclosed_overdensity(200)
+
+    @property
+    def r178(self):
+        r"""
+        The radius at which the mean spherical density reaches 178 times
+        the critical density, :math:`R_{178c}`. Returns `numpy.nan` if the
+        estimate fails.
+
+        Returns
+        -------
+        r178 : float
+            The R178c radius
+        """
+        return self.radius_enclosed_overdensity(178)
+
+    @property
+    def r500(self):
+        r"""
+        The radius at which the mean spherical density reaches 500 times
+        the critical density, :math:`R_{500c}`. Returns `numpy.nan` if the
+        estimate fails.
+
+        Returns
+        -------
+        r500 : float
+            The R500c radius
+        """
+        return self.radius_enclosed_overdensity(500)
+
+    @property
+    def m200(self):
+        r"""
+        The mass enclosed within the :math:`R_{200c}` region, obtained from
+        `self.r200`. Returns `numpy.nan` if the radius estimate fails.
+
+        Returns
+        -------
+        m200 : float
+            The M200 mass
+        """
+        r200 = self.radius_enclosed_overdensity(200)
+        return self.enclosed_spherical_mass(r200)
+
+    @property
+    def m178(self):
+        r"""
+        The mass enclosed within the :math:`R_{178c}` region, obtained from
+        `self.r178`. This is approximately the virial mass, though this notion
+        depends on the dynamical state of the clump. Returns `numpy.nan` if
+        the radius estimate fails.
+
+        Returns
+        -------
+        m178 : float
+            The M178 mass
+        """
+        r178 = self.radius_enclosed_overdensity(178)
+        return self.enclosed_spherical_mass(r178)
+
+    @property
+    def m500(self):
+        r"""
+        The mass enclosed within the :math:`R_{500c}` region, obtained from
+        `self.r500`. Returns `numpy.nan` if the radius estimate fails.
+
+        Returns
+        -------
+        m500 : float
+            The M500 mass
+        """
+        r500 = self.radius_enclosed_overdensity(500)
+        return self.enclosed_spherical_mass(r500)
+
     @classmethod
-    def from_arrays(cls, particles, clump):
+    def from_arrays(cls, particles, clump, rhoc=None):
         """
         Initialises `Halo` from `particles` containing the relevant particle
         information and its `clump` information.
@@ -452,14 +684,15 @@ class Clump:
 
         Returns
         -------
-        halo : `Halo`
+        clump : `Clump`
-            An initialised halo object.
+            An initialised clump object.
         """
         x, y, z, m = (particles[p] for p in ["x", "y", "z", "M"])
-        x0, y0, z0, cl_mass = (
+        x0, y0, z0, cl_mass, hindex = (
-            clump[p] for p in ["peak_x", "peak_y", "peak_z", "mass_cl"])
+            clump[p] for p in ["peak_x", "peak_y", "peak_z", "mass_cl",
+                               "index"])
         try:
             vx, vy, vz = (particles[p] for p in ["vx", "vy", "vz"])
         except ValueError:
             vx, vy, vz = None, None, None
-        return cls(x, y, z, m, x0, y0, z0, cl_mass, vx, vy, vz)
+        return cls(x, y, z, m, x0, y0, z0, cl_mass, vx, vy, vz, hindex, rhoc)

csiborgtools/fits/haloprofile.py

@@ -43,7 +43,8 @@ class NFWProfile:
     def __init__(self):
         pass
 
-    def profile(self, r, Rs, rho0):
+    @staticmethod
+    def profile(r, Rs, rho0):
         r"""
         Halo profile evaluated at :math:`r`.
 
@@ -64,7 +65,8 @@ class NFWProfile:
         x = r / Rs
         return rho0 / (x * (1 + x)**2)
 
-    def logprofile(self, r, Rs, rho0):
+    @staticmethod
+    def logprofile(r, Rs, rho0):
         r"""
         Natural logarithm of the halo profile evaluated at :math:`r`.
 
@@ -85,7 +87,8 @@ class NFWProfile:
         x = r / Rs
         return numpy.log(rho0) - numpy.log(x) - 2 * numpy.log(1 + x)
 
-    def enclosed_mass(self, r, Rs, rho0):
+    @staticmethod
+    def enclosed_mass(r, Rs, rho0):
         r"""
         Enclosed mass  of a NFW profile in radius :math:`r`.
 
@@ -198,16 +201,18 @@ class NFWProfile:
 
 class NFWPosterior(NFWProfile):
     r"""
-    Posterior of for fitting the NFW profile in the range specified by the
+    Posterior for fitting the NFW profile in the range specified by the
-    closest and further particle. The likelihood is calculated as
+    closest particle and the :math:`r_{200c}` radius. The likelihood is
+    calculated as
 
     .. math::
-        \frac{4\pi r^2 \rho(r)} {M(r_\min, r_\max)} \frac{m}{M / N}
+        \frac{4\pi r^2 \rho(r)} {M(r_{\min} r_{200c})} \frac{m}{M / N}
 
-    where :math:`M(r_\min, r_\max)` is the enclosed mass between the closest
+    where :math:`M(r_{\min} r_{200c}))` is the NFW enclosed mass between the
-    and further particle as expected from a NFW profile, :math:`m` is the
+    closest particle and the :math:`r_{200c}` radius, :math:`m` is the particle
-    particle mass, :math:`M` is the sum of the particle masses and :math:`N`
+    mass, :math:`M` is the sum of the particle masses and :math:`N` is the
-    is the number of particles.
+    number of particles. Calculated only using particles within the
+    above-mentioned range.
 
     Paramaters
     ----------
@@ -215,10 +220,12 @@ class NFWPosterior(NFWProfile):
         Clump object containing the particles and clump information.
     """
     _clump = None
-    _N = None
+    _binsguess = 10
+    _r = None
+    _Npart = None
+    _m = None
     _rmin = None
     _rmax = None
-    _binsguess = 10
 
     def __init__(self, clump):
         # Initialise the NFW profile
@@ -228,7 +235,8 @@ class NFWPosterior(NFWProfile):
     @property
     def clump(self):
         """
-        Clump object.
+        Clump object containig all particles, i.e. ones beyond :math:`R_{200c}`
+        as well.
 
         Returns
         -------
@@ -237,44 +245,49 @@ class NFWPosterior(NFWProfile):
         """
         return self._clump
 
-    @clump.setter
-    def clump(self, clump):
-        """Sets `clump` and precalculates useful things."""
-        if not isinstance(clump, Clump):
-            raise TypeError(
-                "`clump` must be :py:class:`csiborgtools.fits.Clump` type. "
-                "Currently `{}`".format(type(clump)))
-        self._clump = clump
-        # Set here the rest of the radial info
-        self._rmax = numpy.max(self.r)
-        # r_min could be zero if particle in the potential minimum.
-        # Set it to the closest particle that is at distance > 0
-        self._rmin = numpy.sort(self.r[self.r > 0])[0]
-        self._logrmin = numpy.log(self.rmin)
-        self._logrmax = numpy.log(self.rmax)
-        self._logprior_volume = numpy.log(self._logrmax - self._logrmin)
-        self._N = self.r.size
-        # Precalculate useful things
-        self._logMtot = numpy.log(numpy.sum(self.clump.m))
-        gamma = 4 * numpy.pi * self.r**2 * self.clump.m * self.N
-        self._ll0 = numpy.sum(numpy.log(gamma)) - self.N * self._logMtot
-
     @property
     def r(self):
-        """
+        r"""
-        Radial distance of particles.
+        Radial distance of particles used to fit the NFW profile, i.e. the ones
+        whose radial distance is less than :math:`R_{\rm 200c}`.
 
         Returns
         -------
         r : 1-dimensional array
-            Radial distance of particles.
+            Array of shape `(n_particles, )`.
         """
-        return self.clump.r
+        return self._r
+
+    @property
+    def Npart(self):
+        r"""
+        Number of particles used to fit the NFW profile, i.e. the ones
+        whose radial distance is less than :math:`R_{\rm 200c}`.
+
+        Returns
+        -------
+        Npart : int
+            Number of particles.
+        """
+        return self._Npart
+
+    @property
+    def m(self):
+        r"""
+        Mass of particles used to fit the NFW profile, i.e. the ones
+        whose radial distance is less than :math:`R_{\rm 200c}`.
+
+        Returns
+        -------
+        r : 1-dimensional array
+            Array of shape `(n_particles, )`.
+        """
+        return self._m
 
     @property
     def rmin(self):
         """
-        Minimum radial distance of a particle belonging to this clump.
+        The minimum radial distance of a particle.
 
         Returns
         -------
@@ -285,30 +298,50 @@ class NFWPosterior(NFWProfile):
 
     @property
     def rmax(self):
-        """
+        r"""
-        Maximum radial distance of a particle belonging to this clump.
+        The maximum radial distance used to fit the profile, here takem to be
+        the :math:`R_{\rm 200c}`.
 
         Returns
         -------
-        rmin : float
+        rmax : float
-            The maximum distance.
+            The R200c radius.
         """
         return self._rmax
 
-    @property
+    @clump.setter
-    def N(self):
+    def clump(self, clump):
-        """
+        """Sets `clump` and precalculates useful things."""
-        The number of particles in this clump.
+        if not isinstance(clump, Clump):
-
+            raise TypeError(
-        Returns
+                "`clump` must be :py:class:`csiborgtools.fits.Clump` type. "
-        -------
+                "Currently `{}`".format(type(clump)))
-        N : int
+        self._clump = clump
-            Number of particles.
+        # The minimum separation
-        """
+        rmin = self.clump.rmin
-        return self._N
+        rmax = self.clump.r200
+        # Set the distances
+        self._rmin = rmin
+        self._rmax = rmax
+        # Set particles that will be used to fit the halo
+        mask_r200 = (self.clump.r >= rmin) & (self.clump.r <= rmax)
+        self._r = self.clump.r[mask_r200]
+        self._m = self.clump.m[mask_r200]
+        self._Npart = self._r.size
+        # Ensure that the minimum separation is > 0 for finite log
+        if self.rmin > 0:
+            self._logrmin = numpy.log10(self.rmin)
+        else:
+            self._logrmin = numpy.log10(numpy.min(self.r[self.r > 0]))
+        self._logrmax = numpy.log10(self.rmax)
+        self._logprior_volume = numpy.log(self._logrmax - self._logrmin)
+        # Precalculate useful things
+        self._logMtot = numpy.log(numpy.sum(self.m))
+        gamma = 4 * numpy.pi * self.r**2 * self.m * self.Npart
+        self._ll0 = numpy.sum(numpy.log(gamma)) - self.Npart * self._logMtot
 
     def rho0_from_logRs(self, logRs):
-        """
+        r"""
         Obtain :math:`\rho_0` of the NFW profile from the integral constraint
         on total mass. Calculated as the ratio between the total particle mass
         and the enclosed NFW profile mass.
@@ -324,8 +357,8 @@ class NFWPosterior(NFWProfile):
             The NFW density parameter.
         """
         Mtot = numpy.exp(self._logMtot)
-        Rs = numpy.exp(logRs)
+        Mnfw_norm = self.bounded_enclosed_mass(self.rmin, self.rmax,
-        Mnfw_norm = self.bounded_enclosed_mass(self.rmin, self.rmax, Rs, 1)
+                                               10**logRs, 1)
         return Mtot / Mnfw_norm
 
     def logprior(self, logRs):
@@ -360,11 +393,12 @@ class NFWPosterior(NFWProfile):
         ll : float
             The logarithmic likelihood.
         """
-        Rs = numpy.exp(logRs)
+        Rs = 10**logRs
         # Expected enclosed mass from a NFW
-        Mnfw = self.bounded_enclosed_mass(self.rmin, self.rmax, Rs, 1)
+        Mnfw = self.bounded_enclosed_mass(self.rmin, self.rmax,
+                                          Rs, 1)
         ll = self._ll0 + numpy.sum(self.logprofile(self.r, Rs, 1))
-        return ll - self.N * numpy.log(Mnfw)
+        return ll - self.Npart * numpy.log(Mnfw)
 
     @property
     def initlogRs(self):
@@ -378,7 +412,8 @@ class NFWPosterior(NFWProfile):
         initlogRs : float
             The initial guess of :math:`\log R_{\rm s}`.
         """
-        bins = numpy.linspace(self.rmin, self.rmax, self._binsguess)
+        bins = numpy.linspace(self.rmin, self.rmax,
+                              self._binsguess)
         counts, edges = numpy.histogram(self.r, bins)
         return numpy.log(edges[numpy.argmax(counts)])
 
@@ -401,61 +436,19 @@ class NFWPosterior(NFWProfile):
             return - numpy.infty
         return self.loglikelihood(logRs) + lp
 
-    def hamiltonian(self, logRs):
-        """
-        Negative logarithmic posterior (i.e. the Hamiltonian).
-
-        Parameters
-        ----------
-        logRs : float
-            Logarithmic scale factor in units matching the coordinates.
-
-        Returns
-        -------
-        neg_lpost : float
-        The Hamiltonian.
-        """
-        return - self(logRs)
-
     def maxpost_logRs(self):
         r"""
         Maximum a-posterio estimate of the scale radius :math:`\log R_{\rm s}`.
+        Returns the scale radius if the fit converged, otherwise `numpy.nan`.
 
         Returns
         -------
-        res : `scipy.optimize.OptimizeResult`
+        res : float
-            Optimisation result.
+            The scale radius.
         """
-        bounds = (self._logrmin, self._logrmax)
+        # Loss function to optimize
-        return minimize_scalar(
+        def loss(logRs):
-            self.hamiltonian, bounds=bounds, method='bounded')
+            return - self(logRs)
-
+        res = minimize_scalar(loss, bounds=(self._logrmin, self._logrmax),
-    @classmethod
+                              method='bounded')
-    def from_coords(cls, x, y, z, m, x0, y0, z0):
+        return res.x if res.success else numpy.nan
-        """
-        Initiate `NFWPosterior` from a set of Cartesian coordinates.
-
-        Parameters
-        ----------
-        x : 1-dimensional array
-            Particle coordinates along the x-axis.
-        y : 1-dimensional array
-            Particle coordinates along the y-axis.
-        z : 1-dimensional array
-            Particle coordinates along the z-axis.
-        m : 1-dimensional array
-            Particle masses.
-        x0 : float
-            Halo center coordinate along the x-axis.
-        y0 : float
-            Halo center coordinate along the y-axis.
-        z0 : float
-            Halo center coordinate along the z-axis.
-
-        Returns
-        -------
-        post : `NFWPosterior`
-            Initiated `NFWPosterior` instance.
-        """
-        r = numpy.sqrt((x - x0)**2 + (y - y0)**2 + (z - z0)**2)
-        return cls(r, m)

csiborgtools/io/__init__.py

@@ -17,6 +17,7 @@ from .readsim import (get_csiborg_ids, get_sim_path, get_snapshots,  # noqa
                       get_snapshot_path, read_info, nparts_to_start_ind,  # noqa
                       open_particle, open_unbinding, read_particle,  # noqa
                       drop_zero_indx,  # noqa
-                      read_clumpid, read_clumps, read_mmain)  # noqa
+                      read_clumpid, read_clumps, read_mmain,  # noqa
+                      merge_mmain_to_clumps)  # noqa
 from .readobs import (read_planck2015, read_2mpp)  # noqa
 from .outsim import (dump_split, combine_splits)  # noqa

csiborgtools/io/outsim.py

@@ -85,9 +85,6 @@ def combine_splits(Nsplits, Nsim, Nsnap, outdir, cols_add, remove_splits=False,
     out : structured array
         Clump array with appended results from the splits.
     """
-    # Will be grabbing these columns from each split
-    cols_add = [("npart", I64), ("totpartmass", F64), ("logRs", F64),
-                ("rho0", F64)]
     # Load clumps to see how many there are and will add to this array
     simpath = get_sim_path(Nsim)
     clumps = read_clumps(Nsnap, simpath, cols=None)

csiborgtools/io/readsim.py

@@ -23,7 +23,7 @@ from os.path import (join, isfile)
 from glob import glob
 from tqdm import tqdm
 
-from ..utils import cols_to_structured
+from ..utils import (cols_to_structured, add_columns)
 
 
 F16 = numpy.float16
@@ -495,3 +495,30 @@ def read_mmain(n, srcdir, fname="Mmain_{}.npy"):
         out[name] = arr[:, i]
 
     return out
+
+
+def merge_mmain_to_clumps(clumps, mmain):
+    """
+    Merge columns from the `mmain` files to the `clump` file, matches them
+    by their halo index while assuming that the indices `index` in both arrays
+    are sorted.
+
+    Parameters
+    ----------
+    clumps : structured array
+        Clumps structured array.
+    mmain : structured array
+        Parent halo array whose information is to be merged into `clumps`.
+
+    Returns
+    -------
+    out : structured array
+        Array with added columns.
+    """
+    X = numpy.full((clumps.size, 2), numpy.nan)
+    # Mask of which clumps have a mmain index
+    mask = numpy.isin(clumps["index"], mmain["index"])
+
+    X[mask, 0] = mmain["mass_cl"]
+    X[mask, 1] = mmain["sub_frac"]
+    return add_columns(clumps, X, ["mass_mmain", "sub_frac"])

csiborgtools/units/box_units.py

@@ -16,24 +16,16 @@
 Simulation box unit transformations.
 """
 
-
+import numpy
 from astropy.cosmology import LambdaCDM
 from astropy import (constants, units)
 from ..io import read_info
 
 
-# Conversion factors
-MSUNCGS = constants.M_sun.cgs.value
-KPC_TO_CM = 3.0856775814913673e+21
-PI = 3.1415926535897932384626433
-
-
 class BoxUnits:
     r"""
     Box units class for converting between box and physical units.
 
-    TODO: check factors of :math:`a` in mass and density transformations
-
     Paramaters
     ----------
     Nsnap : int
@@ -41,6 +33,7 @@ class BoxUnits:
     simpath : str
         Path to the simulation where its snapshot index folders are stored.
     """
+    _cosmo = None
 
     def __init__(self, Nsnap, simpath):
         """
@@ -51,21 +44,105 @@ class BoxUnits:
                 "omega_m", "omega_l", "omega_k", "omega_b",
                 "unit_l", "unit_d", "unit_t"]
         for par in pars:
-            setattr(self, par, float(info[par]))
+            setattr(self, "_" + par, float(info[par]))
 
-        self.h = self.H0 / 100
+        self._cosmo = LambdaCDM(H0=self._H0, Om0=self._omega_m,
-        self.cosmo = LambdaCDM(H0=self.H0, Om0=self.omega_m, Ode0=self.omega_l,
+                                Ode0=self._omega_l, Tcmb0=2.725 * units.K,
-                               Tcmb0=2.725 * units.K, Ob0=self.omega_b)
+                                Ob0=self._omega_b)
-        # Constants in box units
+        self._Msuncgs = constants.M_sun.cgs.value  # Solar mass in grams
-        self.G = constants.G.cgs.value * (self.unit_d * self.unit_t ** 2)
+
-        self.H0 = self.H0 * 1e5 / (1e3 * KPC_TO_CM) * self.unit_t
+    @property
-        self.c = constants.c.cgs.value * self.unit_t / self.unit_l
+    def cosmo(self):
-        self.rho_crit = 3 * self.H0 ** 2 / (8 * PI * self.G)
+        """
+        The  box cosmology.
+
+        Returns
+        -------
+        cosmo : `astropy.cosmology.LambdaCDM`
+            The CSiBORG cosmology.
+        """
+        return self._cosmo
+
+    @property
+    def H0(self):
+        r"""
+        The Hubble parameter at the time of the snapshot
+        in :math:`\mathrm{Mpc} / \mathrm{km} / \mathrm{s}`.
+
+        Returns
+        -------
+        H0 : float
+            Hubble constant.
+        """
+        return self._H0
+
+    @property
+    def h(self):
+        r"""
+        The little 'h` parameter at the time of the snapshot.
+
+        Returns
+        -------
+        h : float
+            The little h
+        """
+        return self._H0 / 100
+
+    @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
+            The gravitational constant.
+        """
+        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
+            The Hubble constant.
+        """
+        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
+            The speed of light.
+        """
+        return constants.c.cgs.value * self._unit_t / self._unit_l
+
+    @property
+    def box_rhoc(self):
+        """
+        Critical density in box units.
+
+        Returns
+        -------
+        rhoc : float
+            The critical density.
+        """
+
+        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`).
+        :math:`h=0.705`). It appears that `self.unit_l` must be in
+        :math:`\mathrm{cm}`.
 
         Parameters
         ----------
@@ -77,7 +154,7 @@ class BoxUnits:
         length : foat
             Length in :math:`\mathrm{ckpc}`
         """
-        return length * self.unit_l / KPC_TO_CM / self.aexp
+        return length * (self._unit_l / units.kpc.to(units.cm) / self._aexp)
 
     def kpc2box(self, length):
         r"""
@@ -94,7 +171,7 @@ class BoxUnits:
         length : foat
             Length in box units.
         """
-        return length / self.unit_l * KPC_TO_CM * self.aexp
+        return length / (self._unit_l / units.kpc.to(units.cm) / self._aexp)
 
     def solarmass2box(self, mass):
         r"""
@@ -110,11 +187,13 @@ class BoxUnits:
         mass : float
             Mass in box units.
         """
-        return mass / self.unit_d / (self.unit_l**3 / MSUNCGS)
+        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
         ----------
@@ -126,7 +205,7 @@ class BoxUnits:
         mass : float
             Mass in :math:`M_\odot`.
         """
-        return mass * self.unit_d * self.unit_l**3 / MSUNCGS
+        return mass * (self._unit_d * self._unit_l**3) / self._Msuncgs
 
     def box2dens(self, density):
         r"""
@@ -143,7 +222,8 @@ class BoxUnits:
         density : float
             Density in :math:`M_\odot / \mathrm{pc}^3`.
         """
-        return density * self.unit_d / MSUNCGS * (KPC_TO_CM * 1e-3)**3
+        return (density * self._unit_d / self._Msuncgs
+                * (units.pc.to(units.cm))**3)
 
     def dens2box(self, density):
         r"""
@@ -160,4 +240,5 @@ class BoxUnits:
         density : float
             Density in box units.
         """
-        return density / self.unit_d * MSUNCGS / (KPC_TO_CM * 1e-3)**3
+        return (density / self._unit_d * self._Msuncgs
+                / (units.pc.to(units.cm))**3)

scripts/run_fit_halos.py

@@ -44,7 +44,10 @@ nproc = comm.Get_size()
 dumpdir = utils.dumpdir
 loaddir = join(utils.dumpdir, "temp")
 cols_collect = [("npart", I64), ("totpartmass", F64), ("logRs", F64),
-                ("rho0", F64)]
+                ("rho0", F64), ("rmin", F64), ("rmax", F64),
+                ("r200", F64), ("r178", F64), ("r500", F64),
+                ("m200", F64), ("m178", F64), ("m500", F64)]
+
 # NOTE later loop over sims too
 Nsim = Nsims[0]
 
@@ -56,7 +59,9 @@ for Nsplit in jobs:
 
     N = clumps.size
     cols = [("index", I64), ("npart", I64), ("totpartmass", F64),
-            ("logRs", F64), ("rho0", F64)]
+            ("logRs", F64), ("rho0", F64), ("rmin", F64), ("rmax", F64),
+            ("r200", F64), ("r178", F64), ("r500", F64),
+            ("m200", F64), ("m178", F64), ("m500", F64)]
     out = csiborgtools.utils.cols_to_structured(N, cols)
     out["index"] = clumps["index"]
 
@@ -65,15 +70,25 @@ for Nsplit in jobs:
         xs = csiborgtools.fits.pick_single_clump(n, parts, part_clumps, clumps)
         clump = csiborgtools.fits.Clump.from_arrays(*xs)
         out["npart"][n] = clump.Npart
+        out["rmin"][n] = clump.rmin
+        out["rmax"][n] = clump.rmax
         out["totpartmass"][n] = clump.total_particle_mass
+        out["r200"][n] = clump.r200
+        out["r178"][n] = clump.r178
+        out["r500"][n] = clump.r500
+        out["m200"][n] = clump.m200
+        out["m178"][n] = clump.m178
+        out["m500"][n] = clump.m200
 
         # NFW profile fit
-        if clump.Npart > 10:
+        if clump.Npart > 10 and numpy.isfinite(out["r200"][n]):
+            # NOTE here it calculates the r200 again, but its fast so does not
+            # matter anyway.
             nfwpost = csiborgtools.fits.NFWPosterior(clump)
             logRs = nfwpost.maxpost_logRs()
-            if logRs.success:
+            if not numpy.isnan(logRs):
-                out["logRs"][n] = logRs.x
+                out["logRs"][n] = logRs
-                out["rho0"][n] = nfwpost.rho0_from_logRs(logRs.x)
+                out["rho0"][n] = nfwpost.rho0_from_logRs(logRs)
 
     csiborgtools.io.dump_split(out, Nsplit, Nsim, Nsnap, dumpdir)