diff --git a/csiborgtools/fits/__init__.py b/csiborgtools/fits/__init__.py index 9575af0..148eb45 100644 --- a/csiborgtools/fits/__init__.py +++ b/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 - number_counts) -from .haloprofile import NFWPosterior, NFWProfile # noqa +from .halo import (Halo, delta2ncells, center_of_mass, # noqa + periodic_distance, shift_to_center_of_box, number_counts) # noqa from .utils import split_jobs # noqa diff --git a/csiborgtools/fits/halo.py b/csiborgtools/fits/halo.py index 790fc9f..561cd0c 100644 --- a/csiborgtools/fits/halo.py +++ b/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] - self.info[p] for p in ps]).T + return numpy.vstack([self[p] for p in ('x', 'y', 'z')]).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 r(self): - """ - 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): + def spherical_overdensity_mass(self, delta_mult, kind="crit", tol=1e-8, + maxiter=100, npart_min=10): r""" - Bullock spin, see Eq. 5 in [1], in a radius of `radius`, which should - define to some overdensity radius. + Calculate spherical overdensity mass and radius via the iterative + shrinking sphere method. Parameters ---------- - radmax : float - Radius in which to calculate the spin. - npart_min : int + delta_mult : int or float + Overdensity multiple. + 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 - if numpy.sum(mask) < npart_min: - return numpy.nan - mass = self.enclosed_mass(radmax) - circvel = numpy.sqrt(self.box.box_G * mass / radmax) - angmom_norm = numpy.linalg.norm(self.angular_momentum(radmax)) - return angmom_norm / (numpy.sqrt(2) * mass * circvel * radmax) + pos = self.pos + mask = periodic_distance(pos, ref, boxsize=1) < rad + mass = numpy.sum(self["M"][mask]) + circvel = numpy.sqrt(self.box.box_G * mass / rad) + angmom_norm = numpy.linalg.norm(self.angular_momentum(ref, rad)) + return angmom_norm / (numpy.sqrt(2) * mass * circvel * rad) - def spherical_overdensity_mass(self, delta_mult, npart_min=10, - kind="crit"): - r""" - Calculate spherical overdensity mass and radius. The mass is defined as - 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`). + def nfw_concentration(self, ref, rad, conc_min=1e-3, npart_min=10): + """ + Calculate the NFW concentration parameter in a given radius around a + reference point. Parameters ---------- - delta_mult : list of int or float - Overdensity multiple. - npart_min : int - Minimum number of enclosed particles for a radius to be - considered trustworthy. - kind : str - Either `crit` or `matter`, for critical or matter overdensity + ref : 1-dimensional array of shape `(3, )` + Reference point. + rad : float + Radius around the reference point. + conc_min : float + Minimum concentration limit. + npart_min : int, optional + Minimum number of enclosed particles to calculate the + concentration. Returns ------- - rx : float - Radius where the enclosed density reaches required value. - mx : float - Corresponding spherical enclosed mass. + conc : float """ - # Quick check of inputs - assert kind in ["crit", "matter"] + pos = self.pos + 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 - rs = self.r - order = numpy.argsort(rs) - rs = rs[order] + dist = dist[mask] + weight = self["M"][mask] + weight /= numpy.mean(weight) - cmass = numpy.cumsum(self["M"][order]) # Cumulative mass - # We calculate the enclosed volume and indices where it is above target - vol = 4 * numpy.pi / 3 * rs**3 + # We do the minimization in log space + def negll_nfw_concentration(log_c, xs, weight): + 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 - if kind == "matter": - target_density *= self.box.cosmo.Om0 - ks = numpy.where(cmass > target_density * vol)[0] - if ks.size == 0: # Never above the threshold? - return numpy.nan, numpy.nan - k = numpy.max(ks) - if k < npart_min: # Too few particles? - return numpy.nan, numpy.nan - return rs[k], cmass[k] + res = minimize(negll_nfw_concentration, x0=1.5, + args=(dist / rad, weight, ), method='Nelder-Mead', + bounds=((numpy.log10(conc_min), 5),)) + + if not res.success: + return numpy.nan + + res = 10**res["x"][0] + if res < conc_min or numpy.isclose(res, conc_min): + 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 - particles halo. + Halo 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. + 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) - Particle array. The first four columns must be `x`, `y`, `z` and `M`. + points : 2-dimensional array of shape (n_particles, 3) + 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, )` - Particle distance from the centre of mass. - cm : len-3 list - Center of mass coordinates. + cm : 1-dimensional array of shape `(3, )` """ - mass = clump[:, 3] - x, y, z = clump[:, 0], clump[:, 1], clump[:, 2] - cmx, cmy, cmz = [numpy.average(xi, weights=mass) for xi in (x, y, z)] - dist = ((x - cmx)**2 + (y - cmy)**2 + (z - cmz)**2)**0.5 - return dist, [cmx, cmy, cmz] + # Convert positions to unit circle coordinates in the complex plane + pos = numpy.exp(2j * numpy.pi * points / boxsize) + # Compute weighted average of these coordinates, convert it back to + # box coordinates and fix any negative positions due to angle calculations. + 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) diff --git a/csiborgtools/fits/haloprofile.py b/csiborgtools/fits/haloprofile.py deleted file mode 100644 index 3a97631..0000000 --- a/csiborgtools/fits/haloprofile.py +++ /dev/null @@ -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 diff --git a/csiborgtools/read/box_units.py b/csiborgtools/read/box_units.py index cf882e0..c454bf9 100644 --- a/csiborgtools/read/box_units.py +++ b/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.") diff --git a/scripts/cluster_tpcf_auto.py b/scripts/cluster_tpcf_auto.py index ec9242d..546f80e 100644 --- a/scripts/cluster_tpcf_auto.py +++ b/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) - - comm.Barrier() - if comm.Get_rank() == 0: - print(f"{datetime.now()}: all finished. Quitting.") + work_delegation(do_work, nsims, comm) diff --git a/scripts/fit_halos.py b/scripts/fit_halos.py index a4d101f..dba21ef 100644 --- a/scripts/fit_halos.py +++ b/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 -CSiBORG realisation must have been processed in advance by `pre_dumppart.py`. +A script to fit FoF halos (concentration, ...). The CSiBORG particle array of +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() -parser.add_argument("--nsims", type=int, nargs="+", default=None, - help="IC realisations. If `-1` processes all simulations.") -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) +def fit_halo(particles, box): + """ + Fit a single halo from the particle array. -cols_collect = [ - ("index", numpy.int32), - ("npart", numpy.int32), - ("totpartmass", numpy.float32), - ("vx", numpy.float32), - ("vy", numpy.float32), - ("vz", numpy.float32), - ("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), - ] + Parameters + ---------- + particles : 2-dimensional array of shape `(n_particles, 3)` + Particle array. The columns must be `x`, `y`, `z`, `vx`, `vy`, `vz`, + `M`. + box : object derived from :py:class`csiborgtools.read.BaseBox` + Box object. - -def fit_halo(particles, clump_info, box): - obj = csiborgtools.fits.Clump(particles, clump_info, box) + Returns + ------- + out : dict + """ + halo = csiborgtools.fits.Halo(particles, box) out = {} - out["npart"] = len(obj) - out["totpartmass"] = numpy.sum(obj["M"]) + out["npart"] = len(halo) + out["totpartmass"] = numpy.sum(halo["M"]) for i, v in enumerate(["vx", "vy", "vz"]): - out[v] = numpy.average(obj.vel[:, i], weights=obj["M"]) - # Overdensity masses - for n in [200, 500]: - out[f"r{n}c"], out[f"m{n}c"] = obj.spherical_overdensity_mass( - n, kind="crit", npart_min=10) - out[f"r{n}m"], out[f"m{n}m"] = obj.spherical_overdensity_mass( - n, kind="matter", npart_min=10) - # NFW fit - 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"]) + out[v] = numpy.average(halo.vel[:, i], weights=halo["M"]) + + m200c, r200c, cm = halo.spherical_overdensity_mass(200, kind="crit", + maxiter=100) + out["m200c"] = m200c + out["r200c"] = r200c + out["lambda200c"] = halo.lambda_bullock(cm, r200c) + out["conc"] = halo.nfw_concentration(cm, r200c) return out -# We MPI loop over all simulations. -jobs = csiborgtools.fits.split_jobs(len(nsims), nproc)[rank] -for nsim in [nsims[i] for i in jobs]: - print(f"{datetime.now()}: rank {rank} calculating simulation `{nsim}`.", - flush=True) +def _main(nsim, simname, verbose): + """ + Fit the FoF halos. + + 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) - indxs = cat["index"] + + out = csiborgtools.read.cols_to_structured(len(cat), cols) 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, - # otherwise it would be NaN converted to integers (-2147483648) and - # yield an error further down. + # Skip if no particles. 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) diff --git a/scripts/fit_hmf.py b/scripts/fit_hmf.py index 786b06a..26cbf58 100644 --- a/scripts/fit_hmf.py +++ b/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) diff --git a/scripts/fit_init.py b/scripts/fit_init.py index 1059a67..29b011e 100644 --- a/scripts/fit_init.py +++ b/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 -comm = MPI.COMM_WORLD -rank = comm.Get_rank() -nproc = comm.Get_size() -verbose = nproc == 1 +def _main(nsim, simname, verbose): + """ + Calculate the Lagrangian halo centre of mass and Lagrangian patch size in + the initial snapshot. -# Argument parser -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 simulations.") -args = parser.parse_args() -paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) -partreader = csiborgtools.read.ParticleReader(paths) + Parameters + ---------- + nsim : int + IC realisation index. + simname : str + Simulation name. + verbose : bool + Verbosity flag. + """ + if simname == "quijote": + raise NotImplementedError("Quijote not implemented yet.") -nsims = get_nsims(args, paths) - -cols_collect = [("index", numpy.int32), - ("x", numpy.float32), - ("y", numpy.float32), - ("z", numpy.float32), - ("lagpatch_size", numpy.float32), - ("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) + paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) + cols = [("index", numpy.int32), + ("x", numpy.float32), + ("y", numpy.float32), + ("z", numpy.float32), + ("lagpatch_size", numpy.float32), + ("lagpatch_ncells", numpy.int32),] 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) - # We enforce a maximum patchsize of 0.075 in box coordinates. - patchsize = min(numpy.percentile(dist, 99), 0.075) + cm = csiborgtools.fits.center_of_mass(pos, mass, boxsize=1.0) + distances = csiborgtools.fits.periodic_distance(pos, cm, boxsize=1.0) 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}`.", - flush=True) + if verbose: + 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) diff --git a/scripts/mv_fofmembership.py b/scripts/mv_fofmembership.py index 8733c8e..d759fdd 100644 --- a/scripts/mv_fofmembership.py +++ b/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, comm) + work_delegation(main, nsims, MPI.COMM_WORLD) diff --git a/scripts/pre_mmain.py b/scripts/old/pre_mmain.py similarity index 100% rename from scripts/pre_mmain.py rename to scripts/old/pre_mmain.py diff --git a/scripts/old/pre_mmain.sh b/scripts/old/pre_mmain.sh new file mode 100644 index 0000000..f43712a --- /dev/null +++ b/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 diff --git a/scripts/pre_dumppart.py b/scripts/pre_dumppart.py index 93d9fb5..1bda078 100644 --- a/scripts/pre_dumppart.py +++ b/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): - main(nsim, args.simname, verbose=verbose) + def _main(nsim): + main(nsim, args.simname, verbose=MPI.COMM_WORLD.Get_size() == 1) work_delegation(_main, nsims, MPI.COMM_WORLD) diff --git a/scripts/pre_sortinit.py b/scripts/pre_sortinit.py index 1067084..5c8bea9 100644 --- a/scripts/pre_sortinit.py +++ b/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)