Within halo work and NFW fit (#4)

* add listing of snapshots

* change distance to comoving

* ignore cp files

* rename nb

* add str to list

* add NFW profile shapes

* add fits imports

* Rename to Nsnap

* in clumps_read only select props

* make clumpid int

* expand doc

* add import

* edit readme

* distribute halos

* add profile & posterior

* add import

* add import

* add documentation

* add rvs and init guess

* update todo

* update nb

* add file

* return end index too

* change clump_ids format to int32

* skeleton of dump particle

* update nb

* add func to drop 0 clump indxs parts

* add import

* add halo dump

* switch to float32

* Update TODO

* update TODO

* add func that loads a split

* add halo object

* Rename to clump

* make post work with a clump

* add optimiser

* add Nsplits

* ignore submission scripts

* ignore .out

* add dumppath

* add job splitting

* add split halos script

* rename file

* renaem files

* rm file

* rename imports

* edit desc

* add pick clump

* add number of particles

* update TODO

* update todo

* add script

* add dumping

* change dumpdir structure

* change dumpdir

* add import

* Remove tqdm

* Increase the number of splits

* rm shuffle option

* Change to remove split

* add emojis

* fix part counts in splits

* change num of splits

* rm with particle cut

* keep splits

* fit only if 10 part and more

* add min distance

* rm warning about not set vels

* update TODO

* calculate rho0 too

* add results collection

* add import

* add func to combine splits

* update TODO

* add extract cols

* update nb

* update TODO

.gitignore

@@ -6,3 +6,8 @@ bin/
 */.ipynb_checkpoints/
 plots/*
 .vscode/settings.json
+csiborgtools/fits/_halo_profile.py
+csiborgtools/fits/_filenames.py
+csiborgtools/fits/analyse_voids_25.py
+scripts/*.sh
+scripts/*.out

README.md

@@ -1 +1,10 @@
-# CSiBORG analysis tools
+# CSiBORG analysis tools :dart:
+
+## TODO :scroll:
+- [ ] Calculate $\left\{M_{\rm vir}, R_{\rm vir}, c\right\}$ from $\left\{R_s, \rho_0, \ldots \right\}$
+- [ ] Calculate $M_{\rm 500c}$ by sphere shrinking
+- [ ] Calculate the cross-correlation in CSiBORG. Should see the scale of the constraints?
+
+
+## Open questions :bulb:
+- Get uncertainty on the fitted $R_{\rm s}$? If so get this directly from JAX.

csiborgtools/__init__.py

@@ -13,4 +13,4 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from csiborgtools import (io, match, utils, units)  # noqa
+from csiborgtools import (io, match, utils, units, fits)  # noqa

csiborgtools/fits/__init__.py

@@ -12,3 +12,8 @@
 # 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 .haloprofile import (NFWProfile, NFWPosterior)  # noqa
+from .halo  import (distribute_halos, clump_with_particles,  # noqa
+                    dump_split_particles, load_split_particles,  # noqa
+                    split_jobs, pick_single_clump, Clump)  # noqa

csiborgtools/fits/halo.py

@@ -0,0 +1,465 @@
+# 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.
+"""
+Tools for splitting the particles and a clump object.
+"""
+
+
+import numpy
+from os import remove
+from warnings import warn
+from os.path import join
+from tqdm import trange
+from ..io import nparts_to_start_ind
+
+
+def clump_with_particles(particle_clumps, clumps):
+    """
+    Count how many particles does each clump have.
+
+    Parameters
+    ----------
+    particle_clumps : 1-dimensional array
+        Array of particles' clump IDs.
+    clumps : structured array
+        The clumps array.
+
+    Returns
+    -------
+    with_particles : 1-dimensional array
+        Array of whether a clump has any particles.
+    """
+    return numpy.isin(clumps["index"], particle_clumps)
+
+
+def distribute_halos(Nsplits, clumps):
+    """
+    Evenly distribute clump indices to smaller splits. Clumps should only be
+    clumps that contain particles.
+
+    Parameters
+    ----------
+    Nsplits : int
+        Number of splits.
+    clumps : structured array
+        The clumps array.
+
+    Returns
+    -------
+    splits : 2-dimensional array
+        Array of starting and ending indices of each CPU of shape `(Njobs, 2)`.
+    """
+    # Make sure these are unique IDs
+    indxs = clumps["index"]
+    if indxs.size > numpy.unique((indxs)).size:
+        raise ValueError("`clump_indxs` constains duplicate indices.")
+    Ntotal = indxs.size
+    Njobs_per_cpu = numpy.ones(Nsplits, dtype=int) * Ntotal // Nsplits
+    # Split the remainder Ntotal % Njobs among the CPU
+    Njobs_per_cpu[:Ntotal % Nsplits] += 1
+    start = nparts_to_start_ind(Njobs_per_cpu)
+    return numpy.vstack([start, start + Njobs_per_cpu]).T
+
+
+def dump_split_particles(particles, particle_clumps, clumps, Nsplits,
+                         dumpfolder, Nsim, Nsnap, verbose=True):
+    """
+    Save the data needed for each split so that a process does not have to load
+    everything.
+
+    Parameters
+    ----------
+    particles : structured array
+        The particle array.
+    particle_clumps : 1-dimensional array
+        Array of particles' clump IDs.
+    clumps : structured array
+        The clumps array.
+    Nsplits : int
+        Number of times to split the clumps.
+    dumpfolder : str
+        Path to the folder where to dump the splits.
+    Nsim : int
+        CSiBORG simulation index.
+    Nsnap : int
+        Snapshot index.
+    verbose : bool, optional
+        Verbosity flag. By default `True`.
+
+    Returns
+    -------
+    None
+    """
+    if particles.size != particle_clumps.size:
+        raise ValueError("`particles` must correspond to `particle_clumps`.")
+    # Calculate which clumps have particles
+    with_particles = clump_with_particles(particle_clumps, clumps)
+    clumps = numpy.copy(clumps)[with_particles]
+    if verbose:
+        warn(r"There are {:.4f}% clumps that have identified particles."
+             .format(with_particles.sum() / with_particles.size * 100))
+
+    # The starting clump index of each split
+    splits = distribute_halos(Nsplits, clumps)
+    fname = join(dumpfolder, "out_{}_snap_{}_{}.npz")
+
+    iters = trange(Nsplits) if verbose else range(Nsplits)
+    tot = 0
+    for n in iters:
+        # Lower and upper array index of the clumps array
+        i, j = splits[n, :]
+        # Clump indices in this split
+        indxs = clumps["index"][i:j]
+        hmin, hmax = indxs.min(), indxs.max()
+        mask = (particle_clumps >= hmin) & (particle_clumps <= hmax)
+        # Check number of clumps
+        npart_unique = numpy.unique(particle_clumps[mask]).size
+        if indxs.size > npart_unique:
+            raise RuntimeError(
+                "Split `{}` contains more unique clumps (`{}`) than there are "
+                "unique particles' clump indices (`{}`)after removing clumps "
+                "with no particles.".format(n, indxs.size, npart_unique))
+        # Dump it!
+        tot += mask.sum()
+        fout = fname.format(Nsim, Nsnap, n)
+        numpy.savez(fout, particles[mask], particle_clumps[mask], clumps[i:j])
+
+    # There are particles whose clump ID is > 1 and have no counterpart in the
+    # clump file. Therefore can save fewer particles, depending on the cut.
+    if tot > particle_clumps.size:
+        raise RuntimeError(
+            "Num. of dumped particles `{}` is greater than the particle file "
+            "size `{}`.".format(tot, particle_clumps.size))
+
+
+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
+
+
+def load_split_particles(Nsplit, dumpfolder, Nsim, Nsnap, remove_split=False):
+    """
+    Load particles of a split saved by `dump_split_particles`.
+
+    Parameters
+    --------
+    Nsplit : int
+        Split index.
+    dumpfolder : str
+        Path to the folder where the splits were dumped.
+    Nsim : int
+        CSiBORG simulation index.
+    Nsnap : int
+        Snapshot index.
+    remove_split : bool, optional
+        Whether to remove the split file. By default `False`.
+
+    Returns
+    -------
+    particles : structured array
+        Particle array of this split.
+    clumps_indxs : 1-dimensional array
+        Array of particles' clump IDs of this split.
+    clumps : 1-dimensional array
+        Clumps belonging to this split.
+    """
+    fname = join(
+        dumpfolder, "out_{}_snap_{}_{}.npz".format(Nsim, Nsnap, Nsplit))
+    file = numpy.load(fname)
+    particles, clump_indxs, clumps = (file[f] for f in file.files)
+    if remove_split:
+        remove(fname)
+    return particles, clump_indxs, clumps
+
+
+def pick_single_clump(n, particles, particle_clumps, clumps):
+    """
+    Get particles belonging to the `n`th clump in `clumps` arrays.
+
+    Parameters
+    ----------
+    n : int
+        Clump position in `clumps` array. Not its halo finder index!
+    particles : structured array
+        Particle array.
+    particle_clumps : 1-dimensional array
+        Array of particles' clump IDs.
+    clumps : structured array
+        Array of clumps.
+
+    Returns
+    -------
+    sel_particles : structured array
+        Particles belonging to the requested clump.
+    sel_clump : array
+        A slice of a `clumps` array corresponding to this clump. Must
+        contain `["peak_x", "peak_y", "peak_z", "mass_cl"]`.
+    """
+    # Clump index on the nth position
+    k = clumps["index"][n]
+    # Mask of which particles belong to this clump
+    mask = particle_clumps == k
+    return particles[mask], clumps[n]
+
+
+class Clump:
+    """
+    A clump (halo) object to handle the particles and their clump's data.
+
+    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
+        Clump center coordinate along the x-axis.
+    y0 : float
+        Clump center coordinate along the y-axis.
+    z0 : float
+        Clump center coordinate along the z-axis.
+    clump_mass : float
+        Mass of the clump.
+    vx : 1-dimensional array
+        Particle velocity along the x-axis.
+    vy : 1-dimensional array
+        Particle velocity along the y-axis.
+    vz : 1-dimensional array
+        Particle velocity along the z-axis.
+    """
+    _r = None
+    _pos = None
+    _clump_pos = None
+    _clump_mass = None
+    _vel = None
+
+    def __init__(self, x, y, z, m, x0, y0, z0, clump_mass=None,
+                 vx=None, vy=None, vz=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
+
+    @property
+    def pos(self):
+        """
+        Cartesian particle coordinates centered at the clump.
+
+        Returns
+        -------
+        pos : 2-dimensional array
+            Array of shape `(n_particles, 3)`.
+        """
+        return self._pos
+
+    @pos.setter
+    def pos(self, X):
+        """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
+
+    @property
+    def Npart(self):
+        """
+        Number of particles associated with this clump.
+
+        Returns
+        -------
+        Npart : int
+            Number of particles.
+        """
+        return self.r.size
+
+    @property
+    def clump_pos(self):
+        """
+        Cartesian clump coordinates.
+
+        Returns
+        -------
+        pos : 1-dimensional array
+            Array of shape `(3, )`.
+        """
+        return self._clump_pos
+
+    @clump_pos.setter
+    def clump_pos(self, pos):
+        """Sets `clump_pos`. Makes sure it is the correct shape."""
+        pos = numpy.asarray(pos)
+        if pos.shape != (3,):
+            raise TypeError("Invalid clump position `{}`".format(pos.shape))
+        self._clump_pos = pos
+
+    @property
+    def clump_mass(self):
+        """
+        Clump mass.
+
+        Returns
+        -------
+        mass : float
+            Clump mass.
+        """
+        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):
+            raise ValueError("`clump_mass` must be a float.")
+        self._clump_mass = mass
+
+    @property
+    def vel(self):
+        """
+        Cartesian particle velocities. Throws an error if they are not set.
+
+        Returns
+        -------
+        vel : 2-dimensional array
+            Array of shape (`n_particles, 3`).
+        """
+        if self._vel is None:
+            raise ValueError("Velocities `vel` have not been set.")
+        return self._vel
+
+    @vel.setter
+    def vel(self, V):
+        """Sets the particle velocities, making sure the shape is OK."""
+        if any(v is None for v in V):
+            return
+        vx, vy, vz = V
+        self._vel = numpy.vstack([vx, vy, vz]).T
+        if self.pos.shape != self.vel.shape:
+            raise ValueError("Different `pos` and `vel` arrays!")
+
+    @property
+    def m(self):
+        """
+        Particle masses.
+
+        Returns
+        -------
+        m : 1-dimensional array
+            Array of shape `(n_particles, )`.
+        """
+        return self._m
+
+    @m.setter
+    def m(self, m):
+        """Sets particle masses `m`, ensuring it is the right size."""
+        if not isinstance(m, numpy.ndarray) and m.size != self.r.size:
+            raise TypeError("`r` and `m` must be equal size 1-dim arrays.")
+        self._m = m
+
+    @property
+    def r(self):
+        """
+        Radial distance of particles from the clump peak.
+
+        Returns
+        -------
+        r : 1-dimensional array
+            Array of shape `(n_particles, )`
+        """
+        return self._r
+
+    @r.setter
+    def r(self, r):
+        """Sets `r`. Again checks the shape."""
+        if not isinstance(r, numpy.ndarray) and r.ndim == 1:
+            raise TypeError("`r` must be a 1-dimensional array.")
+        if not numpy.all(r >= 0):
+            raise ValueError("`r` larger than zero.")
+        self._r = r
+
+    @property
+    def total_particle_mass(self):
+        """
+        Total mass of all particles.
+
+        Returns
+        -------
+        tot_mass : float
+            The summed mass.
+        """
+        return numpy.sum(self.m)
+
+    @property
+    def mean_particle_pos(self):
+        """
+        Mean Cartesian particle coordinate. Not centered at the halo!
+
+        Returns
+        -------
+        pos : 1-dimensional array
+            Array of shape `(3, )`.
+        """
+        return numpy.mean(self.pos + self.clump_pos, axis=0)
+
+    @classmethod
+    def from_arrays(cls, particles, clump):
+        """
+        Initialises `Halo` from `particles` containing the relevant particle
+        information and its `clump` information.
+
+        Paramaters
+        ----------
+        particles : structured array
+            Array of particles belonging to this clump. Must contain
+            `["x", "y", "z", "M"]` and optionally also `["vx", "vy", "vz"]`.
+        clump : array
+            A slice of a `clumps` array corresponding to this clump. Must
+            contain `["peak_x", "peak_y", "peak_z", "mass_cl"]`.
+
+        Returns
+        -------
+        halo : `Halo`
+            An initialised halo object.
+        """
+        x, y, z, m = (particles[p] for p in ["x", "y", "z", "M"])
+        x0, y0, z0, cl_mass = (
+            clump[p] for p in ["peak_x", "peak_y", "peak_z", "mass_cl"])
+        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)

csiborgtools/fits/haloprofile.py

@@ -0,0 +1,461 @@
+# 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 defined as
+
+    .. math::
+        \rho(r) = \frac{\rho_0}{x(1 + x)^2}
+
+    where :math:`x = r / R_s` with free parameters :math:`R_s, \rho_0`.
+
+    Parameters
+    ----------
+    Rs : float
+        Scale radius :math:`R_s`.
+    rho0 : float
+        NFW density parameter :math:`\rho_0`.
+    """
+
+    def __init__(self):
+        pass
+
+    def profile(self, r, Rs, rho0):
+        r"""
+        Halo profile evaluated at :math:`r`.
+
+        Parameters
+        ----------
+        r : float or 1-dimensional array
+            Radial distance :math:`r`.
+        Rs : float
+            Scale radius :math:`R_s`.
+        rho0 : float
+            NFW density parameter :math:`\rho_0`.
+
+        Returns
+        -------
+        density : float or 1-dimensional array
+            Density of the NFW profile at :math:`r`.
+        """
+        x = r / Rs
+        return rho0 / (x * (1 + x)**2)
+
+    def logprofile(self, r, Rs, rho0):
+        r"""
+        Natural logarithm of the halo profile evaluated at :math:`r`.
+
+        Parameters
+        ----------
+        r : float or 1-dimensional array
+            Radial distance :math:`r`.
+        Rs : float
+            Scale radius :math:`R_s`.
+        rho0 : float
+            NFW density parameter :math:`\rho_0`.
+
+        Returns
+        -------
+        logdensity : float or 1-dimensional array
+            Logarithmic density of the NFW profile at :math:`r`.
+        """
+        x = r / Rs
+        return numpy.log(rho0) - numpy.log(x) - 2 * numpy.log(1 + x)
+
+    def enclosed_mass(self, r, Rs, rho0):
+        r"""
+        Enclosed mass  of a NFW profile in radius :math:`r`.
+
+        Parameters
+        ----------
+        r : float or 1-dimensional array
+            Radial distance :math:`r`.
+        Rs : float
+            Scale radius :math:`R_s`.
+        rho0 : float
+            NFW density parameter :math:`\rho_0`.
+
+        Returns
+        -------
+        M : float or 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_enclosed_mass(self, rmin, rmax, Rs, rho0):
+        """
+        Calculate the enclosed mass between :math:`r_min <= r <= r_max`.
+
+        Parameters
+        ----------
+        rmin : float
+            The minimum radius.
+        rmax : float
+            The maximum radius.
+        Rs : float
+            Scale radius :math:`R_s`.
+        rho0 : float
+            NFW density parameter :math:`\rho_0`.
+
+        Returns
+        -------
+        M : float
+            The enclosed mass within the radial range.
+        """
+        return (self.enclosed_mass(rmax, Rs, rho0)
+                - self.enclosed_mass(rmin, Rs, rho0))
+
+    def pdf(self, r, Rs, rmin, rmax):
+        r"""
+        The radial probability density function of the NFW profile calculated
+        as
+
+        .. 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.
+
+        Parameters
+        ----------
+        r : float or 1-dimensional array
+            Radial distance :math:`r`.
+        Rs : float
+            Scale radius :math:`R_s`.
+        rmin : float
+            The minimum radius.
+        rmax : float
+            The maximum radius.
+
+        Returns
+        -------
+        pdf : float or 1-dimensional array
+            Probability density of the NFW profile at :math:`r`.
+        """
+
+        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, N=1):
+        """
+        Generate random samples from the NFW profile via rejection sampling.
+
+        Parameters
+        ----------
+        rmin : float
+            The minimum radius.
+        rmax : float
+            The maximum radius.
+        Rs : float
+            Scale radius :math:`R_s`.
+        N : 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(N, numpy.nan)
+        for i in range(N):
+            while True:
+                r = gen.rvs()
+                if self.pdf(r, Rs, rmin, rmax) > numpy.random.rand():
+                    samples[i] = r
+                    break
+
+        if N == 1:
+            return samples[0]
+        return samples
+
+
+class NFWPosterior(NFWProfile):
+    r"""
+    Posterior of for fitting the NFW profile in the range specified by the
+    closest and further particle. The likelihood is calculated as
+
+    .. math::
+        \frac{4\pi r^2 \rho(r)} {M(r_\min, r_\max)} \frac{m}{M / N}
+
+    where :math:`M(r_\min, r_\max)` is the enclosed mass between the closest
+    and further particle as expected from a NFW profile, :math:`m` is the
+    particle mass, :math:`M` is the sum of the particle masses and :math:`N`
+    is the number of particles.
+
+    Paramaters
+    ----------
+    clump : `Clump`
+        Clump object containing the particles and clump information.
+    """
+    _clump = None
+    _N = None
+    _rmin = None
+    _rmax = None
+    _binsguess = 10
+
+    def __init__(self, clump):
+        # Initialise the NFW profile
+        super().__init__()
+        self.clump = clump
+
+    @property
+    def clump(self):
+        """
+        Clump object.
+
+        Returns
+        -------
+        clump : `Clump`
+            The clump object.
+        """
+        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):
+        """
+        Radial distance of particles.
+
+        Returns
+        -------
+        r : 1-dimensional array
+            Radial distance of particles.
+        """
+        return self.clump.r
+
+    @property
+    def rmin(self):
+        """
+        Minimum radial distance of a particle belonging to this clump.
+
+        Returns
+        -------
+        rmin : float
+            The minimum distance.
+        """
+        return self._rmin
+
+    @property
+    def rmax(self):
+        """
+        Maximum radial distance of a particle belonging to this clump.
+
+        Returns
+        -------
+        rmin : float
+            The maximum distance.
+        """
+        return self._rmax
+
+    @property
+    def N(self):
+        """
+        The number of particles in this clump.
+
+        Returns
+        -------
+        N : int
+            Number of particles.
+        """
+        return self._N
+
+    def rho0_from_logRs(self, logRs):
+        """
+        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
+        ----------
+        logRs : float
+            Logarithmic scale factor in units matching the coordinates.
+
+        Returns
+        -------
+        rho0: float
+            The NFW density parameter.
+        """
+        Mtot = numpy.exp(self._logMtot)
+        Rs = numpy.exp(logRs)
+        Mnfw_norm = self.bounded_enclosed_mass(self.rmin, self.rmax, Rs, 1)
+        return Mtot / Mnfw_norm
+
+    def logprior(self, logRs):
+        r"""
+        Logarithmic uniform prior on :math:`\log R_{\rm s}`.
+
+        Parameters
+        ----------
+        logRs : float
+            Logarithmic scale factor in units matching the coordinates.
+
+        Returns
+        -------
+        ll : float
+            The logarithmic prior.
+        """
+        if not self._logrmin < logRs < self._logrmax:
+            return - numpy.infty
+        return - self._logprior_volume
+
+    def loglikelihood(self, logRs):
+        """
+        Logarithmic likelihood.
+
+        Parameters
+        ----------
+        logRs : float
+            Logarithmic scale factor in units matching the coordinates.
+
+        Returns
+        -------
+        ll : float
+            The logarithmic likelihood.
+        """
+        Rs = numpy.exp(logRs)
+        # Expected enclosed mass from a NFW
+        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)
+
+    @property
+    def initlogRs(self):
+        r"""
+        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}`.
+
+        Returns
+        -------
+        initlogRs : float
+            The initial guess of :math:`\log R_{\rm s}`.
+        """
+        bins = numpy.linspace(self.rmin, self.rmax, self._binsguess)
+        counts, edges = numpy.histogram(self.r, bins)
+        return numpy.log(edges[numpy.argmax(counts)])
+
+    def __call__(self, logRs):
+        """
+        Logarithmic posterior. Sum of the logarithmic prior and likelihood.
+
+        Parameters
+        ----------
+        logRs : float
+            Logarithmic scale factor in units matching the coordinates.
+
+        Returns
+        -------
+        lpost : float
+            The logarithmic posterior.
+        """
+        lp = self.logprior(logRs)
+        if not numpy.isfinite(lp):
+            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
+        -------
+        res : `scipy.optimize.OptimizeResult`
+            Optimisation result.
+        """
+        bounds = (self._logrmin, self._logrmax)
+        return minimize_scalar(
+            self.hamiltonian, bounds=bounds, method='bounded')
+
+    @classmethod
+    def from_coords(cls, x, y, z, m, x0, y0, z0):
+        """
+        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

@@ -13,8 +13,10 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from .readsim import (get_csiborg_ids, get_sim_path, get_snapshot_path,  # noqa
+from .readsim import (get_csiborg_ids, get_sim_path, get_snapshots,  # noqa
-                      read_info,  # 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
 from .readobs import (read_planck2015, read_2mpp)  # noqa
+from .outsim import (dump_split, combine_splits)  # noqa

csiborgtools/io/outsim.py

@@ -0,0 +1,124 @@
+# Copyright (C) 2022 Richard Stiskalek, Harry Desmond
+# 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.
+"""
+I/O functions for analysing the CSiBORG realisations.
+"""
+
+
+import numpy
+from os.path import join
+from os import remove
+from tqdm import trange
+from .readsim import (get_sim_path, read_clumps)
+
+I64 = numpy.int64
+F64 = numpy.float64
+
+
+def dump_split(arr, Nsplit, Nsim, Nsnap, outdir):
+    """
+    Dump an array from a split.
+
+    Parameters
+    ----------
+    arr : n-dimensional or structured array
+        Array to be saved.
+    Nsplit : int
+        The split index.
+    Nsim : int
+        The CSiBORG realisation index.
+    Nsnap : int
+        The index of a redshift snapshot.
+    outdir : string
+        Directory where to save the temporary files.
+
+    Returns
+    -------
+    None
+    """
+    Nsim = str(Nsim).zfill(5)
+    Nsnap = str(Nsnap).zfill(5)
+    fname = join(outdir, "ramses_out_{}_{}_{}.npy".format(Nsim, Nsnap, Nsplit))
+    numpy.save(fname, arr)
+
+
+def combine_splits(Nsplits, Nsim, Nsnap, outdir, cols_add, remove_splits=False,
+                   verbose=True):
+    """
+    Combine results of many splits saved from `dump_split`. Identifies to which
+    clump the clumps in the split correspond to by matching their index.
+    Returns an array that contains the original clump data along with the newly
+    calculated quantities.
+
+    Paramaters
+    ----------
+    Nsplits : int
+        The total number of clump splits.
+    Nsim : int
+        The CSiBORG realisation index.
+    Nsnap : int
+        The index of a redshift snapshot.
+    outdir : str
+        Directory where to save the new array.
+    cols_add : list of `(str, dtype)`
+        Colums to add. Must be formatted as, for example,
+        `[("npart", numpy.float64), ("totpartmass", numpy.float64)]`.
+    remove_splits : bool, optional
+        Whether to remove the splits files. By default `False`.
+    verbose : bool, optional
+        Verbosity flag. By default `True`.
+
+    Returns
+    -------
+    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)
+    # Get the old + new dtypes and create an empty array
+    descr = clumps.dtype.descr + cols_add
+    out = numpy.full(clumps.size, numpy.nan, dtype=descr)
+    # Now put the old values into the array
+    for par in clumps.dtype.names:
+        out[par] = clumps[par]
+
+    # Filename of splits data
+    froot = "ramses_out_{}_{}".format(str(Nsim).zfill(5), str(Nsnap).zfill(5))
+    fname = join(outdir, froot + "_{}.npy")
+
+    # Iterate over splits and add to the output array
+    cols_add_names = [col[0] for col in cols_add]
+    iters = trange(Nsplits) if verbose else range(Nsplits)
+    for n in iters:
+        fnamesplit = fname.format(n)
+        arr = numpy.load(fnamesplit)
+
+        # Check that all halo indices from the split are in the clump file
+        if not numpy.alltrue(numpy.isin(arr["index"], out["index"])):
+            raise KeyError("....")
+        # Mask of where to put the values from the split
+        mask = numpy.isin(out["index"], arr["index"])
+        for par in cols_add_names:
+            out[par][mask] = arr[par]
+
+        # Now remove this split
+        if remove_splits:
+            remove(fnamesplit)
+
+    return out

csiborgtools/io/readsim.py

@@ -12,8 +12,9 @@
 # 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.
-
+"""
-"""Functions to read in the particle and clump files."""
+Functions to read in the particle and clump files.
+"""
 
 import numpy
 from scipy.io import FortranFile
@@ -84,6 +85,27 @@ def get_sim_path(n, fname="ramses_out_{}", srcdir="/mnt/extraspace/hdesmond"):
     return join(srcdir, fname.format(n))
 
 
+def get_snapshots(simpath):
+    """
+    Get the list of snapshots for the given IC realisation.
+
+    Parameters
+    ----------
+    simpath : str
+        Path to the CSiBORG IC realisation.
+
+    Returns
+    -------
+    snapshots : 1-dimensional array
+        Array of snapshot IDs.
+    """
+    # Get all files in simpath that start with output_
+    snaps = glob(join(simpath, "output_*"))
+    # Take just the last _00XXXX from each file  and strip zeros
+    snaps = [int(snap.split('_')[-1].lstrip('0')) for snap in snaps]
+    return numpy.sort(snaps)
+
+
 def get_snapshot_path(Nsnap, simpath):
     """
     Get a path to a CSiBORG IC realisation snapshot.
@@ -135,13 +157,13 @@ def read_info(Nsnap, simpath):
     return {key: val for key, val in zip(keys, vals)}
 
 
-def open_particle(n, simpath, verbose=True):
+def open_particle(Nsnap, simpath, verbose=True):
     """
     Open particle files to a given CSiBORG simulation.
 
     Parameters
     ----------
-    n : int
+    Nsnap : int
         The index of a redshift snapshot.
     simpath : str
         The complete path to the CSiBORG simulation.
@@ -156,9 +178,9 @@ def open_particle(n, simpath, verbose=True):
         Opened part files.
     """
     # Zeros filled snapshot number and the snapshot path
-    nout = str(n).zfill(5)
+    nout = str(Nsnap).zfill(5)
-    snappath = get_snapshot_path(n, simpath)
+    snappath = get_snapshot_path(Nsnap, simpath)
-    ncpu = int(read_info(n, simpath)["ncpu"])
+    ncpu = int(read_info(Nsnap, simpath)["ncpu"])
 
     if verbose:
         print("Reading in output `{}` with ncpu = `{}`.".format(nout, ncpu))
@@ -245,7 +267,7 @@ def nparts_to_start_ind(nparts):
     return numpy.hstack([[0], numpy.cumsum(nparts[:-1])])
 
 
-def read_particle(pars_extract, n, simpath, verbose=True):
+def read_particle(pars_extract, Nsnap, simpath, verbose=True):
     """
     Read particle files of a simulation at a given snapshot and return
     values of `pars_extract`.
@@ -254,7 +276,7 @@ def read_particle(pars_extract, n, simpath, verbose=True):
     ----------
     pars_extract : list of str
         Parameters to be extacted.
-    n : int
+    Nsnap : int
         The index of the redshift snapshot.
     simpath : str
         The complete path to the CSiBORG simulation.
@@ -267,17 +289,19 @@ def read_particle(pars_extract, n, simpath, verbose=True):
         The data read from the particle file.
     """
     # Open the particle files
-    nparts, partfiles = open_particle(n, simpath)
+    nparts, partfiles = open_particle(Nsnap, simpath)
     if verbose:
         print("Opened {} particle files.".format(nparts.size))
     ncpu = nparts.size
     # Order in which the particles are written in the FortranFile
-    forder = [("x", F16), ("y", F16), ("z", F16),
+    forder = [("x", F32), ("y", F32), ("z", F32),
-              ("vx", F16), ("vy", F16), ("vz", F16),
+              ("vx", F32), ("vy", F32), ("vz", F32),
               ("M", F32), ("ID", I32), ("level", I32)]
     fnames = [fp[0] for fp in forder]
     fdtypes = [fp[1] for fp in forder]
     # Check there are no strange parameters
+    if isinstance(pars_extract, str):
+        pars_extract = [pars_extract]
     for p in pars_extract:
         if p not in fnames:
             raise ValueError("Undefined parameter `{}`. Must be one of `{}`."
@@ -305,7 +329,7 @@ def read_particle(pars_extract, n, simpath, verbose=True):
     return out
 
 
-def open_unbinding(cpu, n, simpath):
+def open_unbinding(cpu, Nsnap, simpath):
     """
     Open particle files to a given CSiBORG simulation. Note that to be
     consistent CPU is incremented by 1.
@@ -314,7 +338,7 @@ def open_unbinding(cpu, n, simpath):
     ----------
     cpu : int
         The CPU index.
-    n : int
+    Nsnap : int
         The index of a redshift snapshot.
     simpath : str
         The complete path to the CSiBORG simulation.
@@ -324,7 +348,7 @@ def open_unbinding(cpu, n, simpath):
     unbinding : `scipy.io.FortranFile`
         The opened unbinding FortranFile.
     """
-    nout = str(n).zfill(5)
+    nout = str(Nsnap).zfill(5)
     cpu = str(cpu + 1).zfill(5)
     fpath = join(simpath, "output_{}".format(nout),
                  "unbinding_{}.out{}".format(nout, cpu))
@@ -332,13 +356,13 @@ def open_unbinding(cpu, n, simpath):
     return FortranFile(fpath)
 
 
-def read_clumpid(n, simpath, verbose=True):
+def read_clumpid(Nsnap, simpath, verbose=True):
     """
-    Read clump IDs from unbinding files.
+    Read clump IDs of halos from unbinding files.
 
     Parameters
     ----------
-    n : int
+    Nsnap : int
         The index of a redshift snapshot.
     simpath : str
         The complete path to the CSiBORG simulation.
@@ -350,52 +374,94 @@ def read_clumpid(n, simpath, verbose=True):
     clumpid : 1-dimensional array
         The array of clump IDs.
     """
-    nparts, __ = open_particle(n, simpath, verbose)
+    nparts, __ = open_particle(Nsnap, simpath, verbose)
     start_ind = nparts_to_start_ind(nparts)
     ncpu = nparts.size
 
-    clumpid = numpy.full(numpy.sum(nparts), numpy.nan)
+    clumpid = numpy.full(numpy.sum(nparts), numpy.nan, dtype=I32)
     iters = tqdm(range(ncpu)) if verbose else range(ncpu)
     for cpu in iters:
         i = start_ind[cpu]
         j = nparts[cpu]
-        ff = open_unbinding(cpu, n, simpath)
+        ff = open_unbinding(cpu, Nsnap, simpath)
         clumpid[i:i + j] = ff.read_ints()
 
     return clumpid
 
 
-def read_clumps(n, simpath):
+def drop_zero_indx(clump_ids, particles):
     """
-    Read in a precomputed clump file `clump_N.dat`.
+    Drop from `clump_ids` and `particles` entries whose clump index is 0.
 
     Parameters
     ----------
-    n : int
+    clump_ids : 1-dimensional array
+        Array of clump IDs.
+    particles : structured array
+        Array of the particle data.
+
+    Returns
+    -------
+    clump_ids : 1-dimensional array
+        The array of clump IDs after removing zero clump ID entries.
+    particles : structured array
+        The particle data after removing zero clump ID entries.
+    """
+    mask = clump_ids != 0
+    return clump_ids[mask], particles[mask]
+
+
+def read_clumps(Nsnap, simpath, cols=None):
+    """
+    Read in a clump file `clump_Nsnap.dat`.
+
+    Parameters
+    ----------
+    Nsnap : int
         The index of a redshift snapshot.
     simpath : str
         The complete path to the CSiBORG simulation.
+    cols : list of str, optional.
+        Columns to extract. By default `None` and all columns are extracted.
 
     Returns
     -------
     out : structured array
         Structured array of the clumps.
     """
-    n = str(n).zfill(5)
+    Nsnap = str(Nsnap).zfill(5)
-    fname = join(simpath, "output_{}".format(n), "clump_{}.dat".format(n))
+    fname = join(simpath, "output_{}".format(Nsnap),
+                 "clump_{}.dat".format(Nsnap))
     # Check the file exists.
     if not isfile(fname):
         raise FileExistsError("Clump file `{}` does not exist.".format(fname))
 
     # Read in the clump array. This is how the columns must be written!
-    arr = numpy.genfromtxt(fname)
+    data = numpy.genfromtxt(fname)
-    cols = [("index", I64), ("level", I64), ("parent", I64), ("ncell", F64),
+    clump_cols = [("index", I64), ("level", I64), ("parent", I64),
-            ("peak_x", F64), ("peak_y", F64), ("peak_z", F64),
+                  ("ncell", F64), ("peak_x", F64), ("peak_y", F64),
-            ("rho-", F64), ("rho+", F64), ("rho_av", F64),
+                  ("peak_z", F64), ("rho-", F64), ("rho+", F64),
-            ("mass_cl", F64), ("relevance", F64)]
+                  ("rho_av", F64), ("mass_cl", F64), ("relevance", F64)]
-    out = cols_to_structured(arr.shape[0], cols)
+    out0 = cols_to_structured(data.shape[0], clump_cols)
-    for i, name in enumerate(out.dtype.names):
+    for i, name in enumerate(out0.dtype.names):
-        out[name] = arr[:, i]
+        out0[name] = data[:, i]
+    # If take all cols then return
+    if cols is None:
+        return out0
+    # Make sure we have a list
+    cols = [cols] if isinstance(cols, str) else cols
+    # Get the indxs of clump_cols to output
+    clump_names = [col[0] for col in clump_cols]
+    indxs = [None] * len(cols)
+    for i, col in enumerate(cols):
+        if col not in clump_names:
+            raise KeyError("...")
+        indxs[i] = clump_names.index(col)
+    # Make an array and fill it
+    out = cols_to_structured(out0.size, [clump_cols[i] for i in indxs])
+    for name in out.dtype.names:
+        out[name] = out0[name]
+
     return out
 
 

csiborgtools/units/box_units.py

@@ -29,9 +29,11 @@ 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
@@ -62,7 +64,7 @@ class BoxUnits:
 
     def box2kpc(self, length):
         r"""
-        Convert length from box units to :math:`\mathrm{kpc}` (with
+        Convert length from box units to :math:`\mathrm{ckpc}` (with
         :math:`h=0.705`).
 
         Parameters
@@ -73,26 +75,26 @@ class BoxUnits:
         Returns
         -------
         length : foat
-            Length in :math:`\mathrm{kpc}`
+            Length in :math:`\mathrm{ckpc}`
         """
-        return length * self.unit_l / KPC_TO_CM
+        return length * self.unit_l / KPC_TO_CM / self.aexp
 
     def kpc2box(self, length):
         r"""
-        Convert length from :math:`\mathrm{kpc}` (with :math:`h=0.705`) to
+        Convert length from :math:`\mathrm{ckpc}` (with :math:`h=0.705`) to
         box units.
 
         Parameters
         ----------
         length : float
-            Length in :math:`\mathrm{kpc}`
+            Length in :math:`\mathrm{ckpc}`
 
         Returns
         -------
         length : foat
             Length in box units.
         """
-        return length / self.unit_l * KPC_TO_CM
+        return length / self.unit_l * KPC_TO_CM * self.aexp
 
     def solarmass2box(self, mass):
         r"""

scripts/run_fit_halos.py

@@ -0,0 +1,92 @@
+# 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.
+"""
+A script to fit halos (concentration, ...). The particle array of each CSiBORG
+realisation must have been split in advance by `run_split_halos`.
+"""
+
+import numpy
+from os.path import join
+from mpi4py import MPI
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
+    import csiborgtools
+import utils
+
+F64 = numpy.float64
+I64 = numpy.int64
+
+# Simulations and their snapshot to analyze
+Nsims = [9844]
+Nsnap = 1016
+
+# Get MPI things
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+nproc = comm.Get_size()
+
+
+dumpdir = utils.dumpdir
+loaddir = join(utils.dumpdir, "temp")
+cols_collect = [("npart", I64), ("totpartmass", F64), ("logRs", F64),
+                ("rho0", F64)]
+# NOTE later loop over sims too
+Nsim = Nsims[0]
+
+jobs = csiborgtools.fits.split_jobs(utils.Nsplits, nproc)[rank]
+for Nsplit in jobs:
+    print("Rank {} working on {}.".format(rank, Nsplit))
+    parts, part_clumps, clumps = csiborgtools.fits.load_split_particles(
+        Nsplit, loaddir, Nsim, Nsnap, remove_split=False)
+
+    N = clumps.size
+    cols = [("index", I64), ("npart", I64), ("totpartmass", F64),
+            ("logRs", F64), ("rho0", F64)]
+    out = csiborgtools.utils.cols_to_structured(N, cols)
+    out["index"] = clumps["index"]
+
+    for n in range(N):
+        # Pick clump and its particles
+        xs = csiborgtools.fits.pick_single_clump(n, parts, part_clumps, clumps)
+        clump = csiborgtools.fits.Clump.from_arrays(*xs)
+        out["npart"][n] = clump.Npart
+        out["totpartmass"][n] = clump.total_particle_mass
+
+        # NFW profile fit
+        if clump.Npart > 10:
+            nfwpost = csiborgtools.fits.NFWPosterior(clump)
+            logRs = nfwpost.maxpost_logRs()
+            if logRs.success:
+                out["logRs"][n] = logRs.x
+                out["rho0"][n] = nfwpost.rho0_from_logRs(logRs.x)
+
+    csiborgtools.io.dump_split(out, Nsplit, Nsim, Nsnap, dumpdir)
+
+# Force all ranks to wait
+comm.Barrier()
+# Use the rank 0 to combine outputs for this CSiBORG realisation
+if rank == 0:
+    print("Collecting results!")
+    out_collected = csiborgtools.io.combine_splits(
+        utils.Nsplits, Nsim, Nsnap, utils.dumpdir, cols_collect,
+        remove_splits=True, verbose=False)
+    fname = join(utils.dumpdir, "ramses_out_{}_{}.npy"
+                 .format(str(Nsim).zfill(5), str(Nsnap).zfill(5)))
+    print("Saving results to `{}`.".format(fname))
+    numpy.save(fname, out_collected)
+    print("All finished! See ya!")

scripts/run_split_halos.py

@@ -0,0 +1,51 @@
+# 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.
+"""
+Script to split particles into smaller files according to their clump
+membership for faster manipulation. Running this will require a lot of memory.
+"""
+
+from tqdm import tqdm
+from os.path import join
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
+    import csiborgtools
+import utils
+
+Nsims = [9844]
+Nsnap = 1016
+partcols = ["x", "y", "z", "M", "level"]
+dumpdir = join(utils.dumpdir, "temp")
+
+for Nsim in tqdm(Nsims):
+    simpath = csiborgtools.io.get_sim_path(Nsim)
+    # Load the clumps, particles' clump IDs and particles.
+    clumps = csiborgtools.io.read_clumps(Nsnap, simpath)
+    particle_clumps = csiborgtools.io.read_clumpid(Nsnap, simpath,
+                                                   verbose=False)
+    particles = csiborgtools.io.read_particle(partcols, Nsnap, simpath,
+                                              verbose=False)
+    # Drop all particles whose clump index is 0 (not assigned to any halo)
+    particle_clumps, particles = csiborgtools.io.drop_zero_indx(
+        particle_clumps, particles)
+    # Dump it!
+    csiborgtools.fits.dump_split_particles(particles, particle_clumps, clumps,
+                                           utils.Nsplits, dumpdir, Nsim, Nsnap,
+                                           verbose=False)
+
+print("All finished!")

scripts/test_paths.ipynb

@@ -1,352 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "5a38ed25",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T21:53:10.387523Z",
-     "start_time": "2022-10-20T21:53:08.558627Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "try:\n",
-    "    import csiborgtools\n",
-    "except ModuleNotFoundError:\n",
-    "    import sys\n",
-    "    sys.path.append(\"../\")\n",
-    "    import csiborgtools\n",
-    "    \n",
-    "    \n",
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "id": "7accd798",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:00:25.350404Z",
-     "start_time": "2022-10-20T22:00:24.775748Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "simpath = csiborgtools.io.get_sim_path(9844)\n",
-    "Nsnap = 1016\n",
-    "# csiborgtools.io.read_info(1016, simpath)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 59,
-   "id": "131c3107",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:25:59.104572Z",
-     "start_time": "2022-10-20T22:25:59.074460Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "box = csiborgtools.units.BoxUnits(Nsnap, simpath)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 61,
-   "id": "b1a11aa6",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:26:12.931468Z",
-     "start_time": "2022-10-20T22:26:12.671686Z"
-    }
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "3.765003070876428e+19"
-      ]
-     },
-     "execution_count": 61,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "box.box2solarmass(1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 55,
-   "id": "19f39a80",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:24:59.563137Z",
-     "start_time": "2022-10-20T22:24:59.530173Z"
-    }
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "3.749e+19"
-      ]
-     },
-     "execution_count": 55,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "3.749e19"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 57,
-   "id": "5d1413fd",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:25:17.251057Z",
-     "start_time": "2022-10-20T22:25:17.223543Z"
-    }
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "0.9957495198343349"
-      ]
-     },
-     "execution_count": 57,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "box.solarmass2box(3.749e19)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "90cf9d61",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b53e8166",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b3eeca5b",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 49,
-   "id": "93d32800",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:09:12.915516Z",
-     "start_time": "2022-10-20T22:09:12.856213Z"
-    }
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "0.9988759675724455"
-      ]
-     },
-     "execution_count": 49,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "box.kpc2box(677.7 / 0.705 * 1000)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 43,
-   "id": "66d3e98f",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T22:06:46.471218Z",
-     "start_time": "2022-10-20T22:06:46.437023Z"
-    }
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "678.4626139789958"
-      ]
-     },
-     "execution_count": 43,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "box.box2kpc(1) / 1000 * 0.705"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ef00a04c",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d201838b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "filename = \"/mnt/extraspace/hdesmond/ramses_out_9844/output_01016/info_01016.txt\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d012a21b",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "547f7944",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T19:52:22.548299Z",
-     "start_time": "2022-10-20T19:52:22.544342Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "results_filepath = '/mnt/extraspace/deaglan/H-AGN/'\n",
-    "sugata_filepath = '/mnt/extraspace/jeg/greenwhale/Sugata/'\n",
-    "\n",
-    "def info_dir():\n",
-    "    return sugata_filepath + 'INFO/'\n",
-    "\n",
-    "def info_file(nout):\n",
-    "    return info_dir() + 'info_' + str(nout).zfill(5) + '.txt'\n",
-    "\n",
-    "def h_agn_particle_directory(nout):\n",
-    "    return sugata_filepath + 'H-AGN/output_' + str(nout).zfill(5) + '/'\n",
-    "\n",
-    "def cosmo_dir():\n",
-    "    return './data/'\n",
-    "\n",
-    "def cosmo_file(nout):\n",
-    "    return cosmo_dir() + 'cosmo_table_' + str(nout).zfill(5) + '.txt'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "30630635",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T19:52:22.777853Z",
-     "start_time": "2022-10-20T19:52:22.774584Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "info_dir()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a3fc21b2",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T19:56:35.946185Z",
-     "start_time": "2022-10-20T19:56:35.943643Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "filename = \"/mnt/extraspace/hdesmond/ramses_out_9844/output_01016/info_01016.txt\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "fae40a32",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2022-10-20T20:05:04.410987Z",
-     "start_time": "2022-10-20T20:05:04.403721Z"
-    },
-    "scrolled": false
-   },
-   "outputs": [],
-   "source": [
-    "e"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "de7e3ae7",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3.8.0 ('venv_galomatch': virtualenv)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.0"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "f29d02a8350410abc2a9fb79641689d10bf7ab64afc03ec87ca3cf6ed2daa499"
-   }
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

scripts/utils.py

@@ -28,6 +28,10 @@ except ModuleNotFoundError:
     sys.path.append("../")
 
 
+Nsplits = 200
+dumpdir = "/mnt/extraspace/rstiskalek/csiborg/"
+
+
 def load_mmain_convert(n):
     srcdir = "/users/hdesmond/Mmain"
     arr = csiborgtools.io.read_mmain(n, srcdir)