Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools_public.git synced 2025-05-22 02:21:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Within halo work and NFW fit (#4)

Browse source 
* add listing of snapshots

* change distance to comoving

* ignore cp files

* rename nb

* add str to list

* add NFW profile shapes

* add fits imports

* Rename to Nsnap

* in clumps_read only select props

* make clumpid int

* expand doc

* add import

* edit readme

* distribute halos

* add profile & posterior

* add import

* add import

* add documentation

* add rvs and init guess

* update todo

* update nb

* add file

* return end index too

* change clump_ids format to int32

* skeleton of dump particle

* update nb

* add func to drop 0 clump indxs parts

* add import

* add halo dump

* switch to float32

* Update TODO

* update TODO

* add func that loads a split

* add halo object

* Rename to clump

* make post work with a clump

* add optimiser

* add Nsplits

* ignore submission scripts

* ignore .out

* add dumppath

* add job splitting

* add split halos script

* rename file

* renaem files

* rm file

* rename imports

* edit desc

* add pick clump

* add number of particles

* update TODO

* update todo

* add script

* add dumping

* change dumpdir structure

* change dumpdir

* add import

* Remove tqdm

* Increase the number of splits

* rm shuffle option

* Change to remove split

* add emojis

* fix part counts in splits

* change num of splits

* rm with particle cut

* keep splits

* fit only if 10 part and more

* add min distance

* rm warning about not set vels

* update TODO

* calculate rho0 too

* add results collection

* add import

* add func to combine splits

* update TODO

* add extract cols

* update nb

* update TODO
This commit is contained in:
Richard Stiskalek 2022-10-30 20:16:56 +00:00 • committed by GitHub
parent 85a6a6d58a
commit 8a56c22813
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 3815 additions and 397 deletions
Download patch file Download diff file Expand all files Collapse all files

2
csiborgtools/__init__.py
View file

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

5
csiborgtools/fits/__init__.py
View file

@ -12,3 +12,8 @@
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from .haloprofile import (NFWProfile, NFWPosterior) # noqa
from .halo import (distribute_halos, clump_with_particles, # noqa
dump_split_particles, load_split_particles, # noqa
split_jobs, pick_single_clump, Clump) # noqa

465
csiborgtools/fits/halo.py Normal file
View file

@ -0,0 +1,465 @@
# Copyright (C) 2022 Richard Stiskalek, Deaglan Bartlett
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Tools for splitting the particles and a clump object.
"""
import numpy
from os import remove
from warnings import warn
from os.path import join
from tqdm import trange
from ..io import nparts_to_start_ind
def clump_with_particles(particle_clumps, clumps):
"""
Count how many particles does each clump have.
Parameters
----------
particle_clumps : 1-dimensional array
Array of particles' clump IDs.
clumps : structured array
The clumps array.
Returns
-------
with_particles : 1-dimensional array
Array of whether a clump has any particles.
"""
return numpy.isin(clumps["index"], particle_clumps)
def distribute_halos(Nsplits, clumps):
"""
Evenly distribute clump indices to smaller splits. Clumps should only be
clumps that contain particles.
Parameters
----------
Nsplits : int
Number of splits.
clumps : structured array
The clumps array.
Returns
-------
splits : 2-dimensional array
Array of starting and ending indices of each CPU of shape `(Njobs, 2)`.
"""
# Make sure these are unique IDs
indxs = clumps["index"]
if indxs.size > numpy.unique((indxs)).size:
raise ValueError("`clump_indxs` constains duplicate indices.")
Ntotal = indxs.size
Njobs_per_cpu = numpy.ones(Nsplits, dtype=int) * Ntotal // Nsplits
# Split the remainder Ntotal % Njobs among the CPU
Njobs_per_cpu[:Ntotal % Nsplits] += 1
start = nparts_to_start_ind(Njobs_per_cpu)
return numpy.vstack([start, start + Njobs_per_cpu]).T
def dump_split_particles(particles, particle_clumps, clumps, Nsplits,
dumpfolder, Nsim, Nsnap, verbose=True):
"""
Save the data needed for each split so that a process does not have to load
everything.
Parameters
----------
particles : structured array
The particle array.
particle_clumps : 1-dimensional array
Array of particles' clump IDs.
clumps : structured array
The clumps array.
Nsplits : int
Number of times to split the clumps.
dumpfolder : str
Path to the folder where to dump the splits.
Nsim : int
CSiBORG simulation index.
Nsnap : int
Snapshot index.
verbose : bool, optional
Verbosity flag. By default `True`.
Returns
-------
None
"""
if particles.size != particle_clumps.size:
raise ValueError("`particles` must correspond to `particle_clumps`.")
# Calculate which clumps have particles
with_particles = clump_with_particles(particle_clumps, clumps)
clumps = numpy.copy(clumps)[with_particles]
if verbose:
warn(r"There are {:.4f}% clumps that have identified particles."
.format(with_particles.sum() / with_particles.size * 100))
# The starting clump index of each split
splits = distribute_halos(Nsplits, clumps)
fname = join(dumpfolder, "out_{}_snap_{}_{}.npz")
iters = trange(Nsplits) if verbose else range(Nsplits)
tot = 0
for n in iters:
# Lower and upper array index of the clumps array
i, j = splits[n, :]
# Clump indices in this split
indxs = clumps["index"][i:j]
hmin, hmax = indxs.min(), indxs.max()
mask = (particle_clumps >= hmin) & (particle_clumps <= hmax)
# Check number of clumps
npart_unique = numpy.unique(particle_clumps[mask]).size
if indxs.size > npart_unique:
raise RuntimeError(
"Split `{}` contains more unique clumps (`{}`) than there are "
"unique particles' clump indices (`{}`)after removing clumps "
"with no particles.".format(n, indxs.size, npart_unique))
# Dump it!
tot += mask.sum()
fout = fname.format(Nsim, Nsnap, n)
numpy.savez(fout, particles[mask], particle_clumps[mask], clumps[i:j])
# There are particles whose clump ID is > 1 and have no counterpart in the
# clump file. Therefore can save fewer particles, depending on the cut.
if tot > particle_clumps.size:
raise RuntimeError(
"Num. of dumped particles `{}` is greater than the particle file "
"size `{}`.".format(tot, particle_clumps.size))
def split_jobs(Njobs, Ncpu):
"""
Split `Njobs` amongst `Ncpu`.
Parameters
----------
Njobs : int
Number of jobs.
Ncpu : int
Number of CPUs.
Returns
-------
jobs : list of lists of integers
Outer list of each CPU and inner lists for CPU's jobs.
"""
njobs_per_cpu, njobs_remainder = divmod(Njobs, Ncpu)
jobs = numpy.arange(njobs_per_cpu * Ncpu).reshape((njobs_per_cpu, Ncpu)).T
jobs = jobs.tolist()
for i in range(njobs_remainder):
jobs[i].append(njobs_per_cpu * Ncpu + i)
return jobs
def load_split_particles(Nsplit, dumpfolder, Nsim, Nsnap, remove_split=False):
"""
Load particles of a split saved by `dump_split_particles`.
Parameters
--------
Nsplit : int
Split index.
dumpfolder : str
Path to the folder where the splits were dumped.
Nsim : int
CSiBORG simulation index.
Nsnap : int
Snapshot index.
remove_split : bool, optional
Whether to remove the split file. By default `False`.
Returns
-------
particles : structured array
Particle array of this split.
clumps_indxs : 1-dimensional array
Array of particles' clump IDs of this split.
clumps : 1-dimensional array
Clumps belonging to this split.
"""
fname = join(
dumpfolder, "out_{}_snap_{}_{}.npz".format(Nsim, Nsnap, Nsplit))
file = numpy.load(fname)
particles, clump_indxs, clumps = (file[f] for f in file.files)
if remove_split:
remove(fname)
return particles, clump_indxs, clumps
def pick_single_clump(n, particles, particle_clumps, clumps):
"""
Get particles belonging to the `n`th clump in `clumps` arrays.
Parameters
----------
n : int
Clump position in `clumps` array. Not its halo finder index!
particles : structured array
Particle array.
particle_clumps : 1-dimensional array
Array of particles' clump IDs.
clumps : structured array
Array of clumps.
Returns
-------
sel_particles : structured array
Particles belonging to the requested clump.
sel_clump : array
A slice of a `clumps` array corresponding to this clump. Must
contain `["peak_x", "peak_y", "peak_z", "mass_cl"]`.
"""
# Clump index on the nth position
k = clumps["index"][n]
# Mask of which particles belong to this clump
mask = particle_clumps == k
return particles[mask], clumps[n]
class Clump:
"""
A clump (halo) object to handle the particles and their clump's data.
Parameters
----------
x : 1-dimensional array
Particle coordinates along the x-axis.
y : 1-dimensional array
Particle coordinates along the y-axis.
z : 1-dimensional array
Particle coordinates along the z-axis.
m : 1-dimensional array
Particle masses.
x0 : float
Clump center coordinate along the x-axis.
y0 : float
Clump center coordinate along the y-axis.
z0 : float
Clump center coordinate along the z-axis.
clump_mass : float
Mass of the clump.
vx : 1-dimensional array
Particle velocity along the x-axis.
vy : 1-dimensional array
Particle velocity along the y-axis.
vz : 1-dimensional array
Particle velocity along the z-axis.
"""
_r = None
_pos = None
_clump_pos = None
_clump_mass = None
_vel = None
def __init__(self, x, y, z, m, x0, y0, z0, clump_mass=None,
vx=None, vy=None, vz=None):
self.pos = (x, y, z, x0, y0, z0)
self.clump_pos = (x0, y0, z0)
self.clump_mass = clump_mass
self.vel = (vx, vy, vz)
self.m = m
@property
def pos(self):
"""
Cartesian particle coordinates centered at the clump.
Returns
-------
pos : 2-dimensional array
Array of shape `(n_particles, 3)`.
"""
return self._pos
@pos.setter
def pos(self, X):
"""Sets `pos` and calculates radial distance."""
x, y, z, x0, y0, z0 = X
self._pos = numpy.vstack([x - x0, y - y0, z - z0]).T
self.r = numpy.sum(self.pos**2, axis=1)**0.5
@property
def Npart(self):
"""
Number of particles associated with this clump.
Returns
-------
Npart : int
Number of particles.
"""
return self.r.size
@property
def clump_pos(self):
"""
Cartesian clump coordinates.
Returns
-------
pos : 1-dimensional array
Array of shape `(3, )`.
"""
return self._clump_pos
@clump_pos.setter
def clump_pos(self, pos):
"""Sets `clump_pos`. Makes sure it is the correct shape."""
pos = numpy.asarray(pos)
if pos.shape != (3,):
raise TypeError("Invalid clump position `{}`".format(pos.shape))
self._clump_pos = pos
@property
def clump_mass(self):
"""
Clump mass.
Returns
-------
mass : float
Clump mass.
"""
return self._clump_mass
@clump_mass.setter
def clump_mass(self, mass):
"""Sets `clump_mass`, making sure it is a float."""
if not isinstance(mass, float):
raise ValueError("`clump_mass` must be a float.")
self._clump_mass = mass
@property
def vel(self):
"""
Cartesian particle velocities. Throws an error if they are not set.
Returns
-------
vel : 2-dimensional array
Array of shape (`n_particles, 3`).
"""
if self._vel is None:
raise ValueError("Velocities `vel` have not been set.")
return self._vel
@vel.setter
def vel(self, V):
"""Sets the particle velocities, making sure the shape is OK."""
if any(v is None for v in V):
return
vx, vy, vz = V
self._vel = numpy.vstack([vx, vy, vz]).T
if self.pos.shape != self.vel.shape:
raise ValueError("Different `pos` and `vel` arrays!")
@property
def m(self):
"""
Particle masses.
Returns
-------
m : 1-dimensional array
Array of shape `(n_particles, )`.
"""
return self._m
@m.setter
def m(self, m):
"""Sets particle masses `m`, ensuring it is the right size."""
if not isinstance(m, numpy.ndarray) and m.size != self.r.size:
raise TypeError("`r` and `m` must be equal size 1-dim arrays.")
self._m = m
@property
def r(self):
"""
Radial distance of particles from the clump peak.
Returns
-------
r : 1-dimensional array
Array of shape `(n_particles, )`
"""
return self._r
@r.setter
def r(self, r):
"""Sets `r`. Again checks the shape."""
if not isinstance(r, numpy.ndarray) and r.ndim == 1:
raise TypeError("`r` must be a 1-dimensional array.")
if not numpy.all(r >= 0):
raise ValueError("`r` larger than zero.")
self._r = r
@property
def total_particle_mass(self):
"""
Total mass of all particles.
Returns
-------
tot_mass : float
The summed mass.
"""
return numpy.sum(self.m)
@property
def mean_particle_pos(self):
"""
Mean Cartesian particle coordinate. Not centered at the halo!
Returns
-------
pos : 1-dimensional array
Array of shape `(3, )`.
"""
return numpy.mean(self.pos + self.clump_pos, axis=0)
@classmethod
def from_arrays(cls, particles, clump):
"""
Initialises `Halo` from `particles` containing the relevant particle
information and its `clump` information.
Paramaters
----------
particles : structured array
Array of particles belonging to this clump. Must contain
`["x", "y", "z", "M"]` and optionally also `["vx", "vy", "vz"]`.
clump : array
A slice of a `clumps` array corresponding to this clump. Must
contain `["peak_x", "peak_y", "peak_z", "mass_cl"]`.
Returns
-------
halo : `Halo`
An initialised halo object.
"""
x, y, z, m = (particles[p] for p in ["x", "y", "z", "M"])
x0, y0, z0, cl_mass = (
clump[p] for p in ["peak_x", "peak_y", "peak_z", "mass_cl"])
try:
vx, vy, vz = (particles[p] for p in ["vx", "vy", "vz"])
except ValueError:
vx, vy, vz = None, None, None
return cls(x, y, z, m, x0, y0, z0, cl_mass, vx, vy, vz)

461
csiborgtools/fits/haloprofile.py Normal file
View file

@ -0,0 +1,461 @@
# Copyright (C) 2022 Richard Stiskalek, Deaglan Bartlett
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Halo profiles functions and posteriors.
"""
import numpy
from scipy.optimize import minimize_scalar
from scipy.stats import uniform
from .halo import Clump
class NFWProfile:
r"""
The Navarro-Frenk-White (NFW) density profile defined as
.. math::
\rho(r) = \frac{\rho_0}{x(1 + x)^2}
where :math:`x = r / R_s` with free parameters :math:`R_s, \rho_0`.
Parameters
----------
Rs : float
Scale radius :math:`R_s`.
rho0 : float
NFW density parameter :math:`\rho_0`.
"""
def __init__(self):
pass
def profile(self, r, Rs, rho0):
r"""
Halo profile evaluated at :math:`r`.
Parameters
----------
r : float or 1-dimensional array
Radial distance :math:`r`.
Rs : float
Scale radius :math:`R_s`.
rho0 : float
NFW density parameter :math:`\rho_0`.
Returns
-------
density : float or 1-dimensional array
Density of the NFW profile at :math:`r`.
"""
x = r / Rs
return rho0 / (x * (1 + x)**2)
def logprofile(self, r, Rs, rho0):
r"""
Natural logarithm of the halo profile evaluated at :math:`r`.
Parameters
----------
r : float or 1-dimensional array
Radial distance :math:`r`.
Rs : float
Scale radius :math:`R_s`.
rho0 : float
NFW density parameter :math:`\rho_0`.
Returns
-------
logdensity : float or 1-dimensional array
Logarithmic density of the NFW profile at :math:`r`.
"""
x = r / Rs
return numpy.log(rho0) - numpy.log(x) - 2 * numpy.log(1 + x)
def enclosed_mass(self, r, Rs, rho0):
r"""
Enclosed mass of a NFW profile in radius :math:`r`.
Parameters
----------
r : float or 1-dimensional array
Radial distance :math:`r`.
Rs : float
Scale radius :math:`R_s`.
rho0 : float
NFW density parameter :math:`\rho_0`.
Returns
-------
M : float or 1-dimensional array
The enclosed mass.
"""
x = r / Rs
out = numpy.log(1 + x) - x / (1 + x)
return 4 * numpy.pi * rho0 * Rs**3 * out
def bounded_enclosed_mass(self, rmin, rmax, Rs, rho0):
"""
Calculate the enclosed mass between :math:`r_min <= r <= r_max`.
Parameters
----------
rmin : float
The minimum radius.
rmax : float
The maximum radius.
Rs : float
Scale radius :math:`R_s`.
rho0 : float
NFW density parameter :math:`\rho_0`.
Returns
-------
M : float
The enclosed mass within the radial range.
"""
return (self.enclosed_mass(rmax, Rs, rho0)
- self.enclosed_mass(rmin, Rs, rho0))
def pdf(self, r, Rs, rmin, rmax):
r"""
The radial probability density function of the NFW profile calculated
as
.. math::
\frac{4\pi r^2 \rho(r)} {M(r_\min, r_\max)}
where :math:`M(r_\min, r_\max)` is the enclosed mass between
:math:`r_\min` and :math:`r_\max'. Note that the dependance on
:math:`\rho_0` is cancelled.
Parameters
----------
r : float or 1-dimensional array
Radial distance :math:`r`.
Rs : float
Scale radius :math:`R_s`.
rmin : float
The minimum radius.
rmax : float
The maximum radius.
Returns
-------
pdf : float or 1-dimensional array
Probability density of the NFW profile at :math:`r`.
"""
norm = self.bounded_enclosed_mass(rmin, rmax, Rs, 1)
return 4 * numpy.pi * r**2 * self.profile(r, Rs, 1) / norm
def rvs(self, rmin, rmax, Rs, N=1):
"""
Generate random samples from the NFW profile via rejection sampling.
Parameters
----------
rmin : float
The minimum radius.
rmax : float
The maximum radius.
Rs : float
Scale radius :math:`R_s`.
N : int, optional
Number of samples to generate. By default 1.
Returns
-------
samples : float or 1-dimensional array
Samples following the NFW profile.
"""
gen = uniform(rmin, rmax-rmin)
samples = numpy.full(N, numpy.nan)
for i in range(N):
while True:
r = gen.rvs()
if self.pdf(r, Rs, rmin, rmax) > numpy.random.rand():
samples[i] = r
break
if N == 1:
return samples[0]
return samples
class NFWPosterior(NFWProfile):
r"""
Posterior of for fitting the NFW profile in the range specified by the
closest and further particle. The likelihood is calculated as
.. math::
\frac{4\pi r^2 \rho(r)} {M(r_\min, r_\max)} \frac{m}{M / N}
where :math:`M(r_\min, r_\max)` is the enclosed mass between the closest
and further particle as expected from a NFW profile, :math:`m` is the
particle mass, :math:`M` is the sum of the particle masses and :math:`N`
is the number of particles.
Paramaters
----------
clump : `Clump`
Clump object containing the particles and clump information.
"""
_clump = None
_N = None
_rmin = None
_rmax = None
_binsguess = 10
def __init__(self, clump):
# Initialise the NFW profile
super().__init__()
self.clump = clump
@property
def clump(self):
"""
Clump object.
Returns
-------
clump : `Clump`
The clump object.
"""
return self._clump
@clump.setter
def clump(self, clump):
"""Sets `clump` and precalculates useful things."""
if not isinstance(clump, Clump):
raise TypeError(
"`clump` must be :py:class:`csiborgtools.fits.Clump` type. "
"Currently `{}`".format(type(clump)))
self._clump = clump
# Set here the rest of the radial info
self._rmax = numpy.max(self.r)
# r_min could be zero if particle in the potential minimum.
# Set it to the closest particle that is at distance > 0
self._rmin = numpy.sort(self.r[self.r > 0])[0]
self._logrmin = numpy.log(self.rmin)
self._logrmax = numpy.log(self.rmax)
self._logprior_volume = numpy.log(self._logrmax - self._logrmin)
self._N = self.r.size
# Precalculate useful things
self._logMtot = numpy.log(numpy.sum(self.clump.m))
gamma = 4 * numpy.pi * self.r**2 * self.clump.m * self.N
self._ll0 = numpy.sum(numpy.log(gamma)) - self.N * self._logMtot
@property
def r(self):
"""
Radial distance of particles.
Returns
-------
r : 1-dimensional array
Radial distance of particles.
"""
return self.clump.r
@property
def rmin(self):
"""
Minimum radial distance of a particle belonging to this clump.
Returns
-------
rmin : float
The minimum distance.
"""
return self._rmin
@property
def rmax(self):
"""
Maximum radial distance of a particle belonging to this clump.
Returns
-------
rmin : float
The maximum distance.
"""
return self._rmax
@property
def N(self):
"""
The number of particles in this clump.
Returns
-------
N : int
Number of particles.
"""
return self._N
def rho0_from_logRs(self, logRs):
"""
Obtain :math:`\rho_0` of the NFW profile from the integral constraint
on total mass. Calculated as the ratio between the total particle mass
and the enclosed NFW profile mass.
Parameters
----------
logRs : float
Logarithmic scale factor in units matching the coordinates.
Returns
-------
rho0: float
The NFW density parameter.
"""
Mtot = numpy.exp(self._logMtot)
Rs = numpy.exp(logRs)
Mnfw_norm = self.bounded_enclosed_mass(self.rmin, self.rmax, Rs, 1)
return Mtot / Mnfw_norm
def logprior(self, logRs):
r"""
Logarithmic uniform prior on :math:`\log R_{\rm s}`.
Parameters
----------
logRs : float
Logarithmic scale factor in units matching the coordinates.
Returns
-------
ll : float
The logarithmic prior.
"""
if not self._logrmin < logRs < self._logrmax:
return - numpy.infty
return - self._logprior_volume
def loglikelihood(self, logRs):
"""
Logarithmic likelihood.
Parameters
----------
logRs : float
Logarithmic scale factor in units matching the coordinates.
Returns
-------
ll : float
The logarithmic likelihood.
"""
Rs = numpy.exp(logRs)
# Expected enclosed mass from a NFW
Mnfw = self.bounded_enclosed_mass(self.rmin, self.rmax, Rs, 1)
ll = self._ll0 + numpy.sum(self.logprofile(self.r, Rs, 1))
return ll - self.N * numpy.log(Mnfw)
@property
def initlogRs(self):
r"""
The most often occuring value of :math:`r` used as initial guess of
:math:`R_{\rm s}` since :math:`r^2 \rho(r)` peaks at
:math:`r = R_{\rm s}`.
Returns
-------
initlogRs : float
The initial guess of :math:`\log R_{\rm s}`.
"""
bins = numpy.linspace(self.rmin, self.rmax, self._binsguess)
counts, edges = numpy.histogram(self.r, bins)
return numpy.log(edges[numpy.argmax(counts)])
def __call__(self, logRs):
"""
Logarithmic posterior. Sum of the logarithmic prior and likelihood.
Parameters
----------
logRs : float
Logarithmic scale factor in units matching the coordinates.
Returns
-------
lpost : float
The logarithmic posterior.
"""
lp = self.logprior(logRs)
if not numpy.isfinite(lp):
return - numpy.infty
return self.loglikelihood(logRs) + lp
def hamiltonian(self, logRs):
"""
Negative logarithmic posterior (i.e. the Hamiltonian).
Parameters
----------
logRs : float
Logarithmic scale factor in units matching the coordinates.
Returns
-------
neg_lpost : float
The Hamiltonian.
"""
return - self(logRs)
def maxpost_logRs(self):
r"""
Maximum a-posterio estimate of the scale radius :math:`\log R_{\rm s}`.
Returns
-------
res : `scipy.optimize.OptimizeResult`
Optimisation result.
"""
bounds = (self._logrmin, self._logrmax)
return minimize_scalar(
self.hamiltonian, bounds=bounds, method='bounded')
@classmethod
def from_coords(cls, x, y, z, m, x0, y0, z0):
"""
Initiate `NFWPosterior` from a set of Cartesian coordinates.
Parameters
----------
x : 1-dimensional array
Particle coordinates along the x-axis.
y : 1-dimensional array
Particle coordinates along the y-axis.
z : 1-dimensional array
Particle coordinates along the z-axis.
m : 1-dimensional array
Particle masses.
x0 : float
Halo center coordinate along the x-axis.
y0 : float
Halo center coordinate along the y-axis.
z0 : float
Halo center coordinate along the z-axis.
Returns
-------
post : `NFWPosterior`
Initiated `NFWPosterior` instance.
"""
r = numpy.sqrt((x - x0)**2 + (y - y0)**2 + (z - z0)**2)
return cls(r, m)

6
csiborgtools/io/__init__.py
View file

@ -13,8 +13,10 @@
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from .readsim import (get_csiborg_ids, get_sim_path, get_snapshot_path, # noqa
read_info, # noqa
from .readsim import (get_csiborg_ids, get_sim_path, get_snapshots, # noqa
get_snapshot_path, read_info, nparts_to_start_ind, # noqa
open_particle, open_unbinding, read_particle, # noqa
drop_zero_indx, # noqa
read_clumpid, read_clumps, read_mmain) # noqa
from .readobs import (read_planck2015, read_2mpp) # noqa
from .outsim import (dump_split, combine_splits) # noqa

124
csiborgtools/io/outsim.py Normal file
View file

@ -0,0 +1,124 @@
# Copyright (C) 2022 Richard Stiskalek, Harry Desmond
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
I/O functions for analysing the CSiBORG realisations.
"""
import numpy
from os.path import join
from os import remove
from tqdm import trange
from .readsim import (get_sim_path, read_clumps)
I64 = numpy.int64
F64 = numpy.float64
def dump_split(arr, Nsplit, Nsim, Nsnap, outdir):
"""
Dump an array from a split.
Parameters
----------
arr : n-dimensional or structured array
Array to be saved.
Nsplit : int
The split index.
Nsim : int
The CSiBORG realisation index.
Nsnap : int
The index of a redshift snapshot.
outdir : string
Directory where to save the temporary files.
Returns
-------
None
"""
Nsim = str(Nsim).zfill(5)
Nsnap = str(Nsnap).zfill(5)
fname = join(outdir, "ramses_out_{}_{}_{}.npy".format(Nsim, Nsnap, Nsplit))
numpy.save(fname, arr)
def combine_splits(Nsplits, Nsim, Nsnap, outdir, cols_add, remove_splits=False,
verbose=True):
"""
Combine results of many splits saved from `dump_split`. Identifies to which
clump the clumps in the split correspond to by matching their index.
Returns an array that contains the original clump data along with the newly
calculated quantities.
Paramaters
----------
Nsplits : int
The total number of clump splits.
Nsim : int
The CSiBORG realisation index.
Nsnap : int
The index of a redshift snapshot.
outdir : str
Directory where to save the new array.
cols_add : list of `(str, dtype)`
Colums to add. Must be formatted as, for example,
`[("npart", numpy.float64), ("totpartmass", numpy.float64)]`.
remove_splits : bool, optional
Whether to remove the splits files. By default `False`.
verbose : bool, optional
Verbosity flag. By default `True`.
Returns
-------
out : structured array
Clump array with appended results from the splits.
"""
# Will be grabbing these columns from each split
cols_add = [("npart", I64), ("totpartmass", F64), ("logRs", F64),
("rho0", F64)]
# Load clumps to see how many there are and will add to this array
simpath = get_sim_path(Nsim)
clumps = read_clumps(Nsnap, simpath, cols=None)
# Get the old + new dtypes and create an empty array
descr = clumps.dtype.descr + cols_add
out = numpy.full(clumps.size, numpy.nan, dtype=descr)
# Now put the old values into the array
for par in clumps.dtype.names:
out[par] = clumps[par]
# Filename of splits data
froot = "ramses_out_{}_{}".format(str(Nsim).zfill(5), str(Nsnap).zfill(5))
fname = join(outdir, froot + "_{}.npy")
# Iterate over splits and add to the output array
cols_add_names = [col[0] for col in cols_add]
iters = trange(Nsplits) if verbose else range(Nsplits)
for n in iters:
fnamesplit = fname.format(n)
arr = numpy.load(fnamesplit)
# Check that all halo indices from the split are in the clump file
if not numpy.alltrue(numpy.isin(arr["index"], out["index"])):
raise KeyError("....")
# Mask of where to put the values from the split
mask = numpy.isin(out["index"], arr["index"])
for par in cols_add_names:
out[par][mask] = arr[par]
# Now remove this split
if remove_splits:
remove(fnamesplit)
return out

134
csiborgtools/io/readsim.py
View file

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

16
csiborgtools/units/box_units.py
View file

@ -29,9 +29,11 @@ PI = 3.1415926535897932384626433
class BoxUnits:
"""
r"""
Box units class for converting between box and physical units.
TODO: check factors of :math:`a` in mass and density transformations
Paramaters
----------
Nsnap : int
@ -62,7 +64,7 @@ class BoxUnits:
def box2kpc(self, length):
r"""
Convert length from box units to :math:`\mathrm{kpc}` (with
Convert length from box units to :math:`\mathrm{ckpc}` (with
:math:`h=0.705`).
Parameters
@ -73,26 +75,26 @@ class BoxUnits:
Returns
-------
length : foat
Length in :math:`\mathrm{kpc}`
Length in :math:`\mathrm{ckpc}`
"""
return length * self.unit_l / KPC_TO_CM
return length * self.unit_l / KPC_TO_CM / self.aexp
def kpc2box(self, length):
r"""
Convert length from :math:`\mathrm{kpc}` (with :math:`h=0.705`) to
Convert length from :math:`\mathrm{ckpc}` (with :math:`h=0.705`) to
box units.
Parameters
----------
length : float
Length in :math:`\mathrm{kpc}`
Length in :math:`\mathrm{ckpc}`
Returns
-------
length : foat
Length in box units.
"""
return length / self.unit_l * KPC_TO_CM
return length / self.unit_l * KPC_TO_CM * self.aexp
def solarmass2box(self, mass):
r"""