Improving halo fits (#76)

* Add periodic distances

* Little corrections

* Fix little bug

* Modernise the script

* Small updates

* Remove clump

* Add new halo routines

* Fix weights

* Modernise the script

* Add check ups on convergence

* More convergence check ups

* Edit bounds

* Add default argument

* Update fit heuristic and NaNs

* Change maxiter

* Switch NFW minimization to log-sapce

* Remove print statement

* Turn convert_from_box abstract property required for all boxes.

* Move files

* Simplify script

* Improve the argument parser

* Remove optinal argument

* Improve argument parser

* Add a minimum concentration limit

csiborgtools/fits/__init__.py

@@ -12,7 +12,6 @@
 # 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.
-from .halo import (Clump, Halo, delta2ncells, dist_centmass,  # noqa
+from .halo import (Halo, delta2ncells, center_of_mass,  # noqa
-                   number_counts)
+                   periodic_distance, shift_to_center_of_box, number_counts)  # noqa
-from .haloprofile import NFWPosterior, NFWProfile  # noqa
 from .utils import split_jobs  # noqa

csiborgtools/fits/halo.py

@@ -12,12 +12,12 @@
 # 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.
-"""A clump object."""
+"""A halo object."""
 from abc import ABC
 
 import numpy
-
 from numba import jit
+from scipy.optimize import minimize
 
 
 class BaseStructure(ABC):
@@ -26,7 +26,6 @@ class BaseStructure(ABC):
     """
 
     _particles = None
-    _info = None
     _box = None
 
     @property
@@ -45,29 +44,14 @@ class BaseStructure(ABC):
         assert particles.ndim == 2 and particles.shape[1] == 7
         self._particles = particles
 
-    @property
-    def info(self):
-        """
-        Array containing information from the clump finder.
-
-        Returns
-        -------
-        info : structured array
-        """
-        return self._info
-
-    @info.setter
-    def info(self, info):
-        self._info = info
-
     @property
     def box(self):
         """
-        CSiBORG box object handling unit conversion.
+        Box object handling unit conversion.
 
         Returns
         -------
-        box : :py:class:`csiborgtools.units.CSiBORGBox`
+        box : Object derived from :py:class:`csiborgtools.units.BaseBox`
         """
         return self._box
 
@@ -82,14 +66,13 @@ class BaseStructure(ABC):
     @property
     def pos(self):
         """
-        Cartesian particle coordinates centered at the object.
+        Cartesian particle coordinates in the box coordinate system.
 
         Returns
         -------
         pos : 2-dimensional array of shape `(n_particles, 3)`.
         """
-        ps = ("x", "y", "z")
+        return numpy.vstack([self[p] for p in ('x', 'y', 'z')]).T
-        return numpy.vstack([self[p] - self.info[p] for p in ps]).T
 
     @property
     def vel(self):
@@ -102,98 +85,130 @@ class BaseStructure(ABC):
         """
         return numpy.vstack([self[p] for p in ("vx", "vy", "vz")]).T
 
-    @property
+    def spherical_overdensity_mass(self, delta_mult, kind="crit", tol=1e-8,
-    def r(self):
+                                   maxiter=100, npart_min=10):
-        """
-        Radial separation of particles from the centre of the object.
-
-        Returns
-        -------
-        r : 1-dimensional array of shape `(n_particles, )`.
-        """
-        return self._get_r(self.pos)
-
-    @staticmethod
-    @jit(nopython=True)
-    def _get_r(pos):
-        return (pos[:, 0]**2 + pos[:, 1]**2 + pos[:, 2]**2)**0.5
-
-    def cmass(self, rmax, rmin):
-        """
-        Calculate Cartesian position components of the object's centre of mass.
-        Note that this is already in a frame centered at the clump's potential
-        minimum, so its distance from origin indicates the separation of the
-        centre of mass and potential minimum.
-
-        Parameters
-        ----------
-        rmax : float
-            Maximum radius for particles to be included in the calculation.
-        rmin : float
-            Minimum radius for particles to be included in the calculation.
-
-        Returns
-        -------
-        cm : 1-dimensional array of shape `(3, )`
-        """
-        r = self.r
-        mask = (r >= rmin) & (r <= rmax)
-        return numpy.average(self.pos[mask], axis=0, weights=self["M"][mask])
-
-    def angular_momentum(self, rmax, rmin=0):
-        """
-        Calculate angular momentum in the box coordinates.
-
-        Parameters
-        ----------
-        rmax : float
-            Maximum radius for particles to be included in the calculation.
-        rmin : float
-            Minimum radius for particles to be included in the calculation.
-
-        Returns
-        -------
-        J : 1-dimensional array or shape `(3, )`
-        """
-        r = self.r
-        mask = (r >= rmin) & (r <= rmax)
-        pos = self.pos[mask] - self.cmass(rmax, rmin)
-        # Velocitities in the object CM frame
-        vel = self.vel[mask]
-        vel -= numpy.average(self.vel[mask], axis=0, weights=self["M"][mask])
-        return numpy.einsum("i,ij->j", self["M"][mask], numpy.cross(pos, vel))
-
-    def enclosed_mass(self, rmax, rmin=0):
-        """
-        Sum of particle masses between two radii.
-
-        Parameters
-        ----------
-        rmax : float
-            Maximum radial distance.
-        rmin : float, optional
-            Minimum radial distance.
-
-        Returns
-        -------
-        enclosed_mass : float
-        """
-        r = self.r
-        return numpy.sum(self["M"][(r >= rmin) & (r <= rmax)])
-
-    def lambda_bullock(self, radmax, npart_min=10):
         r"""
-        Bullock spin, see Eq. 5 in [1], in a radius of `radius`, which should
+        Calculate spherical overdensity mass and radius via the iterative
-        define to some overdensity radius.
+        shrinking sphere method.
 
         Parameters
         ----------
-        radmax : float
+        delta_mult : int or float
-            Radius in which to calculate the spin.
+            Overdensity multiple.
-        npart_min : int
+        kind : str, optional
+            Either `crit` or `matter`, for critical or matter overdensity
+        tol : float, optional
+            Tolerance for the change in the center of mass or radius.
+        maxiter : int, optional
+            Maximum number of iterations.
+        npart_min : int, optional
+            Minimum number of enclosed particles to reset the iterator.
+
+        Returns
+        -------
+        mass :  float
+            The requested spherical overdensity mass.
+        rad : float
+            The radius of the sphere enclosing the requested overdensity.
+        cm : 1-dimensional array of shape `(3, )`
+            The center of mass of the sphere enclosing the requested
+            overdensity.
+        """
+        assert kind in ["crit", "matter"]
+        rho = delta_mult * self.box.box_rhoc
+        if kind == "matter":
+            rho *= self.box.box_Om
+        pos = self.pos
+        mass = self["M"]
+
+        # Initial guesses
+        init_cm = center_of_mass(pos, mass, boxsize=1)
+        init_rad = mass_to_radius(numpy.sum(mass), rho) * 1.5
+
+        rad = init_rad
+        cm = numpy.copy(init_cm)
+
+        success = False
+        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.
+            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)
+                rad = init_rad * (0.75 + numpy.random.rand())
+                dist = periodic_distance(pos, cm, boxsize=1)
+                within_rad = dist <= rad
+
+            # Calculate the enclosed mass for the current CM and radius.
+            enclosed_mass = numpy.sum(mass[within_rad])
+
+            # Calculate the new CM and radius from this mass.
+            new_rad = mass_to_radius(enclosed_mass, rho)
+            new_cm = center_of_mass(pos[within_rad], mass[within_rad],
+                                    boxsize=1)
+
+            # Update the CM and radius
+            prev_cm, cm = cm, new_cm
+            prev_rad, rad = rad, new_rad
+
+            # 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)
+
+        return enclosed_mass, rad, cm
+
+    def angular_momentum(self, ref, rad, npart_min=10):
+        """
+        Calculate angular momentum around a reference point using all particles
+        within a radius. The angular momentum is returned in box units.
+
+        Parameters
+        ----------
+        ref : 1-dimensional array of shape `(3, )`
+            Reference point.
+        rad : float
+            Radius around the reference point.
+        npart_min : int, optional
             Minimum number of enclosed particles for a radius to be
             considered trustworthy.
 
+        Returns
+        -------
+        angmom : 1-dimensional array or shape `(3, )`
+        """
+        pos = self.pos
+        mask = periodic_distance(pos, ref, boxsize=1) < rad
+        if numpy.sum(mask) < npart_min:
+            return numpy.full(3, numpy.nan)
+
+        mass = self["M"][mask]
+        pos = pos[mask]
+        vel = self.vel[mask]
+        # Velocitities in the object CM frame
+        vel -= numpy.average(vel, axis=0, weights=mass)
+        return numpy.sum(mass[:, numpy.newaxis] * numpy.cross(pos, vel),
+                         axis=0)
+
+    def lambda_bullock(self, ref, rad):
+        r"""
+        Bullock spin, see Eq. 5 in [1], in a given radius around a reference
+        point.
+
+        Parameters
+        ----------
+        ref : 1-dimensional array of shape `(3, )`
+            Reference point.
+        rad : float
+            Radius around the reference point.
+
         Returns
         -------
         lambda_bullock : float
@@ -204,61 +219,64 @@ class BaseStructure(ABC):
         Bullock, J. S.;  Dekel, A.;  Kolatt, T. S.;  Kravtsov, A. V.;
         Klypin, A. A.;  Porciani, C.;  Primack, J. R.
         """
