mirror of
https://github.com/Richard-Sti/csiborgtools_public.git
synced 2025-06-08 18:01:11 +00:00
Quijote snapshots support (#77)
* Renaming * Edit docs * Delete old function * Add a blank space * Rename particle reader * Add comments * Rename * Rename * edit get_snapshots * More renaming * Remove old correction * Add import * Add basics of the Quijote reader * Add a blank space * Fix paths * Rename function * Fix HID and path * Add more FoF reading * Move definition * Adding arguments * Renaming * Add kwargs for backward comp * FoF Quijote return only hids * Add sorting of quijote * Add path to CSiBORG ICs snapshot * Add support for Quijote * initmatch paths for quijote * Add kwargs * Fix blank lines * Rename kwarg * Remove unused import * Remove hardcoded numbers * Update for Quijote * Do not store velocities in QUijote ICs * Box units mass Quijote * Fix typo * Ensure particles are not right at the edge * Add structfit paths for QUuijote * Basic CSiBORG units * Add more quijote halo reading * Add Quijote fitting * Docs changes * Docs changes
This commit is contained in:
parent
e08c741fc8
commit
fb4b4edf19
18 changed files with 800 additions and 300 deletions
|
@ -117,7 +117,7 @@ class BaseStructure(ABC):
|
|||
assert kind in ["crit", "matter"]
|
||||
rho = delta_mult * self.box.box_rhoc
|
||||
if kind == "matter":
|
||||
rho *= self.box.box_Om
|
||||
rho *= self.box.Om
|
||||
pos = self.pos
|
||||
mass = self["M"]
|
||||
|
||||
|
@ -216,8 +216,8 @@ class BaseStructure(ABC):
|
|||
References
|
||||
----------
|
||||
[1] A Universal Angular Momentum Profile for Galactic Halos; 2001;
|
||||
Bullock, J. S.; Dekel, A.; Kolatt, T. S.; Kravtsov, A. V.;
|
||||
Klypin, A. A.; Porciani, C.; Primack, J. R.
|
||||
Bullock, J. S.; Dekel, A.; Kolatt, T. S.; Kravtsov, A. V.;
|
||||
Klypin, A. A.; Porciani, C.; Primack, J. R.
|
||||
"""
|
||||
pos = self.pos
|
||||
mask = periodic_distance(pos, ref, boxsize=1) < rad
|
||||
|
|
|
@ -15,6 +15,7 @@
|
|||
"""
|
||||
Support for matching halos between CSiBORG IC realisations.
|
||||
"""
|
||||
from abc import ABC
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from math import ceil
|
||||
|
@ -26,20 +27,72 @@ from tqdm import tqdm, trange
|
|||
|
||||
from ..read import load_halo_particles
|
||||
|
||||
BCKG_HALFSIZE = 475
|
||||
BOX_SIZE = 2048
|
||||
|
||||
class BaseMatcher(ABC):
|
||||
"""
|
||||
Base class for `RealisationsMatcher` and `ParticleOverlap`.
|
||||
"""
|
||||
_box_size = None
|
||||
_bckg_halfsize = None
|
||||
|
||||
@property
|
||||
def box_size(self):
|
||||
"""
|
||||
Number of cells in the box.
|
||||
|
||||
Returns
|
||||
-------
|
||||
box_size : int
|
||||
"""
|
||||
if self._box_size is None:
|
||||
raise RuntimeError("`box_size` is not set.")
|
||||
return self._box_size
|
||||
|
||||
@box_size.setter
|
||||
def box_size(self, value):
|
||||
assert isinstance(value, int)
|
||||
assert value > 0
|
||||
self._box_size = value
|
||||
|
||||
@property
|
||||
def bckg_halfsize(self):
|
||||
"""
|
||||
Number of to each side of the centre of the box to calculate the
|
||||
density field. This is because in CSiBORG we are only interested in the
|
||||
high-resolution region.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bckg_halfsize : int
|
||||
"""
|
||||
if self._bckg_halfsize is None:
|
||||
raise RuntimeError("`bckg_halfsize` is not set.")
|
||||
return self._bckg_halfsize
|
||||
|
||||
@bckg_halfsize.setter
|
||||
def bckg_halfsize(self, value):
|
||||
assert isinstance(value, int)
|
||||
assert value > 0
|
||||
self._bckg_halfsize = value
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Realisations matcher for calculating overlaps #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class RealisationsMatcher:
|
||||
class RealisationsMatcher(BaseMatcher):
|
||||
"""
|
||||
A tool to match haloes between IC realisations.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
box_size : int
|
||||
Number of cells in the box.
|
||||
bckg_halfsize : int
|
||||
Number of to each side of the centre of the box to calculate the
|
||||
density field. This is because in CSiBORG we are only interested in the
|
||||
high-resolution region.
|
||||
nmult : float or int, optional
|
||||
Multiple of the sum of pair initial Lagrangian patch sizes
|
||||
within which to return neighbours. By default 1.
|
||||
|
@ -51,20 +104,22 @@ class RealisationsMatcher:
|
|||
catalogue key. By default `totpartmass`, i.e. the total particle
|
||||
mass associated with a halo.
|
||||
"""
|
||||
|
||||
_nmult = None
|
||||
_dlogmass = None
|
||||
_mass_kind = None
|
||||
_overlapper = None
|
||||
|
||||
def __init__(self, nmult=1.0, dlogmass=2.0, mass_kind="totpartmass"):
|
||||
def __init__(self, box_size, bckg_halfsize, nmult=1.0, dlogmass=2.0,
|
||||
mass_kind="totpartmass"):
|
||||
assert nmult > 0
|
||||
assert dlogmass > 0
|
||||
assert isinstance(mass_kind, str)
|
||||
self.box_size = box_size
|
||||
self.halfsize = bckg_halfsize
|
||||
self._nmult = nmult
|
||||
self._dlogmass = dlogmass
|
||||
self._mass_kind = mass_kind
|
||||
self._overlapper = ParticleOverlap()
|
||||
self._overlapper = ParticleOverlap(box_size, bckg_halfsize)
|
||||
|
||||
@property
|
||||
def nmult(self):
|
||||
|
@ -181,7 +236,7 @@ class RealisationsMatcher:
|
|||
@lru_cache(maxsize=cache_size)
|
||||
def load_cached_halox(hid):
|
||||
return load_processed_halo(hid, particlesx, halo_mapx, hid2mapx,
|
||||
nshift=0, ncells=BOX_SIZE)
|
||||
nshift=0, ncells=self.box_size)
|
||||
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: calculating overlaps.", flush=True)
|
||||
|
@ -195,7 +250,8 @@ class RealisationsMatcher:
|
|||
# Next, we find this halo's particles, total mass, minimum and
|
||||
# maximum cells and convert positions to cells.
|
||||
pos0, mass0, totmass0, mins0, maxs0 = load_processed_halo(
|
||||
k0, particles0, halo_map0, hid2map0, nshift=0, ncells=BOX_SIZE)
|
||||
k0, particles0, halo_map0, hid2map0, nshift=0,
|
||||
ncells=self.box_size)
|
||||
|
||||
# We now loop over matches of this halo and calculate their
|
||||
# overlap, storing them in `_cross`.
|
||||
|
@ -267,7 +323,7 @@ class RealisationsMatcher:
|
|||
@lru_cache(maxsize=cache_size)
|
||||
def load_cached_halox(hid):
|
||||
return load_processed_halo(hid, particlesx, halo_mapx, hid2mapx,
|
||||
nshift=nshift, ncells=BOX_SIZE)
|
||||
nshift=nshift, ncells=self.box_size)
|
||||
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: calculating smoothed overlaps.",
|
||||
|
@ -277,7 +333,7 @@ class RealisationsMatcher:
|
|||
for i, k0 in enumerate(tqdm(indxs) if verbose else indxs):
|
||||
pos0, mass0, __, mins0, maxs0 = load_processed_halo(
|
||||
k0, particles0, halo_map0, hid2map0, nshift=nshift,
|
||||
ncells=BOX_SIZE)
|
||||
ncells=self.box_size)
|
||||
|
||||
# Now loop over the matches and calculate the smoothed overlap.
|
||||
_cross = numpy.full(match_indxs[i].size, numpy.nan, numpy.float32)
|
||||
|
@ -327,13 +383,26 @@ def cosine_similarity(x, y):
|
|||
return out
|
||||
|
||||
|
||||
class ParticleOverlap:
|
||||
class ParticleOverlap(BaseMatcher):
|
||||
r"""
|
||||
Class to calculate halo overlaps. The density field calculation is based on
|
||||
the nearest grid position particle assignment scheme, with optional
|
||||
Gaussian smoothing.
|
||||
Halo overlaps calculator. The density field calculation is based on the
|
||||
nearest grid position particle assignment scheme, with optional Gaussian
|
||||
smoothing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
box_size : int
|
||||
Number of cells in the box.
|
||||
bckg_halfsize : int
|
||||
Number of to each side of the centre of the box to calculate the
|
||||
density field. This is because in CSiBORG we are only interested in the
|
||||
high-resolution region.
|
||||
"""
|
||||
|
||||
def __init__(self, box_size, bckg_halfsize):
|
||||
self.box_size = box_size
|
||||
self.bckg_halfsize = bckg_halfsize
|
||||
|
||||
def make_bckg_delta(self, particles, halo_map, hid2map, halo_cat,
|
||||
delta=None, verbose=False):
|
||||
"""
|
||||
|
@ -363,8 +432,8 @@ class ParticleOverlap:
|
|||
-------
|
||||
delta : 3-dimensional array
|
||||
"""
|
||||
cellmin = BOX_SIZE // 2 - BCKG_HALFSIZE
|
||||
cellmax = BOX_SIZE // 2 + BCKG_HALFSIZE
|
||||
cellmin = self.box_size // 2 - self.bckg_halfsize
|
||||
cellmax = self.box_size // 2 + self.bckg_halfsize
|
||||
ncells = cellmax - cellmin
|
||||
# We then pre-allocate the density field/check it is of the right shape
|
||||
if delta is None:
|
||||
|
@ -379,7 +448,7 @@ class ParticleOverlap:
|
|||
continue
|
||||
|
||||
pos, mass = pos[:, :3], pos[:, 3]
|
||||
pos = pos2cell(pos, BOX_SIZE)
|
||||
pos = pos2cell(pos, self.box_size)
|
||||
# We mask out particles outside the cubical high-resolution region
|
||||
mask = numpy.all((cellmin <= pos) & (pos < cellmax), axis=1)
|
||||
pos = pos[mask]
|
||||
|
@ -413,7 +482,7 @@ class ParticleOverlap:
|
|||
delta : 3-dimensional array
|
||||
"""
|
||||
nshift = read_nshift(smooth_kwargs)
|
||||
cells = pos2cell(pos, BOX_SIZE)
|
||||
cells = pos2cell(pos, self.box_size)
|
||||
# Check that minima and maxima are integers
|
||||
if not (mins is None and maxs is None):
|
||||
assert mins.dtype.char in numpy.typecodes["AllInteger"]
|
||||
|
@ -421,12 +490,12 @@ class ParticleOverlap:
|
|||
|
||||
if subbox:
|
||||
if mins is None or maxs is None:
|
||||
mins, maxs = get_halolims(cells, BOX_SIZE, nshift)
|
||||
mins, maxs = get_halolims(cells, self.box_size, nshift)
|
||||
|
||||
ncells = maxs - mins + 1 # To get the number of cells
|
||||
else:
|
||||
mins = [0, 0, 0]
|
||||
ncells = (BOX_SIZE, ) * 3
|
||||
ncells = (self.box_size, ) * 3
|
||||
|
||||
# Preallocate and fill the array
|
||||
delta = numpy.zeros(ncells, dtype=numpy.float32)
|
||||
|
@ -472,8 +541,8 @@ class ParticleOverlap:
|
|||
Calculated only if no smoothing is applied, otherwise `None`.
|
||||
"""
|
||||
nshift = read_nshift(smooth_kwargs)
|
||||
pos1 = pos2cell(pos1, BOX_SIZE)
|
||||
pos2 = pos2cell(pos2, BOX_SIZE)
|
||||
pos1 = pos2cell(pos1, self.box_size)
|
||||
pos2 = pos2cell(pos2, self.box_size)
|
||||
xc1, yc1, zc1 = [pos1[:, i] for i in range(3)]
|
||||
xc2, yc2, zc2 = [pos2[:, i] for i in range(3)]
|
||||
|
||||
|
@ -490,7 +559,7 @@ class ParticleOverlap:
|
|||
ymax = max(numpy.max(yc1), numpy.max(yc2)) + nshift
|
||||
zmax = max(numpy.max(zc1), numpy.max(zc2)) + nshift
|
||||
# Make sure shifting does not go beyond boundaries
|
||||
xmax, ymax, zmax = [min(px, BOX_SIZE - 1)
|
||||
xmax, ymax, zmax = [min(px, self.box_size - 1)
|
||||
for px in (xmax, ymax, zmax)]
|
||||
else:
|
||||
xmin, ymin, zmin = [min(mins1[i], mins2[i]) for i in range(3)]
|
||||
|
@ -573,12 +642,14 @@ class ParticleOverlap:
|
|||
smooth_kwargs=smooth_kwargs)
|
||||
|
||||
if smooth_kwargs is not None:
|
||||
return calculate_overlap(delta1, delta2, cellmins, delta_bckg)
|
||||
return calculate_overlap(delta1, delta2, cellmins, delta_bckg,
|
||||
self.box_size, self.bckg_halfsize)
|
||||
# Calculate masses not given
|
||||
totmass1 = numpy.sum(mass1) if totmass1 is None else totmass1
|
||||
totmass2 = numpy.sum(mass2) if totmass2 is None else totmass2
|
||||
return calculate_overlap_indxs(
|
||||
delta1, delta2, cellmins, delta_bckg, nonzero, totmass1, totmass2)
|
||||
return calculate_overlap_indxs(delta1, delta2, cellmins, delta_bckg,
|
||||
nonzero, totmass1, totmass2,
|
||||
self.box_size, self.bckg_halfsize)
|
||||
|
||||
|
||||
###############################################################################
|
||||
|
@ -725,7 +796,8 @@ def get_halolims(pos, ncells, nshift=None):
|
|||
|
||||
|
||||
@jit(nopython=True)
|
||||
def calculate_overlap(delta1, delta2, cellmins, delta_bckg):
|
||||
def calculate_overlap(delta1, delta2, cellmins, delta_bckg, box_size,
|
||||
bckg_halfsize):
|
||||
r"""
|
||||
Overlap between two halos whose density fields are evaluated on the
|
||||
same grid. This is a JIT implementation, hence it is outside of the main
|
||||
|
@ -743,6 +815,12 @@ def calculate_overlap(delta1, delta2, cellmins, delta_bckg):
|
|||
Summed background density field of the reference and cross simulations
|
||||
calculated with particles assigned to halos at the final snapshot.
|
||||
Assumed to only be sampled in cells :math:`[512, 1536)^3`.
|
||||
box_size : int
|
||||
Number of cells in the box.
|
||||
bckg_halfsize : int
|
||||
Number of to each side of the centre of the box to calculate the
|
||||
density field. This is because in CSiBORG we are only interested in the
|
||||
high-resolution region.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
@ -751,8 +829,8 @@ def calculate_overlap(delta1, delta2, cellmins, delta_bckg):
|
|||
totmass = 0.0 # Total mass of halo 1 and halo 2
|
||||
intersect = 0.0 # Weighted intersecting mass
|
||||
i0, j0, k0 = cellmins # Unpack things
|
||||
bckg_size = 2 * BCKG_HALFSIZE
|
||||
bckg_offset = BOX_SIZE // 2 - BCKG_HALFSIZE
|
||||
bckg_size = 2 * bckg_halfsize
|
||||
bckg_offset = box_size // 2 - bckg_halfsize
|
||||
imax, jmax, kmax = delta1.shape
|
||||
|
||||
for i in range(imax):
|
||||
|
@ -777,7 +855,7 @@ def calculate_overlap(delta1, delta2, cellmins, delta_bckg):
|
|||
|
||||
@jit(nopython=True)
|
||||
def calculate_overlap_indxs(delta1, delta2, cellmins, delta_bckg, nonzero,
|
||||
mass1, mass2):
|
||||
mass1, mass2, box_size, bckg_halfsize):
|
||||
r"""
|
||||
Overlap between two haloes whose density fields are evaluated on the
|
||||
same grid and `nonzero1` enumerates the non-zero cells of `delta1. This is
|
||||
|
@ -801,6 +879,12 @@ def calculate_overlap_indxs(delta1, delta2, cellmins, delta_bckg, nonzero,
|
|||
mass1, mass2 : floats, optional
|
||||
Total masses of the two haloes, respectively. Optional. If not provided
|
||||
calculcated directly from the density field.
|
||||
box_size : int
|
||||
Number of cells in the box.
|
||||
bckg_halfsize : int
|
||||
Number of to each side of the centre of the box to calculate the
|
||||
density field. This is because in CSiBORG we are only interested in the
|
||||
high-resolution region.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
@ -808,8 +892,8 @@ def calculate_overlap_indxs(delta1, delta2, cellmins, delta_bckg, nonzero,
|
|||
"""
|
||||
intersect = 0.0 # Weighted intersecting mass
|
||||
i0, j0, k0 = cellmins # Unpack cell minimas
|
||||
bckg_size = 2 * BCKG_HALFSIZE
|
||||
bckg_offset = BOX_SIZE // 2 - BCKG_HALFSIZE
|
||||
bckg_size = 2 * bckg_halfsize
|
||||
bckg_offset = box_size // 2 - bckg_halfsize
|
||||
|
||||
for n in range(nonzero.shape[0]):
|
||||
i, j, k = nonzero[n, :]
|
||||
|
|
|
@ -23,8 +23,8 @@ from .overlap_summary import (NPairsOverlap, PairOverlap, # noqa
|
|||
binned_resample_mean, get_cross_sims) # noqa
|
||||
from .paths import Paths # noqa
|
||||
from .pk_summary import PKReader # noqa
|
||||
from .readsim import (MmainReader, ParticleReader, halfwidth_mask, # noqa
|
||||
load_halo_particles, read_initcm) # noqa
|
||||
from .readsim import (MmainReader, CSiBORGReader, QuijoteReader, halfwidth_mask, # noqa
|
||||
load_halo_particles) # noqa
|
||||
from .tpcf_summary import TPCFReader # noqa
|
||||
from .utils import (M200_to_R200, cartesian_to_radec, # noqa
|
||||
cols_to_structured, radec_to_cartesian, read_h5,
|
||||
|
|
|
@ -21,9 +21,9 @@ import numpy
|
|||
from astropy import constants, units
|
||||
from astropy.cosmology import LambdaCDM
|
||||
|
||||
from .readsim import ParticleReader
|
||||
from .readsim import CSiBORGReader, QuijoteReader
|
||||
|
||||
|
||||
# Map of CSiBORG unit conversions
|
||||
CSIBORG_CONV_NAME = {
|
||||
"length": ["x", "y", "z", "peak_x", "peak_y", "peak_z", "Rs", "rmin",
|
||||
"rmax", "r200c", "r500c", "r200m", "r500m", "x0", "y0", "z0",
|
||||
|
@ -31,8 +31,14 @@ CSIBORG_CONV_NAME = {
|
|||
"velocity": ["vx", "vy", "vz"],
|
||||
"mass": ["mass_cl", "totpartmass", "m200c", "m500c", "mass_mmain", "M",
|
||||
"m200m", "m500m"],
|
||||
"density": ["rho0"]}
|
||||
"density": ["rho0"]
|
||||
}
|
||||
|
||||
QUIJOTE_CONV_NAME = {
|
||||
"length": ["x", "y", "z", "x0", "y0", "z0", "Rs", "r200c", "r500c",
|
||||
"r200m", "r500m", "lagpatch_size"],
|
||||
"mass": ["group_mass", "totpartmass", "m200c", "m500c", "m200m", "m500m"],
|
||||
}
|
||||
|
||||
###############################################################################
|
||||
# Base box #
|
||||
|
@ -74,7 +80,7 @@ class BaseBox(ABC):
|
|||
@property
|
||||
def h(self):
|
||||
r"""
|
||||
The little 'h` parameter at the time of the snapshot.
|
||||
The little 'h' parameter at the time of the snapshot.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
@ -157,12 +163,12 @@ class CSiBORGBox(BaseBox):
|
|||
"""
|
||||
Read in the snapshot info file and set the units from it.
|
||||
"""
|
||||
partreader = ParticleReader(paths)
|
||||
partreader = CSiBORGReader(paths)
|
||||
info = partreader.read_info(nsnap, nsim)
|
||||
pars = ["boxlen", "time", "aexp", "H0", "omega_m", "omega_l",
|
||||
"omega_k", "omega_b", "unit_l", "unit_d", "unit_t"]
|
||||
for par in pars:
|
||||
setattr(self, "_" + par, float(info[par]))
|
||||
setattr(self, "_" + par, info[par])
|
||||
|
||||
self._cosmo = LambdaCDM(H0=self._H0, Om0=self._omega_m,
|
||||
Ode0=self._omega_l, Tcmb0=2.725 * units.K,
|
||||
|
@ -226,7 +232,7 @@ class CSiBORGBox(BaseBox):
|
|||
|
||||
Returns
|
||||
-------
|
||||
length : foat
|
||||
length : float
|
||||
Length in :math:`\mathrm{ckpc}`
|
||||
"""
|
||||
return length * (self._unit_l / units.kpc.to(units.cm) / self._aexp)
|
||||
|
@ -243,7 +249,7 @@ class CSiBORGBox(BaseBox):
|
|||
|
||||
Returns
|
||||
-------
|
||||
length : foat
|
||||
length : float
|
||||
Length in box units.
|
||||
"""
|
||||
return length / (self._unit_l / units.kpc.to(units.cm) / self._aexp)
|
||||
|
@ -260,7 +266,7 @@ class CSiBORGBox(BaseBox):
|
|||
|
||||
Returns
|
||||
-------
|
||||
length : foat
|
||||
length : float
|
||||
Length in box units.
|
||||
"""
|
||||
return self.kpc2box(length * 1e3)
|
||||
|
@ -277,7 +283,7 @@ class CSiBORGBox(BaseBox):
|
|||
|
||||
Returns
|
||||
-------
|
||||
length : foat
|
||||
length : float
|
||||
Length in :math:`\mathrm{ckpc}`
|
||||
"""
|
||||
return self.box2kpc(length) * 1e-3
|
||||
|
@ -419,17 +425,110 @@ class QuijoteBox(BaseBox):
|
|||
Empty keyword arguments. For backwards compatibility.
|
||||
"""
|
||||
|
||||
def __init__(self, nsnap, **kwargs):
|
||||
def __init__(self, nsnap, nsim, paths):
|
||||
zdict = {4: 0.0, 3: 0.5, 2: 1.0, 1: 2.0, 0: 3.0}
|
||||
assert nsnap in zdict.keys(), f"`nsnap` must be in {zdict.keys()}."
|
||||
self._aexp = 1 / (1 + zdict[nsnap])
|
||||
|
||||
self._cosmo = LambdaCDM(H0=67.11, Om0=0.3175, Ode0=0.6825,
|
||||
Tcmb0=2.725 * units.K, Ob0=0.049)
|
||||
info = QuijoteReader(paths).read_info(nsnap, nsim)
|
||||
self._cosmo = LambdaCDM(H0=info["Hubble"], Om0=info["Omega_m"],
|
||||
Ode0=info["Omega_l"], Tcmb0=2.725 * units.K)
|
||||
self._info = info
|
||||
|
||||
@property
|
||||
def boxsize(self):
|
||||
return 1000. / (self._cosmo.H0.value / 100)
|
||||
return self._info["BoxSize"]
|
||||
|
||||
def box2mpc(self, length):
|
||||
r"""
|
||||
Convert length from box units to :math:`\mathrm{cMpc} / h`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
length : float
|
||||
Length in box units.
|
||||
|
||||
Returns
|
||||
-------
|
||||
length : float
|
||||
Length in :math:`\mathrm{cMpc} / h`
|
||||
"""
|
||||
return length * self.boxsize
|
||||
|
||||
def mpc2box(self, length):
|
||||
r"""
|
||||
Convert length from :math:`\mathrm{cMpc} / h` to box units.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
length : float
|
||||
Length in :math:`\mathrm{cMpc} / h`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
length : float
|
||||
Length in box units.
|
||||
"""
|
||||
return length / self.boxsize
|
||||
|
||||
def solarmass2box(self, mass):
|
||||
r"""
|
||||
Convert mass from :math:`M_\odot / h` to box units.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mass : float
|
||||
Mass in :math:`M_\odot`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mass : float
|
||||
Mass in box units.
|
||||
"""
|
||||
return mass / self._info["TotMass"]
|
||||
|
||||
def box2solarmass(self, mass):
|
||||
r"""
|
||||
Convert mass from box units to :math:`M_\odot / h`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mass : float
|
||||
Mass in box units.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mass : float
|
||||
Mass in :math:`M_\odot / h`.
|
||||
"""
|
||||
return mass * self._info["TotMass"]
|
||||
|
||||
def convert_from_box(self, data, names):
|
||||
raise NotImplementedError("Conversion not implemented for Quijote.")
|
||||
names = [names] if isinstance(names, str) else names
|
||||
transforms = {"length": self.box2mpc,
|
||||
"mass": self.box2solarmass,
|
||||
# "velocity": self.box2vel,
|
||||
# "density": self.box2dens,
|
||||
}
|
||||
|
||||
for name in names:
|
||||
if name not in data.dtype.names:
|
||||
continue
|
||||
|
||||
# Convert
|
||||
found = False
|
||||
for unittype, suppnames in QUIJOTE_CONV_NAME.items():
|
||||
if name in suppnames:
|
||||
data[name] = transforms[unittype](data[name])
|
||||
found = True
|
||||
continue
|
||||
# If nothing found
|
||||
if not found:
|
||||
raise NotImplementedError(
|
||||
f"Conversion of `{name}` is not defined.")
|
||||
|
||||
# # Center at the observer
|
||||
# if name in ["x0", "y0", "z0"]:
|
||||
# data[name] -= transforms["length"](0.5)
|
||||
|
||||
return data
|
||||
|
|
|
@ -22,7 +22,6 @@ from copy import deepcopy
|
|||
from functools import lru_cache
|
||||
from itertools import product
|
||||
from math import floor
|
||||
from os.path import join
|
||||
|
||||
import numpy
|
||||
from readfof import FoF_catalog
|
||||
|
@ -30,14 +29,14 @@ from sklearn.neighbors import NearestNeighbors
|
|||
|
||||
from .box_units import CSiBORGBox, QuijoteBox
|
||||
from .paths import Paths
|
||||
from .readsim import ParticleReader
|
||||
from .readsim import CSiBORGReader
|
||||
from .utils import (add_columns, cartesian_to_radec, cols_to_structured,
|
||||
flip_cols, radec_to_cartesian, real2redshift)
|
||||
|
||||
|
||||
class BaseCatalogue(ABC):
|
||||
"""
|
||||
Base (sub)halo catalogue.
|
||||
Base halo catalogue.
|
||||
"""
|
||||
_data = None
|
||||
_paths = None
|
||||
|
@ -125,11 +124,8 @@ class BaseCatalogue(ABC):
|
|||
|
||||
def position(self, in_initial=False, cartesian=True):
|
||||
r"""
|
||||
Position components. If Cartesian, then in :math:`\mathrm{cMpc}`. If
|
||||
spherical, then radius is in :math:`\mathrm{cMpc}`, RA in
|
||||
:math:`[0, 360)` degrees and DEC in :math:`[-90, 90]` degrees. Note
|
||||
that the position is defined as the minimum of the gravitationl
|
||||
potential.
|
||||
Position components. If not Cartesian, then RA is in :math:`[0, 360)`
|
||||
degrees and DEC is in :math:`[-90, 90]` degrees.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
@ -399,33 +395,33 @@ class CSiBORGHaloCatalogue(BaseCatalogue):
|
|||
names and the items are a len-2 tuple of (min, max) values. In case of
|
||||
no minimum or maximum, use `None`. For radial distance from the origin
|
||||
use `dist`.
|
||||
with_lagpatch : bool, optional
|
||||
Whether to only load halos with a resolved Lagrangian patch.
|
||||
load_fitted : bool, optional
|
||||
Whether to load fitted quantities.
|
||||
load_initial : bool, optional
|
||||
Whether to load initial positions.
|
||||
with_lagpatch : bool, optional
|
||||
Whether to only load halos with a resolved Lagrangian patch.
|
||||
rawdata : bool, optional
|
||||
Whether to return the raw data. In this case applies no cuts and
|
||||
transformations.
|
||||
"""
|
||||
|
||||
def __init__(self, nsim, paths, bounds={"dist": (0, 155.5 / 0.705)},
|
||||
with_lagpatch=True, load_fitted=True, load_initial=True,
|
||||
load_fitted=True, load_initial=True, with_lagpatch=True,
|
||||
rawdata=False):
|
||||
self.nsim = nsim
|
||||
self.paths = paths
|
||||
reader = ParticleReader(paths)
|
||||
reader = CSiBORGReader(paths)
|
||||
self._data = reader.read_fof_halos(self.nsim)
|
||||
|
||||
if load_fitted:
|
||||
fits = numpy.load(paths.structfit(self.nsnap, nsim))
|
||||
fits = numpy.load(paths.structfit(self.nsnap, nsim, "csiborg"))
|
||||
cols = [col for col in fits.dtype.names if col != "index"]
|
||||
X = [fits[col] for col in cols]
|
||||
self._data = add_columns(self._data, X, cols)
|
||||
|
||||
if load_initial:
|
||||
fits = numpy.load(paths.initmatch(nsim, "fit"))
|
||||
fits = numpy.load(paths.initmatch(nsim, "csiborg", "fit"))
|
||||
X, cols = [], []
|
||||
for col in fits.dtype.names:
|
||||
if col == "index":
|
||||
|
@ -465,7 +461,7 @@ class CSiBORGHaloCatalogue(BaseCatalogue):
|
|||
|
||||
@property
|
||||
def nsnap(self):
|
||||
return max(self.paths.get_snapshots(self.nsim))
|
||||
return max(self.paths.get_snapshots(self.nsim, "csiborg"))
|
||||
|
||||
@property
|
||||
def box(self):
|
||||
|
@ -497,28 +493,35 @@ class QuijoteHaloCatalogue(BaseCatalogue):
|
|||
nsnap : int
|
||||
Snapshot index.
|
||||
origin : len-3 tuple, optional
|
||||
Where to place the origin of the box. By default the centre of the box.
|
||||
In units of :math:`cMpc`.
|
||||
Where to place the origin of the box. In units of :math:`cMpc / h`.
|
||||
bounds : dict
|
||||
Parameter bounds to apply to the catalogue. The keys are the parameter
|
||||
names and the items are a len-2 tuple of (min, max) values. In case of
|
||||
no minimum or maximum, use `None`. For radial distance from the origin
|
||||
use `dist`.
|
||||
load_initial : bool, optional
|
||||
Whether to load initial positions.
|
||||
with_lagpatch : bool, optional
|
||||
Whether to only load halos with a resolved Lagrangian patch.
|
||||
rawdata : bool, optional
|
||||
Whether to return the raw data. In this case applies no cuts and
|
||||
transformations.
|
||||
**kwargs : dict
|
||||
Keyword arguments for backward compatibility.
|
||||
"""
|
||||
_nsnap = None
|
||||
_origin = None
|
||||
|
||||
def __init__(self, nsim, paths, nsnap,
|
||||
origin=[500 / 0.6711, 500 / 0.6711, 500 / 0.6711],
|
||||
bounds=None, **kwargs):
|
||||
def __init__(self, nsim, paths, nsnap, origin=[0., 0., 0.],
|
||||
bounds=None, load_initial=True, with_lagpatch=True,
|
||||
rawdata=False, **kwargs):
|
||||
self.paths = paths
|
||||
self.nsnap = nsnap
|
||||
self.origin = origin
|
||||
self._boxwidth = 1000 / 0.6711
|
||||
self._box = QuijoteBox(nsnap, nsim, paths)
|
||||
self._boxwidth = self.box.boxsize
|
||||
|
||||
fpath = join(self.paths.quijote_dir, "halos", str(nsim))
|
||||
fpath = self.paths.fof_cat(nsim, "quijote")
|
||||
fof = FoF_catalog(fpath, self.nsnap, long_ids=False, swap=False,
|
||||
SFR=False, read_IDs=False)
|
||||
|
||||
|
@ -529,18 +532,44 @@ class QuijoteHaloCatalogue(BaseCatalogue):
|
|||
("index", numpy.int32)]
|
||||
data = cols_to_structured(fof.GroupLen.size, cols)
|
||||
|
||||
pos = fof.GroupPos / 1e3 / self.box.h
|
||||
pos = self.box.mpc2box(fof.GroupPos / 1e3)
|
||||
vel = fof.GroupVel * (1 + self.redshift)
|
||||
for i, p in enumerate(["x", "y", "z"]):
|
||||
data[p] = pos[:, i] - self.origin[i]
|
||||
data["v" + p] = vel[:, i]
|
||||
data["group_mass"] = fof.GroupMass * 1e10 / self.box.h
|
||||
data["group_mass"] = self.box.solarmass2box(fof.GroupMass * 1e10)
|
||||
data["npart"] = fof.GroupLen
|
||||
data["index"] = numpy.arange(data.size, dtype=numpy.int32)
|
||||
# We want to start indexing from 1. Index 0 is reserved for
|
||||
# particles unassigned to any FoF group.
|
||||
data["index"] = 1 + numpy.arange(data.size, dtype=numpy.int32)
|
||||
|
||||
if load_initial:
|
||||
fits = numpy.load(paths.initmatch(nsim, "quijote", "fit"))
|
||||
X, cols = [], []
|
||||
for col in fits.dtype.names:
|
||||
if col == "index":
|
||||
continue
|
||||
if col in ['x', 'y', 'z']:
|
||||
cols.append(col + "0")
|
||||
else:
|
||||
cols.append(col)
|
||||
X.append(fits[col])
|
||||
data = add_columns(data, X, cols)
|
||||
|
||||
self._data = data
|
||||
if bounds is not None:
|
||||
self.apply_bounds(bounds)
|
||||
if not rawdata:
|
||||
if with_lagpatch:
|
||||
mask = numpy.isfinite(self._data["lagpatch_size"])
|
||||
self._data = self._data[mask]
|
||||
|
||||
names = ["x", "y", "z", "group_mass"]
|
||||
self._data = self.box.convert_from_box(self._data, names)
|
||||
if load_initial:
|
||||
names = ["x0", "y0", "z0", "lagpatch_size"]
|
||||
self._data = self.box.convert_from_box(self._data, names)
|
||||
|
||||
if bounds is not None:
|
||||
self.apply_bounds(bounds)
|
||||
|
||||
@property
|
||||
def nsnap(self):
|
||||
|
@ -579,7 +608,7 @@ class QuijoteHaloCatalogue(BaseCatalogue):
|
|||
-------
|
||||
box : instance of :py:class:`csiborgtools.units.BaseBox`
|
||||
"""
|
||||
return QuijoteBox(self.nsnap)
|
||||
return self._box
|
||||
|
||||
@property
|
||||
def origin(self):
|
||||
|
@ -629,6 +658,7 @@ class QuijoteHaloCatalogue(BaseCatalogue):
|
|||
cat.apply_bounds({"dist": (0, rmax)})
|
||||
return cat
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Utility functions for halo catalogues #
|
||||
###############################################################################
|
||||
|
|
|
@ -186,7 +186,7 @@ class Paths:
|
|||
"""
|
||||
return join(self.borg_dir, "mcmc", f"mcmc_{nsim}.h5")
|
||||
|
||||
def fof_membership(self, nsim, sorted=False):
|
||||
def fof_membership(self, nsim, simname, sorted=False):
|
||||
"""
|
||||
Path to the file containing the FoF particle membership.
|
||||
|
||||
|
@ -194,10 +194,15 @@ class Paths:
|
|||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
sorted : bool, optional
|
||||
Whether to return path to the file that is sorted in the same
|
||||
order as the PHEW output.
|
||||
"""
|
||||
assert simname in ["csiborg", "quijote"]
|
||||
if simname == "quijote":
|
||||
raise RuntimeError("Quijote FoF membership is in the FoF cats..")
|
||||
fdir = join(self.postdir, "FoF_membership", )
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
|
@ -207,20 +212,25 @@ class Paths:
|
|||
fout = fout.replace(".npy", "_sorted.npy")
|
||||
return fout
|
||||
|
||||
def fof_cat(self, nsim):
|
||||
"""
|
||||
Path to the FoF halo catalogue file.
|
||||
def fof_cat(self, nsim, simname):
|
||||
r"""
|
||||
Path to the :math:`z = 0` FoF halo catalogue.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
"""
|
||||
fdir = join(self.postdir, "FoF_membership", )
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
return join(fdir, f"halo_catalog_{nsim}_FOF.txt")
|
||||
assert simname in ["csiborg", "quijote"]
|
||||
if simname == "csiborg":
|
||||
fdir = join(self.postdir, "FoF_membership", )
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning)
|
||||
return join(fdir, f"halo_catalog_{nsim}_FOF.txt")
|
||||
return join(self.quijote_dir, "Halos_fiducial", str(nsim))
|
||||
|
||||
def mmain(self, nsnap, nsim):
|
||||
"""
|
||||
|
@ -244,7 +254,7 @@ class Paths:
|
|||
return join(fdir,
|
||||
f"mmain_{str(nsim).zfill(5)}_{str(nsnap).zfill(5)}.npz")
|
||||
|
||||
def initmatch(self, nsim, kind):
|
||||
def initmatch(self, nsim, simname, kind):
|
||||
"""
|
||||
Path to the `initmatch` files where the halo match between the
|
||||
initial and final snapshot of a CSiBORG realisaiton is stored.
|
||||
|
@ -253,19 +263,30 @@ class Paths:
|
|||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
kind : str
|
||||
Type of match. Must be one of `["particles", "fit", "halomap"]`.
|
||||
Type of match. Must be one of `particles` or `fit`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
path : str
|
||||
"""
|
||||
assert kind in ["particles", "fit", "halomap"]
|
||||
assert simname in ["csiborg", "quijote"]
|
||||
assert kind in ["particles", "fit"]
|
||||
ftype = "npy" if kind == "fit" else "h5"
|
||||
fdir = join(self.postdir, "initmatch")
|
||||
|
||||
if simname == "csiborg":
|
||||
fdir = join(self.postdir, "initmatch")
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning)
|
||||
return join(fdir, f"{kind}_{str(nsim).zfill(5)}.{ftype}")
|
||||
|
||||
fdir = join(self.quijote_dir, "initmatch")
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning)
|
||||
return join(fdir, f"{kind}_{str(nsim).zfill(5)}.{ftype}")
|
||||
|
||||
def get_ics(self, simname):
|
||||
|
@ -276,7 +297,7 @@ class Paths:
|
|||
Parameters
|
||||
----------
|
||||
simname : str
|
||||
Simulation name. Must be one of `["csiborg", "quijote"]`.
|
||||
Simulation name. Must be `csiborg` or `quijote`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
@ -295,75 +316,127 @@ class Paths:
|
|||
except ValueError:
|
||||
pass
|
||||
return numpy.sort(ids)
|
||||
else:
|
||||
# TODO here later read this from the catalogues instead.
|
||||
return numpy.arange(100, dtype=int)
|
||||
|
||||
def ic_path(self, nsim, tonew=False):
|
||||
files = glob("/mnt/extraspace/rstiskalek/Quijote/Snapshots_fiducial/*")
|
||||
files = [int(f.split("/")[-1]) for f in files]
|
||||
return numpy.sort(files)
|
||||
|
||||
def snapshots(self, nsim, simname, tonew=False):
|
||||
"""
|
||||
Path to a CSiBORG IC realisation folder.
|
||||
Path to an IC snapshots folder.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
tonew : bool, optional
|
||||
Whether to return the path to the '_new' IC realisation.
|
||||
Whether to return the path to the '_new' IC realisation of
|
||||
CSiBORG. Ignored for Quijote.
|
||||
|
||||
Returns
|
||||
-------
|
||||
path : str
|
||||
"""
|
||||
fname = "ramses_out_{}"
|
||||
if tonew:
|
||||
fname += "_new"
|
||||
return join(self.postdir, "output", fname.format(nsim))
|
||||
assert simname in ["csiborg", "quijote"]
|
||||
if simname == "csiborg":
|
||||
fname = "ramses_out_{}"
|
||||
if tonew:
|
||||
fname += "_new"
|
||||
return join(self.postdir, "output", fname.format(nsim))
|
||||
return join(self.srcdir, fname.format(nsim))
|
||||
|
||||
return join(self.srcdir, fname.format(nsim))
|
||||
return join(self.quijote_dir, "Snapshots_fiducial", str(nsim))
|
||||
|
||||
def get_snapshots(self, nsim):
|
||||
def get_snapshots(self, nsim, simname):
|
||||
"""
|
||||
List of available snapshots of a CSiBORG IC realisation.
|
||||
List of available snapshots of simulation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
snapshots : 1-dimensional array
|
||||
"""
|
||||
simpath = self.ic_path(nsim, tonew=False)
|
||||
# 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]
|
||||
simpath = self.snapshots(nsim, simname, tonew=False)
|
||||
if simname == "csiborg":
|
||||
# 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]
|
||||
else:
|
||||
snaps = glob(join(simpath, "snapdir_*"))
|
||||
snaps = [int(snap.split("/")[-1].split("snapdir_")[-1])
|
||||
for snap in snaps]
|
||||
return numpy.sort(snaps)
|
||||
|
||||
def snapshot(self, nsnap, nsim):
|
||||
def snapshot(self, nsnap, nsim, simname):
|
||||
"""
|
||||
Path to a CSiBORG IC realisation snapshot.
|
||||
Path to an IC realisation snapshot.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsnap : int
|
||||
Snapshot index.
|
||||
Snapshot index. For Quijote, `-1` indicates the IC snapshot.
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
snappath : str
|
||||
"""
|
||||
tonew = nsnap == 1
|
||||
simpath = self.ic_path(nsim, tonew=tonew)
|
||||
return join(simpath, f"output_{str(nsnap).zfill(5)}")
|
||||
simpath = self.snapshots(nsim, simname, tonew=nsnap == 1)
|
||||
if simname == "csiborg":
|
||||
return join(simpath, f"output_{str(nsnap).zfill(5)}")
|
||||
else:
|
||||
if nsnap == -1:
|
||||
return join(simpath, "ICs", "ics")
|
||||
nsnap = str(nsnap).zfill(3)
|
||||
return join(simpath, f"snapdir_{nsnap}", f"snap_{nsnap}")
|
||||
|
||||
def structfit(self, nsnap, nsim):
|
||||
def particles(self, nsim, simname):
|
||||
"""
|
||||
Path to the clump or halo catalogue from `fit_halos.py`. Only CSiBORG
|
||||
is supported.
|
||||
Path to the files containing all particles of a CSiBORG realisation at
|
||||
:math:`z = 0`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
path : str
|
||||
"""
|
||||
assert simname in ["csiborg", "quijote"]
|
||||
if simname == "csiborg":
|
||||
fdir = join(self.postdir, "particles")
|
||||
if not isdir(fdir):
|
||||
makedirs(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning)
|
||||
fname = f"parts_{str(nsim).zfill(5)}.h5"
|
||||
return join(fdir, fname)
|
||||
|
||||
fdir = join(self.quijote_dir, "Particles_fiducial")
|
||||
if not isdir(fdir):
|
||||
makedirs(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning)
|
||||
fname = f"parts_{str(nsim).zfill(5)}.h5"
|
||||
return join(fdir, fname)
|
||||
|
||||
def structfit(self, nsnap, nsim, simname):
|
||||
"""
|
||||
Path to the halo catalogue from `fit_halos.py`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
@ -371,15 +444,26 @@ class Paths:
|
|||
Snapshot index.
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
simname : str
|
||||
Simulation name. Must be one of `csiborg` or `quijote`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
path : str
|
||||
"""
|
||||
fdir = join(self.postdir, "structfit")
|
||||
assert simname in ["csiborg", "quijote"]
|
||||
if simname == "csiborg":
|
||||
fdir = join(self.postdir, "structfit")
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
fname = f"out_{str(nsim).zfill(5)}_{str(nsnap).zfill(5)}.npy"
|
||||
return join(fdir, fname)
|
||||
|
||||
fdir = join(self.quijote_dir, "structfit")
|
||||
if not isdir(fdir):
|
||||
mkdir(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning)
|
||||
fname = f"out_{str(nsim).zfill(5)}_{str(nsnap).zfill(5)}.npy"
|
||||
return join(fdir, fname)
|
||||
|
||||
|
@ -409,27 +493,6 @@ class Paths:
|
|||
fname = fname.replace("overlap", "overlap_smoothed")
|
||||
return join(fdir, fname)
|
||||
|
||||
def particles(self, nsim):
|
||||
"""
|
||||
Path to the files containing all particles of a CSiBORG realisation at
|
||||
:math:`z = 0`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
|
||||
Returns
|
||||
-------
|
||||
path : str
|
||||
"""
|
||||
fdir = join(self.postdir, "particles")
|
||||
if not isdir(fdir):
|
||||
makedirs(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
fname = f"parts_{str(nsim).zfill(5)}.h5"
|
||||
return join(fdir, fname)
|
||||
|
||||
def field(self, kind, MAS, grid, nsim, in_rsp, smooth_scale=None):
|
||||
r"""
|
||||
Path to the files containing the calculated density fields in CSiBORG.
|
||||
|
|
|
@ -15,34 +15,28 @@
|
|||
"""
|
||||
Functions to read in the particle and clump files.
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from gc import collect
|
||||
from os.path import isfile, join
|
||||
from warnings import warn
|
||||
|
||||
import numpy
|
||||
import readfof
|
||||
import readgadget
|
||||
from scipy.io import FortranFile
|
||||
from tqdm import tqdm, trange
|
||||
|
||||
from .paths import Paths
|
||||
from .utils import cols_to_structured
|
||||
|
||||
###############################################################################
|
||||
# Fortran particle reader #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class ParticleReader:
|
||||
class BaseReader(ABC):
|
||||
"""
|
||||
Object to read in particle files along with their corresponding haloes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
paths : py:class`csiborgtools.read.Paths`
|
||||
Base class for all readers.
|
||||
"""
|
||||
_paths = None
|
||||
|
||||
def __init__(self, paths):
|
||||
self.paths = paths
|
||||
|
||||
@property
|
||||
def paths(self):
|
||||
"""
|
||||
|
@ -59,9 +53,10 @@ class ParticleReader:
|
|||
assert isinstance(paths, Paths)
|
||||
self._paths = paths
|
||||
|
||||
@abstractmethod
|
||||
def read_info(self, nsnap, nsim):
|
||||
"""
|
||||
Read CSiBORG simulation snapshot info.
|
||||
Read simulation snapshot info.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
@ -73,10 +68,63 @@ class ParticleReader:
|
|||
Returns
|
||||
-------
|
||||
info : dict
|
||||
Dictionary of information paramaters. Note that both keys and
|
||||
values are strings.
|
||||
Dictionary of information paramaters.
|
||||
"""
|
||||
snappath = self.paths.snapshot(nsnap, nsim)
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def read_particle(self, nsnap, nsim, pars_extract, return_structured=True,
|
||||
verbose=True):
|
||||
"""
|
||||
Read particle files of a simulation at a given snapshot and return
|
||||
values of `pars_extract`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsnap : int
|
||||
Snapshot index.
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
pars_extract : list of str
|
||||
Parameters to be extracted.
|
||||
return_structured : bool, optional
|
||||
Whether to return a structured array or a 2-dimensional array. If
|
||||
the latter, then the order of the columns is the same as the order
|
||||
of `pars_extract`. However, enforces single-precision floating
|
||||
point format for all columns.
|
||||
verbose : bool, optional
|
||||
Verbosity flag while for reading in the files.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : structured array or 2-dimensional array
|
||||
Particle information.
|
||||
pids : 1-dimensional array
|
||||
Particle IDs.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
###############################################################################
|
||||
# CSiBORG particle reader #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class CSiBORGReader:
|
||||
"""
|
||||
Object to read in CSiBORG snapshots from the binary files and halo
|
||||
catalogues.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
paths : py:class`csiborgtools.read.Paths`
|
||||
"""
|
||||
|
||||
def __init__(self, paths):
|
||||
self.paths = paths
|
||||
|
||||
def read_info(self, nsnap, nsim):
|
||||
snappath = self.paths.snapshot(nsnap, nsim, "csiborg")
|
||||
filename = join(snappath, "info_{}.txt".format(str(nsnap).zfill(5)))
|
||||
with open(filename, "r") as f:
|
||||
info = f.read().split()
|
||||
|
@ -87,7 +135,7 @@ class ParticleReader:
|
|||
|
||||
keys = info[eqs - 1]
|
||||
vals = info[eqs + 1]
|
||||
return {key: val for key, val in zip(keys, vals)}
|
||||
return {key: convert_str_to_num(val) for key, val in zip(keys, vals)}
|
||||
|
||||
def open_particle(self, nsnap, nsim, verbose=True):
|
||||
"""
|
||||
|
@ -109,7 +157,7 @@ class ParticleReader:
|
|||
partfiles : list of `scipy.io.FortranFile`
|
||||
Opened part files.
|
||||
"""
|
||||
snappath = self.paths.snapshot(nsnap, nsim)
|
||||
snappath = self.paths.snapshot(nsnap, nsim, "csiborg")
|
||||
ncpu = int(self.read_info(nsnap, nsim)["ncpu"])
|
||||
nsnap = str(nsnap).zfill(5)
|
||||
if verbose:
|
||||
|
@ -192,33 +240,6 @@ class ParticleReader:
|
|||
|
||||
def read_particle(self, nsnap, nsim, pars_extract, return_structured=True,
|
||||
verbose=True):
|
||||
"""
|
||||
Read particle files of a simulation at a given snapshot and return
|
||||
values of `pars_extract`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsnap : int
|
||||
Snapshot index.
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
pars_extract : list of str
|
||||
Parameters to be extracted.
|
||||
return_structured : bool, optional
|
||||
Whether to return a structured array or a 2-dimensional array. If
|
||||
the latter, then the order of the columns is the same as the order
|
||||
of `pars_extract`. However, enforces single-precision floating
|
||||
point format for all columns.
|
||||
verbose : bool, optional
|
||||
Verbosity flag while for reading the CPU outputs.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : structured array or 2-dimensional array
|
||||
Particle information.
|
||||
pids : 1-dimensional array
|
||||
Particle IDs.
|
||||
"""
|
||||
# Open the particle files
|
||||
nparts, partfiles = self.open_particle(nsnap, nsim, verbose=verbose)
|
||||
if verbose:
|
||||
|
@ -299,11 +320,11 @@ class ParticleReader:
|
|||
"""
|
||||
nsnap = str(nsnap).zfill(5)
|
||||
cpu = str(cpu + 1).zfill(5)
|
||||
fpath = join(self.paths.ic_path(nsim, tonew=False), f"output_{nsnap}",
|
||||
f"unbinding_{nsnap}.out{cpu}")
|
||||
fpath = join(self.paths.snapshots(nsim, "csiborg", tonew=False),
|
||||
f"output_{nsnap}", f"unbinding_{nsnap}.out{cpu}")
|
||||
return FortranFile(fpath)
|
||||
|
||||
def read_clumpid(self, nsnap, nsim, verbose=True):
|
||||
def read_phew_clumpid(self, nsnap, nsim, verbose=True):
|
||||
"""
|
||||
Read PHEW clump IDs of particles from unbinding files. This halo finder
|
||||
was used when running the catalogue.
|
||||
|
@ -337,7 +358,7 @@ class ParticleReader:
|
|||
|
||||
return clumpid
|
||||
|
||||
def read_clumps(self, nsnap, nsim, cols=None):
|
||||
def read_phew_clups(self, nsnap, nsim, cols=None):
|
||||
"""
|
||||
Read in a PHEW clump file `clump_xxXXX.dat`.
|
||||
|
||||
|
@ -356,7 +377,7 @@ class ParticleReader:
|
|||
out : structured array
|
||||
"""
|
||||
nsnap = str(nsnap).zfill(5)
|
||||
fname = join(self.paths.ic_path(nsim, tonew=False),
|
||||
fname = join(self.paths.snapshots(nsim, "csiborg", tonew=False),
|
||||
"output_{}".format(nsnap),
|
||||
"clump_{}.dat".format(nsnap))
|
||||
if not isfile(fname):
|
||||
|
@ -387,7 +408,7 @@ class ParticleReader:
|
|||
out[col] = data[:, clump_cols[col][0]]
|
||||
return out
|
||||
|
||||
def read_fof_hids(self, nsim):
|
||||
def read_fof_hids(self, nsim, **kwargs):
|
||||
"""
|
||||
Read in the FoF particle halo membership IDs that are sorted to match
|
||||
the PHEW output.
|
||||
|
@ -396,13 +417,16 @@ class ParticleReader:
|
|||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
**kwargs : dict
|
||||
Keyword arguments for backward compatibility.
|
||||
|
||||
Returns
|
||||
-------
|
||||
hids : 1-dimensional array
|
||||
Halo IDs of particles.
|
||||
"""
|
||||
return numpy.load(self.paths.fof_membership(nsim, sorted=True))
|
||||
return numpy.load(self.paths.fof_membership(nsim, "csiborg",
|
||||
sorted=True))
|
||||
|
||||
def read_fof_halos(self, nsim):
|
||||
"""
|
||||
|
@ -417,7 +441,7 @@ class ParticleReader:
|
|||
-------
|
||||
cat : structured array
|
||||
"""
|
||||
fpath = self.paths.fof_cat(nsim)
|
||||
fpath = self.paths.fof_cat(nsim, "csiborg")
|
||||
hid = numpy.genfromtxt(fpath, usecols=0, dtype=numpy.int32)
|
||||
pos = numpy.genfromtxt(fpath, usecols=(1, 2, 3), dtype=numpy.float32)
|
||||
totmass = numpy.genfromtxt(fpath, usecols=4, dtype=numpy.float32)
|
||||
|
@ -437,7 +461,7 @@ class ParticleReader:
|
|||
|
||||
|
||||
###############################################################################
|
||||
# Summed substructure catalogue #
|
||||
# Summed substructure PHEW catalogue for CSiBORG #
|
||||
###############################################################################
|
||||
|
||||
|
||||
|
@ -467,8 +491,8 @@ class MmainReader:
|
|||
Parameters
|
||||
----------
|
||||
clumparr : structured array
|
||||
Clump array. Read from `ParticleReader.read_clumps`. Must contain
|
||||
`index` and `parent` columns.
|
||||
Clump array. Read from `CSiBORGReader.read_phew_clups`. Must
|
||||
contain `index` and `parent` columns.
|
||||
verbose : bool, optional
|
||||
Verbosity flag.
|
||||
|
||||
|
@ -522,10 +546,10 @@ class MmainReader:
|
|||
The ultimate parent halo index for every clump, i.e. referring to
|
||||
its ultimate parent clump.
|
||||
"""
|
||||
nsnap = max(self.paths.get_snapshots(nsim))
|
||||
partreader = ParticleReader(self.paths)
|
||||
nsnap = max(self.paths.get_snapshots(nsim, "csiborg"))
|
||||
partreader = CSiBORGReader(self.paths)
|
||||
cols = ["index", "parent", "mass_cl", 'x', 'y', 'z']
|
||||
clumparr = partreader.read_clumps(nsnap, nsim, cols)
|
||||
clumparr = partreader.read_phew_clups(nsnap, nsim, cols)
|
||||
|
||||
ultimate_parent = self.find_parents(clumparr, verbose=verbose)
|
||||
mask_main = clumparr["index"] == clumparr["parent"]
|
||||
|
@ -548,37 +572,162 @@ class MmainReader:
|
|||
out["subfrac"] = 1 - clumparr["mass_cl"][mask_main] / out["M"]
|
||||
return out, ultimate_parent
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Supplementary reading functions #
|
||||
# Quijote particle reader #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def read_initcm(nsim, srcdir, fname="clump_{}_cm.npy"):
|
||||
class QuijoteReader:
|
||||
"""
|
||||
Read `clump_cm`, i.e. the center of mass of a clump at redshift z = 70.
|
||||
If the file does not exist returns `None`.
|
||||
Object to read in Quijote snapshots from the binary files.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
srcdir : str
|
||||
Path to the folder containing the files.
|
||||
fname : str, optional
|
||||
File name convention. By default `clump_cm_{}.npy`, where the
|
||||
substituted value is `nsim`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : structured array
|
||||
paths : py:class`csiborgtools.read.Paths`
|
||||
"""
|
||||
fpath = join(srcdir, fname.format(nsim))
|
||||
try:
|
||||
return numpy.load(fpath)
|
||||
except FileNotFoundError:
|
||||
warn("File {} does not exist.".format(fpath), UserWarning,
|
||||
stacklevel=1)
|
||||
return None
|
||||
|
||||
def __init__(self, paths):
|
||||
self.paths = paths
|
||||
|
||||
def read_info(self, nsnap, nsim):
|
||||
snapshot = self.paths.snapshot(nsnap, nsim, "quijote")
|
||||
header = readgadget.header(snapshot)
|
||||
out = {"BoxSize": header.boxsize / 1e3, # Mpc/h
|
||||
"Nall": header.nall[1], # Tot num of particles
|
||||
"PartMass": header.massarr[1] * 1e10, # Part mass in Msun/h
|
||||
"Omega_m": header.omega_m,
|
||||
"Omega_l": header.omega_l,
|
||||
"h": header.hubble,
|
||||
"redshift": header.redshift,
|
||||
}
|
||||
out["TotMass"] = out["Nall"] * out["PartMass"]
|
||||
out["Hubble"] = (100.0 * numpy.sqrt(
|
||||
header.omega_m * (1.0 + header.redshift)**3 + header.omega_l))
|
||||
return out
|
||||
|
||||
def read_particle(self, nsnap, nsim, pars_extract=None,
|
||||
return_structured=True, verbose=True):
|
||||
assert pars_extract in [None, "pids"]
|
||||
snapshot = self.paths.snapshot(nsnap, nsim, "quijote")
|
||||
info = self.read_info(nsnap, nsim)
|
||||
ptype = [1] # DM in Gadget speech
|
||||
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: reading particle IDs.")
|
||||
pids = readgadget.read_block(snapshot, "ID ", ptype)
|
||||
|
||||
if pars_extract == "pids":
|
||||
return None, pids
|
||||
|
||||
if return_structured:
|
||||
dtype = {"names": ['x', 'y', 'z', 'vx', 'vy', 'vz', 'M'],
|
||||
"formats": [numpy.float32] * 7}
|
||||
out = numpy.full(info["Nall"], numpy.nan, dtype=dtype)
|
||||
else:
|
||||
out = numpy.full((info["Nall"], 7), numpy.nan, dtype=numpy.float32)
|
||||
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: reading particle positions.")
|
||||
pos = readgadget.read_block(snapshot, "POS ", ptype) / 1e3 # Mpc/h
|
||||
pos /= info["BoxSize"] # Box units
|
||||
|
||||
for i, p in enumerate(['x', 'y', 'z']):
|
||||
if return_structured:
|
||||
out[p] = pos[:, i]
|
||||
else:
|
||||
out[:, i] = pos[:, i]
|
||||
del pos
|
||||
collect()
|
||||
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: reading particle velocities.")
|
||||
# NOTE convert to box units.
|
||||
vel = readgadget.read_block(snapshot, "VEL ", ptype) # km/s
|
||||
vel *= (1 + info["redshift"])
|
||||
|
||||
for i, v in enumerate(['vx', 'vy', 'vz']):
|
||||
if return_structured:
|
||||
out[v] = vel[:, i]
|
||||
else:
|
||||
out[:, i + 3] = vel[:, i]
|
||||
del vel
|
||||
collect()
|
||||
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: reading particle masses.")
|
||||
if return_structured:
|
||||
out["M"] = info["PartMass"] / info["TotMass"]
|
||||
else:
|
||||
out[:, 6] = info["PartMass"] / info["TotMass"]
|
||||
|
||||
return out, pids
|
||||
|
||||
def read_fof_hids(self, nsnap, nsim, verbose=True, **kwargs):
|
||||
"""
|
||||
Read the FoF group membership of particles. Unassigned particles have
|
||||
FoF group ID 0.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nsnap : int
|
||||
Snapshot index.
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
verbose : bool, optional
|
||||
Verbosity flag.
|
||||
**kwargs : dict
|
||||
Keyword arguments for backward compatibility.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : 1-dimensional array of shape `(nparticles, )`
|
||||
Group membership of particles.
|
||||
"""
|
||||
redshift = {4: 0.0, 3: 0.5, 2: 1.0, 1: 2.0, 0: 3.0}.get(nsnap, None)
|
||||
if redshift is None:
|
||||
raise ValueError(f"Redshift of snapshot {nsnap} is not known.")
|
||||
path = self.paths.fof_cat(nsim, "quijote")
|
||||
cat = readfof.FoF_catalog(path, nsnap)
|
||||
|
||||
# Read the particle IDs of the snapshot
|
||||
__, pids = self.read_particle(nsnap, nsim, pars_extract="pids",
|
||||
verbose=verbose)
|
||||
|
||||
# Read the FoF particle membership. These are only assigned particles.
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: reading the FoF particle membership.",
|
||||
flush=True)
|
||||
group_pids = cat.GroupIDs
|
||||
group_len = cat.GroupLen
|
||||
|
||||
# Create a mapping from particle ID to FoF group ID.
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: creating the particle to FoF ID to map.",
|
||||
flush=True)
|
||||
ks = numpy.insert(numpy.cumsum(group_len), 0, 0)
|
||||
pid2hid = numpy.full((group_pids.size, 2), numpy.nan,
|
||||
dtype=numpy.uint32)
|
||||
for i, (k0, kf) in enumerate(zip(ks[:-1], ks[1:])):
|
||||
pid2hid[k0:kf, 0] = i + 1
|
||||
pid2hid[k0:kf, 1] = group_pids[k0:kf]
|
||||
pid2hid = {pid: hid for hid, pid in pid2hid}
|
||||
|
||||
# Create the final array of hids matchign the snapshot array.
|
||||
# Unassigned particles have hid 0.
|
||||
if verbose:
|
||||
print(f"{datetime.now()}: creating the final hid array.",
|
||||
flush=True)
|
||||
hids = numpy.full(pids.size, 0, dtype=numpy.uint32)
|
||||
for i in trange(pids.size) if verbose else range(pids.size):
|
||||
hids[i] = pid2hid.get(pids[i], 0)
|
||||
|
||||
return hids
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Supplementary reading functions #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def halfwidth_mask(pos, hw):
|
||||
|
@ -627,3 +776,27 @@ def load_halo_particles(hid, particles, halo_map, hid2map):
|
|||
return particles[k0:kf + 1, :]
|
||||
except KeyError:
|
||||
return None
|
||||
|
||||
|
||||
def convert_str_to_num(s):
|
||||
"""
|
||||
Convert a string representation of a number to its appropriate numeric type
|
||||
(int or float).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
s : str
|
||||
The string representation of the number.
|
||||
|
||||
Returns
|
||||
-------
|
||||
num : int or float
|
||||
"""
|
||||
try:
|
||||
return int(s)
|
||||
except ValueError:
|
||||
try:
|
||||
return float(s)
|
||||
except ValueError:
|
||||
warn(f"Cannot convert string '{s}' to number", UserWarning)
|
||||
return s
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue