mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-12-22 16:48:01 +00:00
Improving halo fits (#76)
* Add periodic distances * Little corrections * Fix little bug * Modernise the script * Small updates * Remove clump * Add new halo routines * Fix weights * Modernise the script * Add check ups on convergence * More convergence check ups * Edit bounds * Add default argument * Update fit heuristic and NaNs * Change maxiter * Switch NFW minimization to log-sapce * Remove print statement * Turn convert_from_box abstract property required for all boxes. * Move files * Simplify script * Improve the argument parser * Remove optinal argument * Improve argument parser * Add a minimum concentration limit
This commit is contained in:
parent
eb8d070fff
commit
e08c741fc8
13 changed files with 460 additions and 735 deletions
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
|
@ -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.")
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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),
|
||||
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.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : dict
|
||||
"""
|
||||
halo = csiborgtools.fits.Halo(particles, box)
|
||||
|
||||
out = {}
|
||||
out["npart"] = len(halo)
|
||||
out["totpartmass"] = numpy.sum(halo["M"])
|
||||
for i, v in enumerate(["vx", "vy", "vz"]):
|
||||
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
|
||||
|
||||
|
||||
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),
|
||||
("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),
|
||||
]
|
||||
("lambda200c", numpy.float32),]
|
||||
|
||||
|
||||
def fit_halo(particles, clump_info, box):
|
||||
obj = csiborgtools.fits.Clump(particles, clump_info, box)
|
||||
|
||||
out = {}
|
||||
out["npart"] = len(obj)
|
||||
out["totpartmass"] = numpy.sum(obj["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"])
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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),
|
||||
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),]
|
||||
|
||||
|
||||
# MPI loop over simulations
|
||||
jobs = csiborgtools.fits.split_jobs(len(nsims), nproc)[rank]
|
||||
for nsim in [nsims[i] for i in jobs]:
|
||||
nsnap = max(paths.get_snapshots(nsim))
|
||||
overlapper = csiborgtools.match.ParticleOverlap()
|
||||
print(f"{datetime.now()}: rank {rank} calculating simulation `{nsim}`.",
|
||||
flush=True)
|
||||
|
||||
parts = csiborgtools.read.read_h5(paths.initmatch(nsim, "particles"))
|
||||
parts = parts['particles']
|
||||
halo_map = csiborgtools.read.read_h5(paths.particles(nsim))
|
||||
halo_map = halo_map["halomap"]
|
||||
|
||||
cat = csiborgtools.read.CSiBORGHaloCatalogue(
|
||||
nsim, paths, rawdata=True, load_fitted=False, load_initial=False)
|
||||
hid2map = {hid: i for i, hid in enumerate(halo_map[:, 0])}
|
||||
|
||||
out = csiborgtools.read.cols_to_structured(len(cat), cols_collect)
|
||||
out = csiborgtools.read.cols_to_structured(len(cat), cols)
|
||||
for i, hid in enumerate(tqdm(cat["index"]) if verbose else cat["index"]):
|
||||
out["index"][i] = hid
|
||||
part = csiborgtools.read.load_halo_particles(hid, parts, halo_map,
|
||||
hid2map)
|
||||
|
||||
# Skip if the halo is too small.
|
||||
# Skip if the halo has no particles or is too small.
|
||||
if part is None or part.size < 100:
|
||||
continue
|
||||
|
||||
pos, mass = part[:, :3], part[:, 3]
|
||||
# Calculate the centre of mass and the Lagrangian patch size.
|
||||
dist, cm = csiborgtools.fits.dist_centmass(part)
|
||||
# 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)
|
||||
|
|
|
@ -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)
|
||||
|
|
14
scripts/old/pre_mmain.sh
Normal file
14
scripts/old/pre_mmain.sh
Normal file
|
@ -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
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in a new issue