-        mask = self.r <= radmax
+        pos = self.pos
-        if numpy.sum(mask) < npart_min:
+        mask = periodic_distance(pos, ref, boxsize=1) < rad
-            return numpy.nan
+        mass = numpy.sum(self["M"][mask])
-        mass = self.enclosed_mass(radmax)
+        circvel = numpy.sqrt(self.box.box_G * mass / rad)
-        circvel = numpy.sqrt(self.box.box_G * mass / radmax)
+        angmom_norm = numpy.linalg.norm(self.angular_momentum(ref, rad))
-        angmom_norm = numpy.linalg.norm(self.angular_momentum(radmax))
+        return angmom_norm / (numpy.sqrt(2) * mass * circvel * rad)
-        return angmom_norm / (numpy.sqrt(2) * mass * circvel * radmax)
 
-    def spherical_overdensity_mass(self, delta_mult, npart_min=10,
+    def nfw_concentration(self, ref, rad, conc_min=1e-3, npart_min=10):
-                                   kind="crit"):
+        """
-        r"""
+        Calculate the NFW concentration parameter in a given radius around a
-        Calculate spherical overdensity mass and radius. The mass is defined as
+        reference point.
-        the enclosed mass within an outermost radius where the mean enclosed
-        spherical density reaches a multiple of the critical density `delta`
-        (times the matter density if `kind` is `matter`).
 
         Parameters
         ----------
-        delta_mult : list of int or float
+        ref : 1-dimensional array of shape `(3, )`
-            Overdensity multiple.
+            Reference point.
-        npart_min : int
+        rad : float
-            Minimum number of enclosed particles for a radius to be
+            Radius around the reference point.
-            considered trustworthy.
+        conc_min : float
-        kind : str
+            Minimum concentration limit.
-            Either `crit` or `matter`, for critical or matter overdensity
+        npart_min : int, optional
+            Minimum number of enclosed particles to calculate the
+            concentration.
 
         Returns
         -------
-        rx : float
+        conc : float
-            Radius where the enclosed density reaches required value.
-        mx :  float
-            Corresponding spherical enclosed mass.
         """
-        # Quick check of inputs
+        pos = self.pos
-        assert kind in ["crit", "matter"]
+        dist = periodic_distance(pos, ref, boxsize=1)
+        mask = dist < rad
+        if numpy.sum(mask) < npart_min:
+            return numpy.nan
 
-        # We first sort the particles in an increasing separation
+        dist = dist[mask]
-        rs = self.r
+        weight = self["M"][mask]
-        order = numpy.argsort(rs)
+        weight /= numpy.mean(weight)
-        rs = rs[order]
 
-        cmass = numpy.cumsum(self["M"][order])  # Cumulative mass
+        # We do the minimization in log space
-        # We calculate the enclosed volume and indices where it is above target
+        def negll_nfw_concentration(log_c, xs, weight):
-        vol = 4 * numpy.pi / 3 * rs**3
+            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))
+            return -ll
 
-        target_density = delta_mult * self.box.box_rhoc
+        res = minimize(negll_nfw_concentration, x0=1.5,
-        if kind == "matter":
+                       args=(dist / rad, weight, ), method='Nelder-Mead',
-            target_density *= self.box.cosmo.Om0
+                       bounds=((numpy.log10(conc_min), 5),))
-        ks = numpy.where(cmass > target_density * vol)[0]
+
-        if ks.size == 0:  # Never above the threshold?
+        if not res.success:
-            return numpy.nan, numpy.nan
+            return numpy.nan
-        k = numpy.max(ks)
+
-        if k < npart_min:  # Too few particles?
+        res = 10**res["x"][0]
-            return numpy.nan, numpy.nan
+        if res < conc_min or numpy.isclose(res, conc_min):
-        return rs[k], cmass[k]
+            return numpy.nan
+
+        return res
 
     def __getitem__(self, key):
         keys = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']
@@ -270,44 +288,23 @@ class BaseStructure(ABC):
         return self.particles.shape[0]
 
 
-class Clump(BaseStructure):
-    """
-    Clump object to handle operations on its particles.
-
-    Parameters
-    ----------
-    particles : structured array
-        Particle array. Must contain `['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']`.
-    info : structured array
-        Array containing information from the clump finder.
-    box : :py:class:`csiborgtools.read.CSiBORGBox`
-        Box units object.
-    """
-
-    def __init__(self, particles, info, box):
-        self.particles = particles
-        self.info = info
-        self.box = box
-
-
 class Halo(BaseStructure):
     """
-    Ultimate halo object to handle operations on its particles, i.e. the summed
+    Halo object to handle operations on its particles.
-    particles halo.
 
     Parameters
     ----------
     particles : structured array
         Particle array. Must contain `['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']`.
     info : structured array
-        Array containing information from the clump finder.
+        Array containing information from the halo finder.
     box : :py:class:`csiborgtools.read.CSiBORGBox`
         Box units object.
     """
 
-    def __init__(self, particles, info, box):
+    def __init__(self, particles, box):
         self.particles = particles
-        self.info = info
+        # self.info = info
         self.box = box
 
 
@@ -315,28 +312,102 @@ class Halo(BaseStructure):
 #                       Other, supplementary functions                        #
 ###############################################################################
 
-@jit(nopython=True)
+
-def dist_centmass(clump):
+def center_of_mass(points, mass, boxsize):
     """
-    Calculate the clump (or halo) particles' distance from the centre of mass.
+    Calculate the center of mass of a halo, while assuming for periodic
+    boundary conditions of a cubical box. Assuming that particle positions are
+    in `[0, boxsize)` range.
 
     Parameters
     ----------
-    clump : 2-dimensional array of shape (n_particles, 7)
+    points : 2-dimensional array of shape (n_particles, 3)
-        Particle array. The first four columns must be `x`, `y`, `z` and `M`.
+        Particle position array.
+    mass : 1-dimensional array of shape `(n_particles, )`
+        Particle mass array.
+    boxsize : float
+        Box size in the same units as `parts` coordinates.
 
     Returns
     -------
-    dist : 1-dimensional array of shape `(n_particles, )`
+    cm : 1-dimensional array of shape `(3, )`
-        Particle distance from the centre of mass.
-    cm : len-3 list
-        Center of mass coordinates.
     """
-    mass = clump[:, 3]
+    # Convert positions to unit circle coordinates in the complex plane
-    x, y, z = clump[:, 0], clump[:, 1], clump[:, 2]
+    pos = numpy.exp(2j * numpy.pi * points / boxsize)
-    cmx, cmy, cmz = [numpy.average(xi, weights=mass) for xi in (x, y, z)]
+    # Compute weighted average of these coordinates, convert it back to
-    dist = ((x - cmx)**2 + (y - cmy)**2 + (z - cmz)**2)**0.5
+    # box coordinates and fix any negative positions due to angle calculations.
-    return dist, [cmx, cmy, cmz]
+    cm = numpy.angle(numpy.average(pos, axis=0, weights=mass))
+    cm *= boxsize / (2 * numpy.pi)
+    cm[cm < 0] += boxsize
+    return cm
+
+
+def periodic_distance(points, reference, boxsize):
+    """
+    Compute the periodic distance between multiple points and a reference
+    point.
+
+    Parameters
+    ----------
+    points : 2-dimensional array of shape `(n_points, 3)`
+        Points to calculate the distance from the reference point.
+    reference : 1-dimensional array of shape `(3, )`
+        Reference point.
+    boxsize : float
+        Box size.
+
+    Returns
+    -------
+    dist : 1-dimensional array of shape `(n_points, )`
+    """
+    delta = numpy.abs(points - reference)
+    delta = numpy.where(delta > boxsize / 2, boxsize - delta, delta)
+    return numpy.linalg.norm(delta, axis=1)
+
+
+def shift_to_center_of_box(points, cm, boxsize, set_cm_to_zero=False):
+    """
+    Shift the positions such that the CM is at the center of the box, while
+    accounting for periodic boundary conditions.
+
+    Parameters
+    ----------
+    points : 2-dimensional array of shape `(n_points, 3)`
+        Points to shift.
+    cm : 1-dimensional array of shape `(3, )`
+        Center of mass.
+    boxsize : float
+        Box size.
+    set_cm_to_zero : bool, optional
+        If `True`, set the CM to zero.
+
+    Returns
+    -------
+    shifted_positions : 2-dimensional array of shape `(n_points, 3)`
+    """
+    pos = (points + (boxsize / 2 - cm)) % boxsize
+    if set_cm_to_zero:
+        pos -= boxsize / 2
+    return pos
+
+
+def mass_to_radius(mass, rho):
+    """
+    Compute the radius of a sphere with a given mass and density.
+
+    Parameters
+    ----------
+    mass : float
+        Mass of the sphere.
+    rho : float
+        Density of the sphere.
+
+    Returns
+    -------
+    rad : float
+        Radius of the sphere.
+    """
+    return ((3 * mass) / (4 * numpy.pi * rho))**(1./3)
 
 
 @jit(nopython=True)

csiborgtools/fits/haloprofile.py

@@ -1,377 +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.
-"""Halo profiles functions and posteriors."""
-import numpy
-from scipy.optimize import minimize_scalar
-from scipy.stats import uniform
-
-from .halo import Clump
-
-
-class NFWProfile:
-    r"""
-    The Navarro-Frenk-White (NFW) density profile.
-
-    .. math::
-        \rho(r) = \frac{\rho_0}{x(1 + x)^2},
-
-    :math:`x = r / R_s` and its free paramaters are :math:`R_s, \rho_0`: scale
-    radius and NFW density parameter.
-    """
-
-    @staticmethod
-    def profile(r, Rs, rho0):
-        """
-        Evaluate the halo profile at `r`.
-
-        Parameters
-        ----------
-        r : 1-dimensional array
-            Radial distance.
-        Rs : float
-            Scale radius.
-        rho0 : float
-            NFW density parameter.
-
-        Returns
-        -------
-        density : 1-dimensional array
-        """
-        x = r / Rs
-        return rho0 / (x * (1 + x) ** 2)
-
-    @staticmethod
-    def _logprofile(r, Rs, rho0):
-        """Natural logarithm of `NFWPprofile.profile(...)`."""
-        x = r / Rs
-        return numpy.log(rho0) - numpy.log(x) - 2 * numpy.log(1 + x)
-
-    @staticmethod
-    def mass(r, Rs, rho0):
-        r"""
-        Calculate the enclosed mass of a NFW profile in radius `r`.
-
-        Parameters
-        ----------
-        r : 1-dimensional array
-            Radial distance.
-        Rs : float
-            Scale radius.
-        rho0 : float
-            NFW density parameter.
-
-        Returns
-        -------
-        M : 1-dimensional array
-            The enclosed mass.
-        """
-        x = r / Rs
-        out = numpy.log(1 + x) - x / (1 + x)
-        return 4 * numpy.pi * rho0 * Rs**3 * out
-
-    def bounded_mass(self, rmin, rmax, Rs, rho0):
-        r"""
-        Calculate the enclosed mass between `rmin` and `rmax`.
-
-        Parameters
-        ----------
-        rmin : float
-            Minimum radius.
-        rmax : float
-            Maximum radius.
-        Rs : float
-            Scale radius :math:`R_s`.
-        rho0 : float
-            NFW density parameter.
-
-        Returns
-        -------
-        M : float
-            Enclosed mass within the radial range.
-        """
-        return self.mass(rmax, Rs, rho0) - self.mass(rmin, Rs, rho0)
-
-    def pdf(self, r, Rs, rmin, rmax):
-        r"""
-        Calculate the radial PDF of the NFW profile, defined below.
-
-        .. math::
-            \frac{4\pi r^2 \rho(r)} {M(r_\min, r_\max)},
-
-        where :math:`M(r_\min, r_\max)` is the enclosed mass between
-        :math:`r_\min` and :math:`r_\max'. Note that the dependance on
-        :math:`\rho_0` is cancelled and must be accounted for in the
-        normalisation term to match the total mass.
-
-        Parameters
-        ----------
-        r : 1-dimensional array
-            Radial distance.
-        Rs : float
-            Scale radius.
-        rmin : float
-            Minimum radius to evaluate the PDF (denominator term).
-        rmax : float
-            Maximum radius to evaluate the PDF (denominator term).
-
-        Returns
-        -------
-        pdf : 1-dimensional array
-        """
-        norm = self.bounded_enclosed_mass(rmin, rmax, Rs, 1)
-        return 4 * numpy.pi * r**2 * self.profile(r, Rs, 1) / norm
-
-    def rvs(self, rmin, rmax, Rs, size=1):
-        """
-        Generate random samples from the NFW profile via rejection sampling.
-
-        Parameters
-        ----------
-        rmin : float
-            Minimum radius.
-        rmax : float
-            Maximum radius.
-        Rs : float
-            Scale radius.
-        size : int, optional
-            Number of samples to generate. By default 1.
-
-        Returns
-        -------
-        samples : float or 1-dimensional array
-            Samples following the NFW profile.
-        """
-        gen = uniform(rmin, rmax - rmin)
-        samples = numpy.full(size, numpy.nan)
-        for i in range(size):
-            while True:
-                r = gen.rvs()
-                if self.pdf(r, Rs, rmin, rmax) > numpy.random.rand():
-                    samples[i] = r
-                    break
-
-        if size == 1:
-            return samples[0]
-        return samples
-
-
-class NFWPosterior(NFWProfile):
-    r"""
-    Posterior for fitting the NFW profile in the range specified by the
-    closest particle and the :math:`r_{200c}` radius, calculated as below.
-
-    .. math::
-        \frac{4\pi r^2 \rho(r)} {M(r_{\min} r_{200c})} \frac{m}{M / N},
-
-    where :math:`M(r_{\min} r_{200c}))` is the NFW enclosed mass between the
-    closest particle and the :math:`r_{200c}` radius, :math:`m` is the particle
-    mass, :math:`M` is the sum of the particle masses and :math:`N` is the
-    number of particles. Calculated only using particles within the
-    above-mentioned range.
-
-    Paramaters
-    ----------
-    clump : `Clump`
-        Clump object containing the particles and clump information.
-    """
-
-    @property
-    def clump(self):
-        """
-        Clump object containig all particles, i.e. ones beyond :math:`R_{200c}`
-        as well.
-
-        Returns
-        -------
-        clump : `Clump`
-        """
-        return self._clump
-
-    @clump.setter
-    def clump(self, clump):
-        assert isinstance(clump, Clump)
-        self._clump = clump
-
-    def rho0_from_Rs(self, Rs, rmin, rmax, mass):
-        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.
-
-        Parameters
-        ----------
-        Rs : float
-            Logarithmic scale factor in units matching the coordinates.
-        rmin : float
-            Minimum radial distance of particles used to fit the profile.
-        rmax : float
-            Maximum radial distance of particles used to fit the profile.
-        mass : float
-            Mass enclosed within the radius used to fit the NFW profile.
-
-        Returns
-        -------
-        rho0: float
-        """
-        return mass / self.bounded_mass(rmin, rmax, Rs, 1)
-
-    def initlogRs(self, r, rmin, rmax, binsguess=10):
-        r"""
-        Calculate the most often occuring value of :math:`r` used as initial
-        guess of :math:`R_{\rm s}` since :math:`r^2 \rho(r)` peaks at
-        :math:`r = R_{\rm s}`.
-
-        Parameters
-        ----------
-        r : 1-dimensional array
-            Radial distance of particles used to fit the profile.
-        rmin : float
-            Minimum radial distance of particles used to fit the profile.
-        rmax : float
-            Maximum radial distance of particles used to fit the profile.
-        binsguess : int
-            Number of bins to initially guess :math:`R_{\rm s}`.
-
-        Returns
-        -------
-        initlogRs : float
-        """
-        bins = numpy.linspace(rmin, rmax, binsguess)
-        counts, edges = numpy.histogram(r, bins)
-        return numpy.log10(edges[numpy.argmax(counts)])
-
-    def logprior(self, logRs, rmin, rmax):
-        r"""
-        Logarithmic uniform prior on :math:`\log R_{\rm s}`. Unnormalised but
-        that does not matter.
-
-        Parameters
-        ----------
-        logRs : float
-            Logarithmic scale factor.
-        rmin : float
-            Minimum radial distance of particles used to fit the profile.
-        rmax : float
-            Maximum radial distance of particles used to fit the profile.
-
-        Returns
-        -------
-        lp : float
-        """
-        if not rmin < 10**logRs < rmax:
-            return -numpy.infty
-        return 0.0
-
-    def loglikelihood(self, logRs, r, rmin, rmax, npart):
-        """
-        Logarithmic likelihood.
-
-        Parameters
-        ----------
-        r : 1-dimensional array
-            Radial distance of particles used to fit the profile.
-        logRs : float
-            Logarithmic scale factor in units matching the coordinates.
-        rmin : float
-            Minimum radial distance of particles used to fit the profile.
-        rmax : float
-            Maximum radial distance of particles used to fit the profile.
-        npart : int
-            Number of particles used to fit the profile.
-
-        Returns
-        -------
-        ll : float
-        """
-        Rs = 10**logRs
-        mnfw = self.bounded_mass(rmin, rmax, Rs, 1)
-        return numpy.sum(self._logprofile(r, Rs, 1)) - npart * numpy.log(mnfw)
-
-    def __call__(self, logRs, r, rmin, rmax, npart):
-        """
-        Logarithmic posterior. Sum of the logarithmic prior and likelihood.
-
-        Parameters
-        ----------
-        logRs : float
-            Logarithmic scale factor in units matching the coordinates.
-        r : 1-dimensional array
-            Radial distance of particles used to fit the profile.
-        rmin : float
-            Minimum radial distance of particles used to fit the profile.
-        rmax : float
-            Maximum radial distance of particles used to fit the profile.
-        npart : int
-            Number of particles used to fit the profile.
-
-        Returns
-        -------
-        lpost : float
-        """
-        lp = self.logprior(logRs, rmin, rmax)
-        if not numpy.isfinite(lp):
-            return -numpy.infty
-        return self.loglikelihood(logRs, r, rmin, rmax, npart) + lp
-
-    def fit(self, clump, eps=1e-4):
-        r"""
-        Fit the NFW profile. If the fit is not converged returns NaNs.
-
-        Checks whether :math:`log r_{\rm max} / R_{\rm s} > \epsilon`,
-        to ensure that the scale radius is not too close to the boundary which
-        occurs if the fit fails.
-
-        Parameters
-        ----------
-        clump : :py:class:`csiborgtools.fits.Clump`
-            Clump being fitted.
-        eps : float
-            Tolerance to ensure we are sufficiently far from math:`R_{200c}`.
-
-        Returns
-        -------
-        Rs: float
-            Best fit scale radius.
-        rho0: float
-            Best fit NFW central density.
-        """
-        assert isinstance(clump, Clump)
-        r = clump.r
-        rmin = numpy.min(r[r > 0])  # First particle that is not at r = 0
-        rmax, mtot = clump.spherical_overdensity_mass(200)
-        mask = (rmin <= r) & (r <= rmax)
-        npart = numpy.sum(mask)
-        r = r[mask]
-
-        def loss(logRs):
-            return -self(logRs, r, rmin, rmax, npart)
-
-        # Define optimisation boundaries. Check whether they are finite and
-        # that rmax > rmin. If not, then return NaN.
-        bounds = (numpy.log10(rmin), numpy.log10(rmax))
-        if not (numpy.all(numpy.isfinite(bounds)) and bounds[0] < bounds[1]):
-            return numpy.nan, numpy.nan
-
-        res = minimize_scalar(loss, bounds=bounds, method="bounded")
-        # Check whether the fit converged to radius sufficienly far from `rmax`
-        # and that its a success. Otherwise return NaNs.
-        if numpy.log10(rmax) - res.x < eps:
-            res.success = False
-        if not res.success:
-            return numpy.nan, numpy.nan
-        # Additionally we also wish to calculate the central density from the
-        # mass (integral) constraint.
-        rho0 = self.rho0_from_Rs(10**res.x, rmin, rmax, mtot)
-        return 10**res.x, rho0

csiborgtools/read/box_units.py

@@ -15,7 +15,7 @@
 """
 Simulation box unit transformations.
 """
-from abc import ABC, abstractproperty
+from abc import ABC, abstractmethod, abstractproperty
 
 import numpy
 from astropy import constants, units
@@ -24,7 +24,7 @@ from astropy.cosmology import LambdaCDM
 from .readsim import ParticleReader
 
 # Map of CSiBORG unit conversions
-CONV_NAME = {
+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"],
@@ -104,6 +104,35 @@ class BaseBox(ABC):
         """
         pass
 
+    @abstractmethod
+    def convert_from_box(self, data, names):
+        r"""
+        Convert columns named `names` in array `data` from box units to
+        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
+            Input array.
+        names : list of str
+            Columns to be converted.
+
+        Returns
+        -------
+        data : structured array
+            Input array with converted columns.
+        """
+        pass
+
 
 ###############################################################################
 #                              CSiBORG box                                    #
@@ -340,31 +369,6 @@ class CSiBORGBox(BaseBox):
                 / (units.Mpc.to(units.cm)) ** 3)
 
     def convert_from_box(self, data, names):
-        r"""
-        Convert columns named `names` in array `data` from box units to
-        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
-            Input array.
-        names : list of str
-            Columns to be converted.
-
-        Returns
-        -------
-        data : structured array
-            Input array with converted columns.
-        """
         names = [names] if isinstance(names, str) else names
         transforms = {"length": self.box2mpc,
                       "mass": self.box2solarmass,
@@ -372,13 +376,12 @@ class CSiBORGBox(BaseBox):
                       "density": self.box2dens}
 
         for name in names:
-            # Check that the name is even in the array
             if name not in data.dtype.names:
-                raise ValueError(f"Name `{name}` not in `data` array.")
+                continue
 
             # Convert
             found = False
-            for unittype, suppnames in CONV_NAME.items():
+            for unittype, suppnames in CSIBORG_CONV_NAME.items():
                 if name in suppnames:
                     data[name] = transforms[unittype](data[name])
                     found = True
@@ -389,7 +392,7 @@ class CSiBORGBox(BaseBox):
                     f"Conversion of `{name}` is not defined.")
 
             # Center at the observer
-            if name in ["peak_x", "peak_y", "peak_z", "x0", "y0", "z0"]:
+            if name in ["x0", "y0", "z0"]:
                 data[name] -= transforms["length"](0.5)
 
         return data
@@ -427,3 +430,6 @@ class QuijoteBox(BaseBox):
     @property
     def boxsize(self):
         return 1000. / (self._cosmo.H0.value / 100)
+
+    def convert_from_box(self, data, names):
+        raise NotImplementedError("Conversion not implemented for Quijote.")

scripts/cluster_tpcf_auto.py

@@ -60,9 +60,9 @@ if __name__ == "__main__":
     parser.add_argument("--nsims", type=int, nargs="+", default=None,
                         help="Indices of simulations to cross. If `-1` processes all simulations.")  # noqa
     parser.add_argument("--Rmax", type=float, default=155/0.705,
-                        help="High-resolution region radius")  # noqa
+                        help="High-resolution region radius.")
     parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
-                        default=False)
+                        default=False, help="Verbosity flag.")
     args = parser.parse_args()
 
     with open("./cluster_tpcf_auto.yml", "r") as file:
@@ -79,8 +79,4 @@ if __name__ == "__main__":
         return do_auto(args, config, cats, nsim, paths)
 
     nsims = list(cats.keys())
-    work_delegation(do_work, nsims, comm, master_verbose=args.verbose)
+    work_delegation(do_work, nsims, comm)
-
-    comm.Barrier()
-    if comm.Get_rank() == 0:
-        print(f"{datetime.now()}: all finished. Quitting.")

scripts/fit_halos.py

@@ -13,14 +13,15 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 """
-A script to fit FoF halos (concentration, ...). The particle array of each
+A script to fit FoF halos (concentration, ...). The CSiBORG particle array of
-CSiBORG realisation must have been processed in advance by `pre_dumppart.py`.
+each realisation must have been processed in advance by `pre_dumppart.py`.
+Quijote is not supported yet
 """
 from argparse import ArgumentParser
-from datetime import datetime
 
 import numpy
 from mpi4py import MPI
+from taskmaster import work_delegation
 from tqdm import trange
 
 from utils import get_nsims
@@ -33,72 +34,67 @@ except ModuleNotFoundError:
     sys.path.append("../")
     import csiborgtools
 
-# Get MPI things
-comm = MPI.COMM_WORLD
-rank = comm.Get_rank()
-nproc = comm.Get_size()
-verbose = nproc == 1
 
-parser = ArgumentParser()
+def fit_halo(particles, box):
-parser.add_argument("--nsims", type=int, nargs="+", default=None,
+    """
-                    help="IC realisations. If `-1` processes all simulations.")
+    Fit a single halo from the particle array.
-args = parser.parse_args()
-paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-partreader = csiborgtools.read.ParticleReader(paths)
-nfwpost = csiborgtools.fits.NFWPosterior()
-nsims = get_nsims(args, paths)
 
-cols_collect = [
+    Parameters
-    ("index", numpy.int32),
+    ----------
-    ("npart", numpy.int32),
+    particles : 2-dimensional array of shape `(n_particles, 3)`
-    ("totpartmass", numpy.float32),
+        Particle array. The columns must be `x`, `y`, `z`, `vx`, `vy`, `vz`,
-    ("vx", numpy.float32),
+        `M`.
-    ("vy", numpy.float32),
+    box : object derived from :py:class`csiborgtools.read.BaseBox`
-    ("vz", numpy.float32),
+        Box object.
-    ("conc", numpy.float32),
-    ("rho0", numpy.float32),
-    ("r200c", numpy.float32),
-    ("r500c", numpy.float32),
-    ("m200c", numpy.float32),
-    ("m500c", numpy.float32),
-    ("lambda200c", numpy.float32),
-    ("r200m", numpy.float32),
-    ("m200m", numpy.float32),
-    ("r500m", numpy.float32),
-    ("m500m", numpy.float32),
-    ]
 
-
+    Returns
-def fit_halo(particles, clump_info, box):
+    -------
-    obj = csiborgtools.fits.Clump(particles, clump_info, box)
+    out : dict
+    """
+    halo = csiborgtools.fits.Halo(particles, box)
 
     out = {}
-    out["npart"] = len(obj)
+    out["npart"] = len(halo)
-    out["totpartmass"] = numpy.sum(obj["M"])
+    out["totpartmass"] = numpy.sum(halo["M"])
     for i, v in enumerate(["vx", "vy", "vz"]):
-        out[v] = numpy.average(obj.vel[:, i], weights=obj["M"])
+        out[v] = numpy.average(halo.vel[:, i], weights=halo["M"])
-    # Overdensity masses
+
-    for n in [200, 500]:
+    m200c, r200c, cm = halo.spherical_overdensity_mass(200, kind="crit",
-        out[f"r{n}c"], out[f"m{n}c"] = obj.spherical_overdensity_mass(
+                                                       maxiter=100)
-            n, kind="crit", npart_min=10)
+    out["m200c"] = m200c
-        out[f"r{n}m"], out[f"m{n}m"] = obj.spherical_overdensity_mass(
+    out["r200c"] = r200c
-            n, kind="matter", npart_min=10)
+    out["lambda200c"] = halo.lambda_bullock(cm, r200c)
-    # NFW fit
+    out["conc"] = halo.nfw_concentration(cm, r200c)
-    if out["npart"] > 10 and numpy.isfinite(out["r200c"]):
-        Rs, rho0 = nfwpost.fit(obj)
-        out["conc"] = out["r200c"] / Rs
-        out["rho0"] = rho0
-    # Spin within R200c
-    if numpy.isfinite(out["r200c"]):
-        out["lambda200c"] = obj.lambda_bullock(out["r200c"])
     return out
 
 
-# We MPI loop over all simulations.
+def _main(nsim, simname, verbose):
-jobs = csiborgtools.fits.split_jobs(len(nsims), nproc)[rank]
+    """
-for nsim in [nsims[i] for i in jobs]:
+    Fit the FoF halos.
-    print(f"{datetime.now()}: rank {rank} calculating simulation `{nsim}`.",
+
-          flush=True)
+    Parameters
+    ----------
+    nsim : int
+        IC realisation index.
+    simname : str
+        Simulation name.
+    verbose : bool
+        Verbosity flag.
+    """
+    if simname == "quijote":
+        raise NotImplementedError("Quijote not implemented yet.")
+
+    cols = [("index", numpy.int32),
+            ("npart", numpy.int32),
+            ("totpartmass", numpy.float32),
+            ("vx", numpy.float32),
+            ("vy", numpy.float32),
+            ("vz", numpy.float32),
+            ("conc", numpy.float32),
+            ("r200c", numpy.float32),
+            ("m200c", numpy.float32),
+            ("lambda200c", numpy.float32),]
+
     nsnap = max(paths.get_snapshots(nsim))
     box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
 
@@ -106,29 +102,44 @@ for nsim in [nsims[i] for i in jobs]:
     f = csiborgtools.read.read_h5(paths.particles(nsim))
     particles = f["particles"]
     halo_map = f["halomap"]
-    hid2map = {clid: i for i, clid in enumerate(halo_map[:, 0])}
+    hid2map = {hid: i for i, hid in enumerate(halo_map[:, 0])}
     cat = csiborgtools.read.CSiBORGHaloCatalogue(
         nsim, paths, with_lagpatch=False, load_initial=False, rawdata=True,
         load_fitted=False)
-    # Even if we are calculating parent halo this index runs over all clumps.
+
-    out = csiborgtools.read.cols_to_structured(len(cat), cols_collect)
+    out = csiborgtools.read.cols_to_structured(len(cat), cols)
-    indxs = cat["index"]
     for i in trange(len(cat)) if verbose else range(len(cat)):
         hid = cat["index"][i]
         out["index"][i] = hid
-
         part = csiborgtools.read.load_halo_particles(hid, particles, halo_map,
                                                      hid2map)
-        # We fit the particles if there are any. If not we assign the index,
+        # Skip if no particles.
-        # otherwise it would be NaN converted to integers (-2147483648) and
-        # yield an error further down.
         if part is None:
             continue
 
-        _out = fit_halo(part, cat[i], box)
+        _out = fit_halo(part, box)
         for key in _out.keys():
             out[key][i] = _out[key]
 
     fout = paths.structfit(nsnap, nsim)
-    print(f"Saving to `{fout}`.", flush=True)
+    if verbose:
+        print(f"Saving to `{fout}`.", flush=True)
     numpy.save(fout, out)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--simname", type=str, default="csiborg",
+                        choices=["csiborg", "quijote", "quijote_full"],
+                        help="Simulation name")
+    parser.add_argument("--nsims", type=int, nargs="+", default=None,
+                        help="IC realisations. If `-1` processes all.")
+    args = parser.parse_args()
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(args, paths)
+
+    def main(nsim):
+        _main(nsim, args.simname, MPI.COMM_WORLD.Get_size() == 1)
+
+    work_delegation(main, nsims, MPI.COMM_WORLD)

scripts/fit_hmf.py

@@ -94,17 +94,13 @@ if __name__ == "__main__":
     parser.add_argument("--nsims", type=int, nargs="+", default=None,
                         help="Indices of simulations to cross. If `-1` processes all simulations.")  # noqa
     parser.add_argument("--Rmax", type=float, default=155/0.705,
-                        help="High-resolution region radius")
+                        help="High-resolution region radius. Ignored for `quijote_full`.")  # noqa
     parser.add_argument("--bw", type=float, default=0.2,
-                        help="Bin width in dex")
+                        help="Bin width in dex.")
     parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
-                        default=False)
+                        default=False, help="Verbosity flag.")
-
     parser_args = parser.parse_args()
-    comm = MPI.COMM_WORLD
+
-    rank = comm.Get_rank()
-    nproc = comm.Get_size()
-    verbose = nproc == 1
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     nsims = get_nsims(parser_args, paths)
     bins = numpy.arange(11., 16., parser_args.bw, dtype=numpy.float32)
@@ -112,4 +108,4 @@ if __name__ == "__main__":
     def do_work(nsim):
         get_counts(nsim, bins, paths, parser_args)
 
-    work_delegation(do_work, nsims, comm, master_verbose=parser_args.verbose)
+    work_delegation(do_work, nsims, MPI.COMM_WORLD)

scripts/fit_init.py

@@ -22,6 +22,7 @@ from datetime import datetime
 
 import numpy
 from mpi4py import MPI
+from taskmaster import work_delegation
 from tqdm import tqdm
 
 from utils import get_nsims
@@ -35,73 +36,83 @@ except ModuleNotFoundError:
     import csiborgtools
 
 
-# Get MPI things
+def _main(nsim, simname, verbose):
-comm = MPI.COMM_WORLD
+    """
-rank = comm.Get_rank()
+    Calculate the Lagrangian halo centre of mass and Lagrangian patch size in
-nproc = comm.Get_size()
+    the initial snapshot.
-verbose = nproc == 1
 
-# Argument parser
+    Parameters
-parser = ArgumentParser()
+    ----------
-parser.add_argument("--simname", type=str, default="csiborg",
+    nsim : int
-                    choices=["csiborg", "quijote"],
+        IC realisation index.
-                    help="Simulation name")
+    simname : str
-parser.add_argument("--nsims", type=int, nargs="+", default=None,
+        Simulation name.
-                    help="IC realisations. If `-1` processes all simulations.")
+    verbose : bool
-args = parser.parse_args()
+        Verbosity flag.
-paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    """
-partreader = csiborgtools.read.ParticleReader(paths)
+    if simname == "quijote":
+        raise NotImplementedError("Quijote not implemented yet.")
 
-nsims = get_nsims(args, paths)
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-
+    cols = [("index", numpy.int32),
-cols_collect = [("index", numpy.int32),
+            ("x", numpy.float32),
-                ("x", numpy.float32),
+            ("y", numpy.float32),
-                ("y", numpy.float32),
+            ("z", numpy.float32),
-                ("z", numpy.float32),
+            ("lagpatch_size", numpy.float32),
-                ("lagpatch_size", numpy.float32),
+            ("lagpatch_ncells", numpy.int32),]
-                ("lagpatch_ncells", numpy.int32),]
-
-
-# MPI loop over simulations
-jobs = csiborgtools.fits.split_jobs(len(nsims), nproc)[rank]
-for nsim in [nsims[i] for i in jobs]:
-    nsnap = max(paths.get_snapshots(nsim))
-    overlapper = csiborgtools.match.ParticleOverlap()
-    print(f"{datetime.now()}: rank {rank} calculating simulation `{nsim}`.",
-          flush=True)
 
     parts = csiborgtools.read.read_h5(paths.initmatch(nsim, "particles"))
     parts = parts['particles']
     halo_map = csiborgtools.read.read_h5(paths.particles(nsim))
     halo_map = halo_map["halomap"]
+
     cat = csiborgtools.read.CSiBORGHaloCatalogue(
         nsim, paths, rawdata=True, load_fitted=False, load_initial=False)
     hid2map = {hid: i for i, hid in enumerate(halo_map[:, 0])}
 
-    out = csiborgtools.read.cols_to_structured(len(cat), cols_collect)
+    out = csiborgtools.read.cols_to_structured(len(cat), cols)
     for i, hid in enumerate(tqdm(cat["index"]) if verbose else cat["index"]):
         out["index"][i] = hid
         part = csiborgtools.read.load_halo_particles(hid, parts, halo_map,
                                                      hid2map)
 
-        # Skip if the halo is too small.
+        # Skip if the halo has no particles or is too small.
         if part is None or part.size < 100:
             continue
 
+        pos, mass = part[:, :3], part[:, 3]
         # Calculate the centre of mass and the Lagrangian patch size.
-        dist, cm = csiborgtools.fits.dist_centmass(part)
+        cm = csiborgtools.fits.center_of_mass(pos, mass, boxsize=1.0)
-        # We enforce a maximum patchsize of 0.075 in box coordinates.
+        distances = csiborgtools.fits.periodic_distance(pos, cm, boxsize=1.0)
-        patchsize = min(numpy.percentile(dist, 99), 0.075)
         out["x"][i], out["y"][i], out["z"][i] = cm
-        out["lagpatch_size"][i] = patchsize
+        out["lagpatch_size"][i] = numpy.percentile(distances, 99)
 
         # Calculate the number of cells with > 0 density.
-        delta = overlapper.make_delta(part[:, :3], part[:, 3], subbox=True)
+        overlapper = csiborgtools.match.ParticleOverlap()
+        delta = overlapper.make_delta(pos, mass, subbox=True)
         out["lagpatch_ncells"][i] = csiborgtools.fits.delta2ncells(delta)
 
     # Now save it
     fout = paths.initmatch(nsim, "fit")
-    print(f"{datetime.now()}: dumping fits to .. `{fout}`.",
+    if verbose:
-          flush=True)
+        print(f"{datetime.now()}: dumping fits to .. `{fout}`.", flush=True)
     with open(fout, "wb") as f:
         numpy.save(f, out)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--simname", type=str, default="csiborg",
+                        choices=["csiborg", "quijote"],
+                        help="Simulation name")
+    parser.add_argument("--nsims", type=int, nargs="+", default=None,
+                        help="IC realisations. If `-1` processes all.")
+    args = parser.parse_args()
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(args, paths)
+
+    def main(nsim):
+        _main(nsim, args.simname, MPI.COMM_WORLD.Get_size() == 1)
+
+    work_delegation(main, nsims, MPI.COMM_WORLD)

scripts/mv_fofmembership.py

@@ -146,6 +146,4 @@ if __name__ == "__main__":
 
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     nsims = get_nsims(args, paths)
-    comm = MPI.COMM_WORLD
+    work_delegation(main, nsims, MPI.COMM_WORLD)
-
-    work_delegation(main, nsims, comm)

scripts/old/pre_mmain.sh

@@ -0,0 +1,14 @@
+nthreads=102
+memory=5
+queue="cmb"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="pre_mmain.py"
+
+# pythoncm="$env $file"
+# $pythoncm
+
+
+cm="addqueue -q $queue -n $nthreads -m $memory $env $file"
+echo "Submitting:"
+echo $cm
+$cm

scripts/pre_dumppart.py

@@ -169,7 +169,7 @@ if __name__ == "__main__":
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     nsims = get_nsims(args, paths)
 
-    def _main(nsim, verbose=MPI.COMM_WORLD.nproc == 1):
+    def _main(nsim):
-        main(nsim, args.simname, verbose=verbose)
+        main(nsim, args.simname, verbose=MPI.COMM_WORLD.Get_size() == 1)
 
     work_delegation(_main, nsims, MPI.COMM_WORLD)

scripts/pre_sortinit.py

@@ -95,6 +95,6 @@ if __name__ == "__main__":
     nsims = get_nsims(args, paths)
 
     def main(nsim):
-        _main(nsim, args.simname, MPI.COMM_WORLD.size == 1)
+        _main(nsim, args.simname, MPI.COMM_WORLD.Get_size() == 1)
 
     work_delegation(main, nsims, MPI.COMM_WORLD)