Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-12-21 14:18:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Clean density calculation (#97)

Browse source 
* Get rid of utils

* Clean up imports

* Move some utils here

* Rename file

* Add simname to boxsize

* Add imports

* Delete old files

* Update README

* Update imports

* Add a new draft of the density calculator

* Update fields

* Draft of new density field calculatiosn

* Add snapshot

* Add boxsizes

* Little updates

* Bring back utils

* Edit docstrings

* Edits imports

* Add progress on snapshots

* edit improts

* add basic snapshot catalogue

* Add support for CSiBORG2 snapshot reader

* add paths to fofcat for csiborg2

* Add more imports

* Add more boxsize

* Add more imports

* Add field readers

* Simplify field paths

* Fix typo

* Add observer vp

* Clean up density field calculation

* Add a short note

* Edit args

* Remove old comments

* Edit docs

* Remove blank line

* Stop flipping RAMSES

* Remove comment

* Edit desc

* Remove normalization

* Remove old dist array

* Remove non-volume weighting

* Remove non-volume weight

* Add ignore of flake8 notebooks

* Fix path typo

* Fix units

* Edit paths docs

* Update nb
This commit is contained in:
Richard Stiskalek 2023-12-18 18:09:08 +01:00 committed by GitHub
parent eeff8f0ab9
commit eb1797e8a9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 1260 additions and 1139 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.flake8 Normal file
View file

@ -0,0 +1,3 @@
[flake8]
exclude = *.ipynb

2
README.md
View file

@ -10,6 +10,8 @@ however with little effort it can support other simulations as well.
## TODO
- [x] Prune old CSiBORG1 merger tree things.
- [x] Add visualiastion of the density field.
- [ ] Clear out `density` support.
- [ ] Add sorting of Gadget4 initial snapshot like final snapshot.
- [ ] Add full support for CSiBORG2 suite of simulations.
- [ ] Add SPH field calculation from cosmotools.

28
csiborgtools/__init__.py
View file

@ -37,6 +37,34 @@ neighbour_kwargs = {"rmax_radial": 155 / 0.705,
"paths_kind": paths_glamdring}
def simname2boxsize(simname):
"""
Return boxsize in `Mpc/h` for a given simname.
Parameters
----------
simname : str
Simulation name.
Returns
-------
boxsize : float
"""
d = {"csiborg1": 677.7,
"csiborg2_main": 676.6,
"csiborg2_varysmall": 676.6,
"csiborg2_random": 676.6,
"quijote": 1000.
}
boxsize = d.get(simname, None)
if boxsize is None:
raise ValueError("Unknown simname: {}".format(simname))
return boxsize
###############################################################################
# Surveys #
###############################################################################

11
csiborgtools/field/__init__.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.
from .density import (DensityField, PotentialField, TidalTensorField, # noqa
VelocityField, power_spectrum) # noqa
from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa
fill_outside, make_sky, observer_peculiar_velocity) # noqa
from .utils import nside2radec, smoothen_field # noqa
from .density import (DensityField, PotentialField, TidalTensorField, # noqa
VelocityField, radial_velocity, power_spectrum, # noqa
overdensity_field) # noqa
from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa
fill_outside, make_sky, observer_peculiar_velocity, # noqa
nside2radec, smoothen_field) # noqa

351
csiborgtools/field/density.py
View file

@ -13,37 +13,36 @@
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Density field and cross-correlation calculations.
Density field, potential and tidal tensor field calculations. Most routines
here do not support SPH-calculated density fields because of the unknown
corrrections necessary when performing the fast Fourier transform.
"""
from abc import ABC
import MAS_library as MASL
import Pk_library as PKL
import numpy
import Pk_library as PKL
from numba import jit
from tqdm import trange
from .interp import divide_nonzero
from .utils import force_single_precision
from .utils import divide_nonzero, force_single_precision
class BaseField(ABC):
"""Base class for density field calculations."""
_box = None
_MAS = None
_boxsize = None
@property
def box(self):
"""Simulation box information and transformations."""
return self._box
def boxsize(self):
"""Size of the box in units matching the particle coordinates."""
return self._boxsize
@box.setter
def box(self, box):
try:
assert box._name == "box_units"
self._box = box
except AttributeError as err:
raise TypeError from err
@boxsize.setter
def boxsize(self, value):
if not isinstance(value, (int, float)):
raise ValueError("`boxsize` must be an integer.")
self._boxsize = value
@property
def MAS(self):
@ -54,7 +53,12 @@ class BaseField(ABC):
@MAS.setter
def MAS(self, MAS):
assert MAS in ["NGP", "CIC", "TSC", "PCS"]
if MAS == "SPH":
raise ValueError("`SPH` is not supported. Use `cosmotool` scripts to calculate the density field with SPH.") # noqa
if MAS not in ["NGP", "CIC", "TSC", "PCS"]:
raise ValueError(f"Invalid `MAS` value: {MAS}")
self._MAS = MAS
@ -69,48 +73,27 @@ class DensityField(BaseField):
Parameters
----------
box : :py:class:`csiborgtools.read.CSiBORG1Box`
The simulation box information and transformations.
boxsize : float
Size of the periodic box.
MAS : str
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
Mass assignment scheme. Options are: 'NGP' (nearest grid point),
'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline).
paths : :py:class:`csiborgtools.read.Paths`
The simulation paths.
References
----------
[1] https://pylians3.readthedocs.io/
"""
def __init__(self, box, MAS):
self.box = box
def __init__(self, boxsize, MAS):
self.boxsize = boxsize
self.MAS = MAS
def overdensity_field(self, delta):
r"""
Calculate the overdensity field from the density field.
Defined as :math:`\rho/ <\rho> - 1`. Overwrites the input array.
Parameters
----------
delta : 3-dimensional array of shape `(grid, grid, grid)`
The density field.
Returns
-------
3-dimensional array of shape `(grid, grid, grid)`.
"""
delta /= delta.mean()
delta -= 1
return delta
def __call__(self, pos, mass, grid, nbatch=30, verbose=True):
"""
Calculate the density field using a Pylians routine [1, 2].
Iteratively loads the particles into memory, flips their `x` and `z`
coordinates. Particles are assumed to be in box units, with positions
in [0, 1]
Calculate the density field using a Pylians routine [1, 2]. Iteratively
loads the particles into memory. Particle coordinates units should
match that of `boxsize`.
Parameters
----------
@ -142,7 +125,7 @@ class DensityField(BaseField):
start = 0
for __ in trange(nbatch + 1, disable=not verbose,
desc="Loading particles for the density field"):
desc="Processing particles for the density field"):
end = min(start + batch_size, nparts)
batch_pos = pos[start:end]
batch_mass = mass[start:end]
@ -150,113 +133,43 @@ class DensityField(BaseField):
batch_pos = force_single_precision(batch_pos)
batch_mass = force_single_precision(batch_mass)
MASL.MA(batch_pos, rho, 1., self.MAS, W=batch_mass, verbose=False)
MASL.MA(batch_pos, rho, self.boxsize, self.MAS, W=batch_mass,
verbose=False)
if end == nparts:
break
start = end
# Divide by the cell volume in (kpc / h)^3
rho /= (self.box.boxsize / grid * 1e3)**3
rho /= (self.boxsize / grid * 1e3)**3
return rho
# class SPHDensityVelocity(BaseField):
# r"""
# Density field calculation. Based primarily on routines of Pylians [1].
#
# Parameters
# ----------
# box : :py:class:`csiborgtools.read.CSiBORG1Box`
# The simulation box information and transformations.
# MAS : str
# Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
# point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
# (piecewise cubic spline).
# paths : :py:class:`csiborgtools.read.Paths`
# The simulation paths.
#
# References
# ----------
# [1] https://pylians3.readthedocs.io/
# """
#
# def __init__(self, box, MAS):
# self.box = box
# self.MAS = MAS
#
# def overdensity_field(self, delta):
# r"""
# Calculate the overdensity field from the density field.
# Defined as :math:`\rho/ <\rho> - 1`. Overwrites the input array.
#
# Parameters
# ----------
# delta : 3-dimensional array of shape `(grid, grid, grid)`
# The density field.
#
# Returns
# -------
# 3-dimensional array of shape `(grid, grid, grid)`.
# """
# delta /= delta.mean()
# delta -= 1
# return delta
#
# def __call__(self, pos, mass, grid, nbatch=30, verbose=True):
# """
# Calculate the density field using a Pylians routine [1, 2].
# Iteratively loads the particles into memory, flips their `x` and `z`
# coordinates. Particles are assumed to be in box units, with positions
# in [0, 1]
#
# Parameters
# ----------
# pos : 2-dimensional array of shape `(n_parts, 3)`
# Particle positions
# mass : 1-dimensional array of shape `(n_parts,)`
# Particle masses
# grid : int
# Grid size.
# nbatch : int, optional
# Number of batches to split the particle loading into.
# verbose : bool, optional
# Verbosity flag.
#
# Returns
# -------
# 3-dimensional array of shape `(grid, grid, grid)`.
#
# References
# ----------
# [1] https://pylians3.readthedocs.io/
# [2] https://github.com/franciscovillaescusa/Pylians3/blob/master
# /library/MAS_library/MAS_library.pyx
# """
# rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
#
# nparts = pos.shape[0]
# batch_size = nparts // nbatch
# start = 0
#
# for __ in trange(nbatch + 1, disable=not verbose,
# desc="Loading particles for the density field"):
# end = min(start + batch_size, nparts)
# batch_pos = pos[start:end]
# batch_mass = mass[start:end]
#
# batch_pos = force_single_precision(batch_pos)
# batch_mass = force_single_precision(batch_mass)
#
# MASL.MA(batch_pos, rho, 1., self.MAS, W=batch_mass, verbose=False)
# if end == nparts:
# break
# start = end
#
# # Divide by the cell volume in (kpc / h)^3
# rho /= (self.box.boxsize / grid * 1e3)**3
#
# return rho
def overdensity_field(delta, make_copy=True):
r"""
Get the overdensity field from the density field as `rho / <rho> - 1`.
Parameters
----------
delta : 3-dimensional array of shape `(grid, grid, grid)`
The density field.
make_copy : bool, optional
Whether to make a copy of the input array.
Returns
-------
3-dimensional array of shape `(grid, grid, grid)`.
"""
if make_copy:
delta = numpy.copy(delta)
delta /= delta.mean()
delta -= 1
return delta
###############################################################################
# Velocity field calculation #
@ -269,11 +182,11 @@ class VelocityField(BaseField):
Parameters
----------
box : :py:class:`csiborgtools.read.CSiBORG1Box`
The simulation box information and transformations.
boxsize : float
Size of the periodic box.
MAS : str
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
Mass assignment scheme. Options are: 'NGP' (nearest grid point),
'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline).
References
@ -281,49 +194,11 @@ class VelocityField(BaseField):
[1] https://pylians3.readthedocs.io/
"""
def __init__(self, box, MAS):
self.box = box
def __init__(self, boxsize, MAS):
self.boxsize = boxsize
self.MAS = MAS
@staticmethod
@jit(nopython=True)
def radial_velocity(rho_vel, observer_velocity):
"""
Calculate the radial velocity field around the observer in the centre
of the box.
Parameters
----------
rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
Velocity field along each axis.
observer_velocity : 3-dimensional array of shape `(3,)`
Observer velocity.
Returns
-------
3-dimensional array of shape `(grid, grid, grid)`.
"""
grid = rho_vel.shape[1]
radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
vx0, vy0, vz0 = observer_velocity
for i in range(grid):
px = i - 0.5 * (grid - 1)
for j in range(grid):
py = j - 0.5 * (grid - 1)
for k in range(grid):
pz = k - 0.5 * (grid - 1)
vx = rho_vel[0, i, j, k] - vx0
vy = rho_vel[1, i, j, k] - vy0
vz = rho_vel[2, i, j, k] - vz0
radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
/ numpy.sqrt(px**2 + py**2 + pz**2))
return radvel
def __call__(self, pos, vel, mass, grid, flip_xz=True, nbatch=30,
verbose=True):
def __call__(self, pos, vel, mass, grid, nbatch=30, verbose=True):
"""
Calculate the velocity field using a Pylians routine [1, 2].
Iteratively loads the particles into memory, flips their `x` and `z`
@ -339,8 +214,6 @@ class VelocityField(BaseField):
Particle masses.
grid : int
Grid size.
flip_xz : bool, optional
Whether to flip the `x` and `z` coordinates.
nbatch : int, optional
Number of batches to split the particle loading into.
verbose : bool, optional
@ -379,11 +252,12 @@ class VelocityField(BaseField):
vel *= mass.reshape(-1, 1)
for i in range(3):
MASL.MA(pos, rho_vel[i], 1., self.MAS, W=vel[:, i],
MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[:, i],
verbose=False)
MASL.MA(pos, cellcounts, 1., self.MAS, W=mass,
MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass,
verbose=False)
if end == nparts:
break
start = end
@ -394,6 +268,43 @@ class VelocityField(BaseField):
return numpy.stack(rho_vel)
@jit(nopython=True)
def radial_velocity(rho_vel, observer_velocity):
"""
Calculate the radial velocity field around the observer in the centre
of the box.
Parameters
----------
rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
Velocity field along each axis.
observer_velocity : 3-dimensional array of shape `(3,)`
Observer velocity.
Returns
-------
3-dimensional array of shape `(grid, grid, grid)`.
"""
grid = rho_vel.shape[1]
radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
vx0, vy0, vz0 = observer_velocity
for i in range(grid):
px = i - 0.5 * (grid - 1)
for j in range(grid):
py = j - 0.5 * (grid - 1)
for k in range(grid):
pz = k - 0.5 * (grid - 1)
vx = rho_vel[0, i, j, k] - vx0
vy = rho_vel[1, i, j, k] - vy0
vz = rho_vel[2, i, j, k] - vz0
radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
/ numpy.sqrt(px**2 + py**2 + pz**2))
return radvel
###############################################################################
# Potential field calculation #
###############################################################################
@ -405,18 +316,18 @@ class PotentialField(BaseField):
Parameters
----------
box : :py:class:`csiborgtools.read.CSiBORG1Box`
The simulation box information and transformations.
boxsize : float
Size of the periodic box.
MAS : str
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
Mass assignment scheme. Options are: 'NGP' (nearest grid point),
'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline).
"""
def __init__(self, box, MAS):
self.box = box
def __init__(self, boxsize, MAS):
self.boxsize = boxsize
self.MAS = MAS
def __call__(self, overdensity_field):
def __call__(self, overdensity_field, omega_m, aexp):
"""
Calculate the potential field.
@ -424,13 +335,16 @@ class PotentialField(BaseField):
----------
overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
The overdensity field.
omega_m : float
TODO
aexp : float
TODO
Returns
-------
3-dimensional array of shape `(grid, grid, grid)`.
"""
return MASL.potential(overdensity_field, self.box._omega_m,
self.box._aexp, self.MAS)
return MASL.potential(overdensity_field, omega_m, aexp, self.MAS)
###############################################################################
@ -444,15 +358,15 @@ class TidalTensorField(BaseField):
Parameters
----------
box : :py:class:`csiborgtools.read.CSiBORG1Box`
The simulation box information and transformations.
boxsize : float
Size of the periodic box.
MAS : str
Mass assignment scheme used to calculate the density field. Options
are: 'NGP' (nearest grid point), 'CIC' (cloud-in-cell), 'TSC'
(triangular-shape cloud), 'PCS' (piecewise cubic spline).
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline).
"""
def __init__(self, box, MAS):
self.box = box
def __init__(self, boxsize, MAS):
self.boxsize = boxsize
self.MAS = MAS
@staticmethod
@ -494,7 +408,7 @@ class TidalTensorField(BaseField):
"""
return eigenvalues_to_environment(eigvals, threshold)
def __call__(self, overdensity_field):
def __call__(self, overdensity_field, omega_m, aexp):
"""
Calculate the tidal tensor field.
@ -502,6 +416,10 @@ class TidalTensorField(BaseField):
----------
overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
The overdensity field.
omega_m : float
TODO
aexp : float
TODO
Returns
-------
@ -509,8 +427,7 @@ class TidalTensorField(BaseField):
Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
the relevant tensor components.
"""
return MASL.tidal_tensor(overdensity_field, self.box._omega_m,
self.box._aexp, self.MAS)
return MASL.tidal_tensor(overdensity_field, omega_m, aexp, self.MAS)
@jit(nopython=True)
@ -606,7 +523,9 @@ def power_spectrum(delta, boxsize, MAS, threads=1, verbose=True):
boxsize : float
The simulation box size in `Mpc / h`.
MAS : str
Mass assignment scheme used to calculate the density field.
Mass assignment scheme used to calculate the density field. Options
are: 'NGP' (nearest grid point), 'CIC' (cloud-in-cell), 'TSC'
(triangular-shape cloud), 'PCS' (piecewise cubic spline).
threads : int, optional
Number of threads to use.
verbose : bool, optional
@ -617,6 +536,8 @@ def power_spectrum(delta, boxsize, MAS, threads=1, verbose=True):
k, Pk : 1-dimensional arrays of shape `(grid,)`
The wavenumbers and the power spectrum.
"""
axis = 2 # Axis along which compute the quadrupole and hexadecapole
# Axis along which compute the quadrupole and hexadecapole, is not used
# for the monopole that we calculat here.
axis = 2
Pk = PKL.Pk(delta, boxsize, axis, MAS, threads, verbose)
return Pk.k3D, Pk.Pk[:, 0]

67
csiborgtools/field/interp.py
View file

@ -15,13 +15,15 @@
"""
Tools for interpolating 3D fields at arbitrary positions.
"""
import healpy
import MAS_library as MASL
import numpy
import smoothing_library as SL
from numba import jit
from tqdm import trange, tqdm
from tqdm import tqdm, trange
from .utils import force_single_precision, smoothen_field
from ..utils import periodic_wrap_grid, radec_to_cartesian
from .utils import divide_nonzero, force_single_precision
###############################################################################
@ -169,7 +171,7 @@ def evaluate_sky(*fields, pos, mpc2box, smooth_scales=None, verbose=False):
smooth_scales=smooth_scales, verbose=verbose)
def make_sky(field, angpos, dist, boxsize, volume_weight=True, verbose=True):
def make_sky(field, angpos, dist, boxsize, verbose=True):
r"""
Make a sky map of a scalar field. The observer is in the centre of the
box the field is evaluated along directions `angpos` (RA [0, 360) deg,
@ -186,8 +188,6 @@ def make_sky(field, angpos, dist, boxsize, volume_weight=True, verbose=True):
Uniformly spaced radial distances to evaluate the field in `Mpc / h`.
boxsize : float
Box size in `Mpc / h`.
volume_weight : bool, optional
Whether to weight the field by the volume of the pixel.
verbose : bool, optional
Verbosity flag.
@ -209,17 +209,30 @@ def make_sky(field, angpos, dist, boxsize, volume_weight=True, verbose=True):
dir_loop[:, 0] = dist
dir_loop[:, 1] = angpos[i, 0]
dir_loop[:, 2] = angpos[i, 1]
if volume_weight:
out[i] = numpy.sum(
dist**2
* evaluate_sky(field, pos=dir_loop, mpc2box=1 / boxsize))
else:
out[i] = numpy.sum(
evaluate_sky(field, pos=dir_loop, mpc2box=1 / boxsize))
out *= dx
out[i] = numpy.sum(
dist**2 * evaluate_sky(field, pos=dir_loop, mpc2box=1 / boxsize))
# Assuming the field is in h^2 Msun / kpc**3, we need to convert Mpc / h
# to kpc / h and multiply by the pixel area.
out *= dx * 1e9 * 4 * numpy.pi / len(angpos)
return out
def nside2radec(nside):
"""
Generate RA [0, 360] deg. and declination [-90, 90] deg. for HEALPix pixel
centres at a given nside.
"""
pixs = numpy.arange(healpy.nside2npix(nside))
theta, phi = healpy.pix2ang(nside, pixs)
ra = 180 / numpy.pi * phi
dec = 90 - 180 / numpy.pi * theta
return numpy.vstack([ra, dec]).T
###############################################################################
# Real-to-redshift space field dragging #
###############################################################################
@ -302,21 +315,6 @@ def field2rsp(field, radvel_field, box, MAS, init_value=0.):
###############################################################################
@jit(nopython=True)
def divide_nonzero(field0, field1):
"""
Perform in-place `field0 /= field1` but only where `field1 != 0`.
"""
assert field0.shape == field1.shape, "Field shapes must match."
imax, jmax, kmax = field0.shape
for i in range(imax):
for j in range(jmax):
for k in range(kmax):
if field1[i, j, k] != 0:
field0[i, j, k] /= field1[i, j, k]
@jit(nopython=True)
def fill_outside(field, fill_value, rmax, boxsize):
"""
@ -339,3 +337,16 @@ def fill_outside(field, fill_value, rmax, boxsize):
if idist2 + jdist2 + kdist2 > rmax_box2:
field[i, j, k] = fill_value
return field
def smoothen_field(field, smooth_scale, boxsize, threads=1, make_copy=False):
"""
Smooth a field with a Gaussian filter.
"""
W_k = SL.FT_filter(boxsize, smooth_scale, field.shape[0], "Gaussian",
threads)
if make_copy:
field = numpy.copy(field)
return SL.field_smoothing(field, W_k, threads)

38
csiborgtools/field/utils.py
View file

@ -13,11 +13,11 @@
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Utility functions for the field module.
Utility functions used in the rest of the `field` module to avoid circular
imports.
"""
import healpy
from numba import jit
import numpy
import smoothing_library as SL
def force_single_precision(x):
@ -29,28 +29,16 @@ def force_single_precision(x):
return x
def smoothen_field(field, smooth_scale, boxsize, threads=1, make_copy=False):
@jit(nopython=True)
def divide_nonzero(field0, field1):
"""
Smooth a field with a Gaussian filter.
Perform in-place `field0 /= field1` but only where `field1 != 0`.
"""
W_k = SL.FT_filter(boxsize, smooth_scale, field.shape[0], "Gaussian",
threads)
assert field0.shape == field1.shape, "Field shapes must match."
if make_copy:
field = numpy.copy(field)
return SL.field_smoothing(field, W_k, threads)
def nside2radec(nside):
"""
Generate RA [0, 360] deg. and declination [-90, 90] deg. for HEALPix pixel
centres at a given nside.
"""
pixs = numpy.arange(healpy.nside2npix(nside))
theta, phi = healpy.pix2ang(nside, pixs)
ra = 180 / numpy.pi * phi
dec = 90 - 180 / numpy.pi * theta
return numpy.vstack([ra, dec]).T
imax, jmax, kmax = field0.shape
for i in range(imax):
for j in range(jmax):
for k in range(kmax):
if field1[i, j, k] != 0:
field0[i, j, k] /= field1[i, j, k]

6
csiborgtools/read/__init__.py
View file

@ -12,8 +12,10 @@
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from .halo_cat import (CSiBORGCatalogue, QuijoteCatalogue, # noqa
fiducial_observers) # noqa
from .catalogue import (CSiBORGCatalogue, QuijoteCatalogue, # noqa
fiducial_observers) # noqa
from .snapshot import (CSIBORG1Snapshot, CSIBORG2Snapshot, QuijoteSnapshot, # noqa
CSiBORG1Field, CSiBORG2Field, QuijoteField) # noqa
from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies, # noqa
TwoMPPGroups, ObservedCluster, match_array_to_no_masking) # noqa
from .paths import Paths # noqa

9
csiborgtools/read/halo_cat.py → csiborgtools/read/catalogue.py
View file

@ -28,8 +28,6 @@ from sklearn.neighbors import NearestNeighbors
from ..utils import (cartesian_to_radec, fprint, great_circle_distance,
number_counts, periodic_distance_two_points,
real2redshift)
# TODO: removing these
# from .box_units import CSiBORG1Box, QuijoteBox
from .paths import Paths
###############################################################################
@ -819,4 +817,9 @@ def load_halo_particles(hid, particles, hid2map):
k0, kf = hid2map[hid]
return particles[k0:kf + 1]
except KeyError:
return None
return None
###############################################################################
# Specific loaders of particles and haloes #
###############################################################################

3
csiborgtools/read/obs.py
View file

@ -109,8 +109,7 @@ class TwoMPPGalaxies(TextSurvey):
cat = cat[cat[:, 12] == 0, :]
# Pre=allocate array and fillt it
cols = [("RA", numpy.float64), ("DEC", numpy.float64),
("Ksmag", numpy.float64), ("ZCMB", numpy.float64),
("DIST", numpy.float64)]
("Ksmag", numpy.float64), ("ZCMB", numpy.float64)]
data = cols_to_structured(cat.shape[0], cols)
data["RA"] = cat[:, 1]
data["DEC"] = cat[:, 2]

155
csiborgtools/read/paths.py
View file

@ -41,23 +41,26 @@ class Paths:
Parameters
----------
# HERE EDIT EVERYTHING
srcdir : str, optional
Path to the folder where the RAMSES outputs are stored.
postdir: str, optional
Path to the folder where post-processed files are stored.
borg_dir : str, optional
Path to the folder where BORG MCMC chains are stored.
quiote_dir : str, optional
Path to the folder where Quijote simulations are stored.
csiborg1_srcdir : str
Path to the CSiBORG1 simulation directory.
csiborg2_main_srcdir : str
Path to the CSiBORG2 main simulation directory.
csiborg2_random_srcdir : str
Path to the CSiBORG2 random simulation directory.
csiborg2_varysmall_srcdir : str
Path to the CSiBORG2 varysmall simulation directory.
postdir : str
Path to the CSiBORG post-processing directory.
quijote_dir : str
Path to the Quijote simulation directory.
"""
def __init__(self,
csiborg1_srcdir=None,
csiborg2_main_srcdir=None,
csiborg2_random_srcdir=None,
csiborg2_varysmall_srcdir=None,
postdir=None,
quijote_dir=None
csiborg1_srcdir,
csiborg2_main_srcdir,
csiborg2_random_srcdir,
csiborg2_varysmall_srcdir,
postdir,
quijote_dir,
):
self.csiborg1_srcdir = csiborg1_srcdir
self.csiborg2_main_srcdir = csiborg2_main_srcdir
@ -117,8 +120,6 @@ class Paths:
-------
snapshots : 1-dimensional array
"""
# simpath = self.snapshots(nsim, simname, tonew=False)
if simname == "csiborg1":
snaps = glob(join(self.csiborg1_srcdir, f"chain_{nsim}",
"snapshot_*"))
@ -176,14 +177,14 @@ class Paths:
return join(self.csiborg1_srcdir, f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(5)}.hdf5")
elif simname == "csiborg2_main":
return join(self.csiborg2_main_srcdir, f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(3)}.hdf5")
return join(self.csiborg2_main_srcdir, f"chain_{nsim}", "output",
f"snapshot_{str(nsnap).zfill(3)}_full.hdf5")
elif simname == "csiborg2_random":
return join(self.csiborg2_random_srcdir, f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(3)}.hdf5")
return join(self.csiborg2_random_srcdir, f"chain_{nsim}", "output",
f"snapshot_{str(nsnap).zfill(3)}_full.hdf5")
elif simname == "csiborg2_varysmall":
return join(self.csiborg2_varysmall_srcdir, f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(3)}.hdf5")
"output", f"snapshot_{str(nsnap).zfill(3)}_full.hdf5")
elif simname == "quijote":
return join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}",
@ -191,6 +192,43 @@ class Paths:
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
def snapshot_catalogue(self, nsnap, nsim, simname):
"""
Path to the halo catalogue of a simulation snapshot.
Parameters
----------
nsnap : int
Snapshot index.
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
"""
if simname == "csiborg1":
return join(self.csiborg1_srcdir, f"chain_{nsim}",
f"fof_{str(nsnap).zfill(5)}.hdf5")
elif simname == "csiborg2_main":
return join(self.csiborg2_main_srcdir, f"chain_{nsim}", "output",
f"fof_subhalo_tab_{str(nsnap).zfill(3)}.hdf5")
elif simname == "csiborg2_random":
return join(self.csiborg2_ranodm_srcdir, f"chain_{nsim}", "output",
f"fof_subhalo_tab_{str(nsnap).zfill(3)}.hdf5")
elif simname == "csiborg2_varysmall":
return join(self.csiborg2_varysmall_srcdir, f"chain_{nsim}",
"output",
f"fof_subhalo_tab_{str(nsnap).zfill(3)}.hdf5")
elif simname == "quijote":
return join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}",
f"fof_{str(nsnap).zfill(3)}.hdf5")
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
def overlap(self, simname, nsim0, nsimx, min_logmass, smoothed):
"""
Path to the overlap files between two CSiBORG simulations.
@ -274,48 +312,77 @@ class Paths:
return join(fdir, fname)
def field(self, kind, MAS, grid, nsim, in_rsp, smooth_scale=None):
def field(self, kind, MAS, grid, nsim, simname):
r"""
Path to the files containing the calculated fields in CSiBORG.
Parameters
----------
kind : str
Field type. Must be one of: `density`, `velocity`, `potential`,
`radvel`, `environment`.
Field type.
MAS : str
Mass-assignment scheme.
grid : int
Grid size.
nsim : int
IC realisation index.
in_rsp : bool
Whether the calculation is performed in redshift space.
smooth_scale : float, optional
Smoothing scale in Mpc/h.
simname : str
Simulation name.
Returns
-------
str
"""
assert kind in ["density", "velocity", "potential", "radvel",
"environment"]
fdir = join(self.postdir, "environment")
if MAS == "SPH":
if kind not in ["density", "velocity"]:
raise ValueError("SPH field must be either `density` or `velocity`.") # noqa
if simname == "csiborg1":
raise ValueError("SPH field not available for CSiBORG1.")
elif simname == "csiborg2_main":
return join(self.csiborg2_main_srcdir, "field",
f"chain_{nsim}_{grid}.hdf5")
elif simname == "csiborg2_random":
return join(self.csiborg2_random_srcdir, "field",
f"chain_{nsim}_{grid}.hdf5")
elif simname == "csiborg2_varysmall":
return join(self.csiborg2_varysmall_srcdir, "field",
f"chain_{nsim}_{grid}.hdf5")
elif simname == "quijote":
raise ValueError("SPH field not available for CSiBORG1.")
fdir = join(self.postdir, "environment")
try_create_directory(fdir)
if in_rsp:
kind = kind + "_rsp"
fname = f"{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npy"
if smooth_scale is not None:
fname = fname.replace(".npy", f"_smooth{smooth_scale}.npy")
fname = f"{kind}_{simname}_{MAS}_{str(nsim).zfill(5)}_{grid}.npy"
return join(fdir, fname)
def field_interpolated(self, survey, kind, MAS, grid, nsim, in_rsp,
smooth_scale=None):
def observer_peculiar_velocity(self, MAS, grid, nsim, simname):
"""
Path to the files containing the observer peculiar velocity.
Parameters
----------
MAS : str
Mass-assignment scheme.
grid : int
Grid size.
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
"""
fdir = join(self.postdir, "environment")
try_create_directory(fdir)
fname = f"observer_peculiar_velocity_{simname}_{MAS}_{str(nsim).zfill(5)}_{grid}.npz" # noqa
return join(fdir, fname)
def field_interpolated(self, survey, kind, MAS, grid, nsim, in_rsp):
"""
Path to the files containing the CSiBORG interpolated field for a given
survey.
@ -335,13 +402,12 @@ class Paths:
IC realisation index.
in_rsp : bool
Whether the calculation is performed in redshift space.
smooth_scale : float, optional
Smoothing scale in Mpc/h.
Returns
-------
str
"""
raise NotImplementedError("This function is not implemented yet.")
assert kind in ["density", "velocity", "potential", "radvel",
"environment"]
fdir = join(self.postdir, "environment_interpolated")
@ -353,9 +419,6 @@ class Paths:
fname = f"{survey}_{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npz"
if smooth_scale is not None:
fname = fname.replace(".npz", f"_smooth{smooth_scale}.npz")
return join(fdir, fname)
def cross_nearest(self, simname, run, kind, nsim=None, nobs=None):

657
csiborgtools/read/snapshot.py Normal file
View file

@ -0,0 +1,657 @@
# Copyright (C) 2023 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Classes for reading in snapshots and unifying the snapshot interface. Here
should be implemented things such as flipping x- and z-axes, to make sure that
observed RA-dec can be mapped into the simulation box.
"""
from abc import ABC, abstractmethod, abstractproperty
import numpy
from h5py import File
###############################################################################
# Base snapshot class #
###############################################################################
class BaseSnapshot(ABC):
"""
Base class for reading snapshots.
"""
def __init__(self, nsim, nsnap, paths):
if not isinstance(nsim, int):
raise TypeError("`nsim` must be an integer")
self._nsim = nsim
if not isinstance(nsnap, int):
raise TypeError("`nsnap` must be an integer")
self._nsnap = nsnap
self._paths = paths
self._hid2offset = None
@property
def nsim(self):
"""
Simulation index.
Returns
-------
int
"""
return self._nsim
@property
def nsnap(self):
"""
Snapshot index.
Returns
-------
int
"""
return self._nsnap
@property
def paths(self):
"""
Paths manager.
Returns
-------
Paths
"""
return self._paths
@abstractproperty
def coordinates(self):
"""
Return the particle coordinates.
Returns
-------
coords : 2-dimensional array
"""
pass
@abstractproperty
def velocities(self):
"""
Return the particle velocities.
Returns
-------
vel : 2-dimensional array
"""
pass
@abstractproperty
def masses(self):
"""
Return the particle masses.
Returns
-------
mass : 1-dimensional array
"""
pass
@abstractproperty
def particle_ids(self):
"""
Return the particle IDs.
Returns
-------
ids : 1-dimensional array
"""
pass
@abstractmethod
def halo_coordinates(self, halo_id, is_group):
"""
Return the halo particle coordinates.
Parameters
----------
halo_id : int
Halo ID.
is_group : bool
If `True`, return the group coordinates. Otherwise, return the
subhalo coordinates.
Returns
-------
coords : 2-dimensional array
"""
pass
@abstractmethod
def halo_velocities(self, halo_id, is_group):
"""
Return the halo particle velocities.
Parameters
----------
halo_id : int
Halo ID.
is_group : bool
If `True`, return the group velocities. Otherwise, return the
subhalo velocities.
Returns
-------
vel : 2-dimensional array
"""
pass
@abstractmethod
def halo_masses(self, halo_id, is_group):
"""
Return the halo particle masses.
Parameters
----------
halo_id : int
Halo ID.
is_group : bool
If `True`, return the group masses. Otherwise, return the
subhalo masses.
Returns
-------
mass : 1-dimensional array
"""
pass
@property
def hid2offset(self):
if self._hid2offset is None:
self._make_hid2offset()
return self._hid2offset
@abstractmethod
def _make_hid2offset(self):
"""
Private class function to make the halo ID to offset dictionary.
"""
pass
###############################################################################
# CSiBORG1 snapshot class #
###############################################################################
class CSIBORG1Snapshot(BaseSnapshot):
"""
CSiBORG1 snapshot class with the FoF halo finder particle assignment.
CSiBORG1 was run with RAMSES.
Parameters
----------
nsim : int
Simulation index.
nsnap : int
Snapshot index.
paths : Paths
Paths object.
"""
def __init__(self, nsim, nsnap, paths):
super().__init__(nsim, nsnap, paths)
self._snapshot_path = self.paths.snapshot(
self.nsnap, self.nsim, "csiborg1")
def _get_particles(self, kind):
with File(self._snapshot_path, "r") as f:
x = f[kind][...]
return x
def coordinates(self):
return self._get_particles("Coordinates")
def velocities(self):
return self._get_particles("Velocities")
def masses(self):
return self._get_particles("Masses")
def particle_ids(self):
with File(self._snapshot_path, "r") as f:
ids = f["ParticleIDs"][...]
return ids
def _get_halo_particles(self, halo_id, kind, is_group):
if not is_group:
raise ValueError("There is no subhalo catalogue for CSiBORG1.")
with File(self._snapshot_path, "r") as f:
i, j = self.hid2offset.get(halo_id, (None, None))
if i is None:
raise ValueError(f"Halo `{halo_id}` not found.")
x = f[kind][i:j + 1]
return x
def halo_coordinates(self, halo_id, is_group=True):
return self._get_halo_particles(halo_id, "Coordinates", is_group)
def halo_velocities(self, halo_id, is_group=True):
return self._get_halo_particles(halo_id, "Velocities", is_group)
def halo_masses(self, halo_id, is_group=True):
return self._get_halo_particles(halo_id, "Masses", is_group)
def _make_hid2offset(self):
catalogue_path = self.paths.snapshot_catalogue(
self.nsnap, self.nsim, "csiborg1")
with File(catalogue_path, "r") as f:
offset = f["GroupOffset"][:]
self._hid2offset = {i: (j, k) for i, j, k in offset}
###############################################################################
# CSiBORG2 snapshot class #
###############################################################################
class CSIBORG2Snapshot(BaseSnapshot):
"""
CSiBORG2 snapshot class with the FoF halo finder particle assignment and
SUBFIND subhalo finder. The simulations were run with Gadget4.
Parameters
----------
nsim : int
Simulation index.
nsnap : int
Snapshot index.
paths : Paths
Paths object.
kind : str
CSiBORG2 run kind. One of `main`, `random`, or `varysmall`.
"""
def __init__(self, nsim, nsnap, paths, kind):
super().__init__(nsim, nsnap, paths)
self.kind = kind
self._snapshot_path = self.paths.snapshot(
self.nsnap, self.nsim, f"csiborg2_{self.kind}")
@property
def kind(self):
"""
CSiBORG2 run kind.
Returns
-------
str
"""
return self._kind
@kind.setter
def kind(self, value):
if value not in ["main", "random", "varysmall"]:
raise ValueError("`kind` must be one of `main`, `random`, or `varysmall`.") # noqa
self._kind = value
def _get_particles(self, kind):
with File(self._snapshot_path, "r") as f:
if kind == "Masses":
npart = f["Header"].attrs["NumPart_Total"][1]
x = numpy.ones(npart, dtype=numpy.float32)
x *= f["Header"].attrs["MassTable"][1]
else:
x = f[f"PartType1/{kind}"][...]
if x.ndim == 1:
x = numpy.hstack([x, f[f"PartType5/{kind}"][...]])
else:
x = numpy.vstack([x, f[f"PartType5/{kind}"][...]])
return x
def coordinates(self):
return self._get_particles("Coordinates")
def velocities(self):
return self._get_particles("Velocities")
def masses(self):
return self._get_particles("Masses") * 1e10
def particle_ids(self):
return self._get_particles("ParticleIDs")
def _get_halo_particles(self, halo_id, kind, is_group):
if not is_group:
raise RuntimeError("While the CSiBORG2 subhalo catalogue exists, it is not currently implemented.") # noqa
with File(self._snapshot_path, "r") as f:
i1, j1 = self.hid2offset["type1"].get(halo_id, (None, None))
i5, j5 = self.hid2offset["type5"].get(halo_id, (None, None))
# Check if this is a valid halo
if i1 is None and i5 is None:
raise ValueError(f"Halo `{halo_id}` not found.")
if j1 - i1 == 0 and j5 - i5 == 0:
raise ValueError(f"Halo `{halo_id}` has no particles.")
if i1 is not None and j1 - i1 > 0:
if kind == "Masses":
x1 = numpy.ones(j1 - i1, dtype=numpy.float32)
x1 *= f["Header"].attrs["MassTable"][1]
else:
x1 = f[f"PartType1/{kind}"][i1:j1]
if i5 is not None and j5 - i5 > 0:
x5 = f[f"PartType5/{kind}"][i5:j5]
if i5 is None or j5 - i5 == 0:
return x1
if i1 is None or j1 - i1 == 0:
return x5
if x1.ndim > 1:
x1 = numpy.vstack([x1, x5])
else:
x1 = numpy.hstack([x1, x5])
return x1
def halo_coordinates(self, halo_id, is_group=True):
return self._get_halo_particles(halo_id, "Coordinates", is_group)
def halo_velocities(self, halo_id, is_group=True):
return self._get_halo_particles(halo_id, "Velocities", is_group)
def halo_masses(self, halo_id, is_group=True):
return self._get_halo_particles(halo_id, "Masses", is_group) * 1e10
def _make_hid2offset(self):
catalogue_path = self.paths.snapshot_catalogue(
self.nsnap, self.nsim, f"csiborg2_{self.kind}")
with File(catalogue_path, "r") as f:
offset = f["Group/GroupOffsetType"][:, 1]
lenghts = f["Group/GroupLenType"][:, 1]
hid2offset_type1 = {i: (offset[i], offset[i] + lenghts[i])
for i in range(len(offset))}
offset = f["Group/GroupOffsetType"][:, 5]
lenghts = f["Group/GroupLenType"][:, 5]
hid2offset_type5 = {i: (offset[i], offset[i] + lenghts[i])
for i in range(len(offset))}
self._hid2offset = {"type1": hid2offset_type1,
"type5": hid2offset_type5,
}
###############################################################################
# CSiBORG2 snapshot class #
###############################################################################
class QuijoteSnapshot(CSIBORG1Snapshot):
"""
Quijote snapshot class with the FoF halo finder particle assignment.
Because of similarities with how the snapshot is processed with CSiBORG1,
it uses the same base class.
Parameters
----------
nsim : int
Simulation index.
nsnap : int
Snapshot index.
paths : Paths
Paths object.
"""
def __init__(self, nsim, nsnap, paths):
super().__init__(nsim, nsnap, paths)
self._snapshot_path = self.paths.snapshot(self.nsnap, self.nsim,
"quijote")
def _make_hid2offset(self):
catalogue_path = self.paths.snapshot_catalogue(
self.nsnap, self.nsim, "quijote")
with File(catalogue_path, "r") as f:
offset = f["GroupOffset"][:]
self._hid2offset = {int(i): (int(j), int(k)) for i, j, k in offset}
###############################################################################
# Base field class #
###############################################################################
class BaseField(ABC):
"""
Base class for reading fields such as density or velocity fields.
"""
def __init__(self, nsim, paths):
if not isinstance(nsim, int):
raise TypeError("`nsim` must be an integer")
self._nsim = nsim
self._paths = paths
@property
def nsim(self):
"""
Simulation index.
Returns
-------
int
"""
return self._nsim
@property
def paths(self):
"""
Paths manager.
Returns
-------
Paths
"""
return self._paths
@abstractmethod
def density_field(self, MAS, grid):
"""
Return the pre-computed density field.
Parameters
----------
MAS : str
Mass assignment scheme.
grid : int
Grid size.
Returns
-------
field : 3-dimensional array
"""
pass
@abstractmethod
def velocity_field(self, MAS, grid):
"""
Return the pre-computed velocity field.
Parameters
----------
MAS : str
Mass assignment scheme.
grid : int
Grid size.
Returns
-------
field : 4-dimensional array
"""
pass
###############################################################################
# CSiBORG1 field class #
###############################################################################
class CSiBORG1Field(BaseField):
"""
CSiBORG1 `z = 0` field class.
Parameters
----------
nsim : int
Simulation index.
paths : Paths
Paths object.
"""
def __init__(self, nsim, paths):
super().__init__(nsim, paths)
def density_field(self, MAS, grid):
fpath = self.paths.field("density", MAS, grid, self.nsim, "csiborg1")
if MAS == "SPH":
with File(fpath, "r") as f:
field = f["density"][:]
else:
field = numpy.load(fpath)
return field
def velocity_field(self, MAS, grid):
fpath = self.paths.field("velocity", MAS, grid, self.nsim, "csiborg1")
if MAS == "SPH":
with File(fpath, "r") as f:
density = f["density"][:]
v0 = f["p0"][:] / density
v1 = f["p1"][:] / density
v2 = f["p2"][:] / density
field = numpy.array([v0, v1, v2])
else:
field = numpy.load(fpath)
return field
###############################################################################
# CSiBORG2 field class #
###############################################################################
class CSiBORG2Field(BaseField):
"""
CSiBORG2 `z = 0` field class.
Parameters
----------
nsim : int
Simulation index.
paths : Paths
Paths object.
kind : str
CSiBORG2 run kind. One of `main`, `random`, or `varysmall`.
"""
def __init__(self, nsim, paths, kind):
super().__init__(nsim, paths)
self.kind = kind
@property
def kind(self):
"""
CSiBORG2 run kind.
Returns
-------
str
"""
return self._kind
@kind.setter
def kind(self, value):
if value not in ["main", "random", "varysmall"]:
raise ValueError("`kind` must be one of `main`, `random`, or `varysmall`.") # noqa
self._kind = value
def density_field(self, MAS, grid):
fpath = self.paths.field("density", MAS, grid, self.nsim,
f"csiborg2_{self.kind}")
if MAS == "SPH":
with File(fpath, "r") as f:
field = f["density"][:]
field *= 1e10 # Convert to Msun / h
field /= (676.6 * 1e3 / 1024)**3 # Convert to h^2 Msun / kpc^3
field = field.T # Flip x- and z-axes
else:
field = numpy.load(fpath)
return field
def velocity_field(self, MAS, grid):
fpath = self.paths.field("velocity", MAS, grid, self.nsim,
f"csiborg2_{self.kind}")
if MAS == "SPH":
with File(fpath, "r") as f:
# TODO: the x and z still have to be flipped.
density = f["density"][:]
v0 = f["p0"][:] / density
v1 = f["p1"][:] / density
v2 = f["p2"][:] / density
field = numpy.array([v0, v1, v2])
else:
field = numpy.load(fpath)
return field
###############################################################################
# Quijote field class #
###############################################################################
class QuijoteField(CSiBORG1Field):
"""
Quijote `z = 0` field class.
Parameters
----------
nsim : int
Simulation index.
paths : Paths
Paths object.
"""
def __init__(self, nsim, paths):
super().__init__(nsim, paths)

108
scripts/field_obs_vp.py
View file

@ -1,108 +0,0 @@
# Copyright (C) 2022 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Script to calculate the peculiar velocity of an observer in the centre of the
CSiBORG box.
"""
from argparse import ArgumentParser
from distutils.util import strtobool
import numpy
from mpi4py import MPI
from taskmaster import work_delegation
from tqdm import tqdm
from utils import get_nsims
try:
import csiborgtools
except ModuleNotFoundError:
import sys
sys.path.append("../")
import csiborgtools
def observer_peculiar_velocity(nsim, parser_args):
"""
Calculate the peculiar velocity of an observer in the centre of the box
for several smoothing scales.
"""
pos = numpy.array([0.5, 0.5, 0.5]).reshape(-1, 3)
boxsize = 677.7
smooth_scales = [0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]
observer_vp = numpy.full((len(smooth_scales), 3), numpy.nan,
dtype=numpy.float32)
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
field_path = paths.field("velocity", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False)
field0 = numpy.load(field_path)
for j, smooth_scale in enumerate(tqdm(smooth_scales,
desc="Smoothing the fields",
disable=not parser_args.verbose)):
if smooth_scale > 0:
field = [None, None, None]
for k in range(3):
field[k] = csiborgtools.field.smoothen_field(
field0[k], smooth_scale, boxsize)
else:
field = field0
v = csiborgtools.field.evaluate_cartesian(
field[0], field[1], field[2], pos=pos)
observer_vp[j, 0] = v[0][0]
observer_vp[j, 1] = v[1][0]
observer_vp[j, 2] = v[2][0]
fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
nsim)
if parser_args.verbose:
print(f"Saving to ... `{fout}`")
numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
return observer_vp
###############################################################################
# Command line interface #
###############################################################################
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--nsims", type=int, nargs="+", default=None,
help="IC realisations. `-1` for all simulations.")
parser.add_argument("--kind", type=str,
choices=["density", "rspdensity", "velocity", "radvel",
"potential", "environment"],
help="What derived field to calculate?")
parser.add_argument("--MAS", type=str,
choices=["NGP", "CIC", "TSC", "PCS"])
parser.add_argument("--grid", type=int, help="Grid resolution.")
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
help="Verbosity flag for reading in particles.")
parser.add_argument("--simname", type=str, default="csiborg",
help="Verbosity flag for reading in particles.")
parser_args = parser.parse_args()
comm = MPI.COMM_WORLD
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = get_nsims(parser_args, paths)
def main(nsim):
return observer_peculiar_velocity(nsim, parser_args)
work_delegation(main, nsims, comm, master_verbose=True)

382
scripts/field_prop.py
View file

@ -12,14 +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.
"""
MPI script to calculate density field-derived fields in the CSiBORG
simulations' final snapshot.
"""
"""MPI script to calculate the various fields."""
from argparse import ArgumentParser
from datetime import datetime
from distutils.util import strtobool
from gc import collect
import numpy
from mpi4py import MPI
@ -29,55 +24,43 @@ import csiborgtools
from utils import get_nsims
###############################################################################
# Cosmotool SPH density & velocity field #
###############################################################################
def cosmotool_sph(nsim, parser_args):
pass
###############################################################################
# Density field #
###############################################################################
def density_field(nsim, parser_args, to_save=True):
"""
Calculate the density field in the CSiBORG simulation.
"""
def density_field(nsim, parser_args):
"""Calculate the density field."""
if parser_args.MAS == "SPH":
raise NotImplementedError("SPH is not implemented here. Use cosmotool")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
if not parser_args.in_rsp:
# TODO I removed this function
snap = csiborgtools.read.read_h5(fname)["snapshot_final"]
pos = snap["pos"]
mass = snap["mass"]
gen = csiborgtools.field.DensityField(box, parser_args.MAS)
field = gen(pos, mass, parser_args.grid, verbose=parser_args.verbose)
# Read in the particle coordinates and masses
if parser_args.simname == "csiborg1":
snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
elif "csiborg2" in parser_args.simname:
kind = parser_args.simname.split("_")[-1]
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
elif parser_args.simname == "quijote":
snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
else:
field = numpy.load(paths.field(
"density", parser_args.MAS, parser_args.grid, nsim, False))
radvel_field = numpy.load(paths.field(
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
if parser_args.verbose:
print(f"{datetime.now()}: converting density field to RSP.",
flush=True)
pos = snapshot.coordinates()
mass = snapshot.masses()
field = csiborgtools.field.field2rsp(field, radvel_field, box,
parser_args.MAS)
# Run the field generator
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
gen = csiborgtools.field.DensityField(boxsize, parser_args.MAS)
field = gen(pos, mass, parser_args.grid)
if to_save:
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
fout = paths.field("density", parser_args.MAS, parser_args.grid,
nsim, parser_args.simname)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
return field
@ -86,31 +69,36 @@ def density_field(nsim, parser_args, to_save=True):
###############################################################################
def velocity_field(nsim, parser_args, to_save=True):
"""
Calculate the velocity field in a CSiBORG simulation.
"""
if parser_args.in_rsp:
raise NotImplementedError("Velocity field in RSP is not implemented.")
def velocity_field(nsim, parser_args):
"""Calculate the velocity field."""
if parser_args.MAS == "SPH":
raise NotImplementedError("SPH is not implemented here. Use cosmotool")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
snap = csiborgtools.read.read_h5(fname)["snapshot_final"]
pos = snap["pos"]
vel = snap["vel"]
mass = snap["mass"]
if parser_args.simname == "csiborg1":
snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
elif "csiborg2" in parser_args.simname:
kind = parser_args.simname.split("_")[-1]
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
elif parser_args.simname == "quijote":
snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
else:
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
field = gen(pos, vel, mass, parser_args.grid, verbose=parser_args.verbose)
pos = snapshot.coordinates()
vel = snapshot.velocities()
mass = snapshot.masses()
if to_save:
fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
gen = csiborgtools.field.VelocityField(boxsize, parser_args.MAS)
field = gen(pos, vel, mass, parser_args.grid)
fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
nsim, parser_args.simname)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
return field
@ -119,125 +107,62 @@ def velocity_field(nsim, parser_args, to_save=True):
###############################################################################
def radvel_field(nsim, parser_args, to_save=True):
"""
Calculate the radial velocity field in the CSiBORG simulation.
"""
if parser_args.in_rsp:
raise NotImplementedError("Radial vel. field in RSP not implemented.")
def radvel_field(nsim, parser_args):
"""Calculate the radial velocity field."""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp))
observer_velocity = csiborgtools.field.observer_vobs(vel)
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
field = gen.radial_velocity(vel, observer_velocity)
if to_save:
fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
return field
###############################################################################
# Potential field #
###############################################################################
def potential_field(nsim, parser_args, to_save=True):
"""
Calculate the potential field in the CSiBORG simulation.
"""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
if not parser_args.in_rsp:
rho = numpy.load(paths.field(
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
rho = density_gen.overdensity_field(rho)
gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
field = gen(rho)
if parser_args.simname == "csiborg1":
field = csiborgtools.read.CSiBORG1Field(nsim, paths)
elif "csiborg2" in parser_args.simname:
kind = parser_args.simname.split("_")[-1]
field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
elif parser_args.simname == "quijote":
field = csiborgtools.read.QuijoteField(nsim, paths)
else:
field = numpy.load(paths.field(
"potential", parser_args.MAS, parser_args.grid, nsim, False))
radvel_field = numpy.load(paths.field(
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
field = csiborgtools.field.field2rsp(field, radvel_field, box,
parser_args.MAS)
vel = field.velocity_field(parser_args.MAS, parser_args.grid)
if to_save:
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
observer_velocity = csiborgtools.field.observer_peculiar_velocity(vel)
radvel = csiborgtools.field.radial_velocity(vel, observer_velocity)
fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
nsim, parser_args.simname)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, radvel)
return field
###############################################################################
# Environment classification #
###############################################################################
def environment_field(nsim, parser_args, to_save=True):
def observer_peculiar_velocity(nsim, parser_args):
"""
Calculate the environmental classification in the CSiBORG simulation.
Calculate the peculiar velocity of an observer in the centre of the box
for several smoothing scales.
"""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
# NOTE thevse values are hard-coded.
smooth_scales = numpy.array([0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
smooth_scales /= boxsize
rho = numpy.load(paths.field(
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
rho = density_gen.overdensity_field(rho)
if parser_args.simname == "csiborg1":
field = csiborgtools.read.CSiBORG1Field(nsim, paths)
elif "csiborg2" in parser_args.simname:
kind = parser_args.simname.split("_")[-1]
field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
elif parser_args.simname == "quijote":
field = csiborgtools.read.QuijoteField(nsim, paths)
else:
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
if parser_args.smooth_scale > 0.0:
rho = csiborgtools.field.smoothen_field(
rho, parser_args.smooth_scale, box.box2mpc(1.))
vel = field.velocity_field(parser_args.MAS, parser_args.grid)
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
field = gen(rho)
observer_vp = csiborgtools.field.observer_peculiar_velocity(
vel, smooth_scales)
del rho
collect()
if parser_args.in_rsp:
radvel_field = numpy.load(paths.field(
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
args = (radvel_field, box, parser_args.MAS)
field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
del radvel_field
collect()
eigvals = gen.tensor_field_eigvals(field)
del field
collect()
env = gen.eigvals_to_environment(eigvals)
if to_save:
fout = paths.field("environment", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp, parser_args.smooth_scale)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, env)
return env
fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
nsim, parser_args.simname)
print(f"Saving to ... `{fout}`")
numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
return observer_vp
###############################################################################
@ -249,39 +174,124 @@ if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--nsims", type=int, nargs="+", default=None,
help="IC realisations. `-1` for all simulations.")
parser.add_argument("--simname", type=str, help="Simulation name.")
parser.add_argument("--kind", type=str,
choices=["density", "rspdensity", "velocity", "radvel",
"potential", "environment"],
choices=["density", "velocity", "radvel", "observer_vp"], # noqa
help="What derived field to calculate?")
parser.add_argument("--MAS", type=str,
choices=["NGP", "CIC", "TSC", "PCS"])
choices=["NGP", "CIC", "TSC", "PCS", "SPH"],
help="Mass assignment scheme.")
parser.add_argument("--grid", type=int, help="Grid resolution.")
parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
help="Calculate in RSP?")
parser.add_argument("--smooth_scale", type=float, default=0.0,
help="Smoothing scale in Mpc / h. Only used for the environment field.") # noqa
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
help="Verbosity flag for reading in particles.")
parser.add_argument("--simname", type=str, default="csiborg",
choices=["csiborg", "csiborg2"],
help="Verbosity flag for reading in particles.")
parser_args = parser.parse_args()
comm = MPI.COMM_WORLD
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = get_nsims(parser_args, paths)
def main(nsim):
if parser_args.kind == "density" or parser_args.kind == "rspdensity":
if parser_args.kind == "density":
density_field(nsim, parser_args)
elif parser_args.kind == "velocity":
velocity_field(nsim, parser_args)
elif parser_args.kind == "radvel":
radvel_field(nsim, parser_args)
elif parser_args.kind == "potential":
potential_field(nsim, parser_args)
elif parser_args.kind == "environment":
environment_field(nsim, parser_args)
elif parser_args.kind == "observer_vp":
observer_peculiar_velocity(nsim, parser_args)
else:
raise RuntimeError(f"Field {parser_args.kind} is not implemented.")
work_delegation(main, nsims, comm, master_verbose=True)
# def potential_field(nsim, parser_args, to_save=True):
# """
# Calculate the potential field in the CSiBORG simulation.
# """
# paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
# nsnap = max(paths.get_snapshots(nsim, "csiborg"))
# box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
#
# if not parser_args.in_rsp:
# rho = numpy.load(paths.field(
# "density", parser_args.MAS, parser_args.grid, nsim,
# in_rsp=False))
# density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
# rho = density_gen.overdensity_field(rho)
#
# gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
# field = gen(rho)
# else:
# field = numpy.load(paths.field(
# "potential", parser_args.MAS, parser_args.grid, nsim, False))
# radvel_field = numpy.load(paths.field(
# "radvel", parser_args.MAS, parser_args.grid, nsim, False))
#
# field = csiborgtools.field.field2rsp(field, radvel_field, box,
# parser_args.MAS)
#
# if to_save:
# fout = paths.field(parser_args.kind, parser_args.MAS,
# parser_args.grid,
# nsim, parser_args.in_rsp)
# print(f"{datetime.now()}: saving output to `{fout}`.")
# numpy.save(fout, field)
# return field
#
#
# #############################################################################
# # Environment classification #
# #############################################################################
#
#
# def environment_field(nsim, parser_args, to_save=True):
# """
# Calculate the environmental classification in the CSiBORG simulation.
# """
# paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
# nsnap = max(paths.get_snapshots(nsim, "csiborg"))
# box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
#
# rho = numpy.load(paths.field(
# "density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
# density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
# rho = density_gen.overdensity_field(rho)
#
# if parser_args.smooth_scale > 0.0:
# rho = csiborgtools.field.smoothen_field(
# rho, parser_args.smooth_scale, box.box2mpc(1.))
#
# gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
# field = gen(rho)
#
# del rho
# collect()
#
# if parser_args.in_rsp:
# radvel_field = numpy.load(paths.field(
# "radvel", parser_args.MAS, parser_args.grid, nsim, False))
# args = (radvel_field, box, parser_args.MAS)
#
# field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
# field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
# field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
# field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
# field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
# field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
#
# del radvel_field
# collect()
#
# eigvals = gen.tensor_field_eigvals(field)
#
# del field
# collect()
#
# env = gen.eigvals_to_environment(eigvals)
#
# if to_save:
# fout = paths.field("environment", parser_args.MAS, parser_args.grid,
# nsim, parser_args.in_rsp,
# parser_args.smooth_scale)
# print(f"{datetime.now()}: saving output to `{fout}`.")
# numpy.save(fout, env)
# return env

7
scripts/field_sample.py
View file

@ -30,7 +30,12 @@ import csiborgtools
from utils import get_nsims
# TODO get rid of this.
MPC2BOX = 1 / 677.7
# MPC2BOX = 1 / 677.7
SIM2BOXSIZE = {"csiborg1": 677.7,
"csiborg2_main": None,
"csiborg2_random": None,
"csiborg2_varysmall": None,
}
def steps(cls, survey_name):

374
scripts/process_snapshot.py
View file

@ -1,374 +0,0 @@
# Copyright (C) 2022 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
r"""
Script to process simulation files and create a single HDF5 file, in which
particles are sorted by the particle halo IDs.
"""
from argparse import ArgumentParser
from gc import collect
import h5py
import numpy
from mpi4py import MPI
import csiborgtools
from csiborgtools import fprint
from numba import jit
from taskmaster import work_delegation
from tqdm import trange, tqdm
from utils import get_nsims
@jit(nopython=True, boundscheck=False)
def minmax_halo(hid, halo_ids, start_loop=0):
"""
Find the start and end index of a halo in a sorted array of halo IDs.
This is much faster than using `numpy.where` and then `numpy.min` and
`numpy.max`.
"""
start = None
end = None
for i in range(start_loop, halo_ids.size):
n = halo_ids[i]
if n == hid:
if start is None:
start = i
end = i
elif n > hid:
break
return start, end
def process_snapshot(nsim, simname, halo_finder, verbose):
"""
Read in the snapshot particles, sort them by their halo ID and dump
into a HDF5 file. Stores the first and last index of each halo in the
particle array for fast slicing of the array to acces particles of a single
halo.
"""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, simname))
if simname == "csiborg":
partreader = csiborgtools.read.CSiBORGReader(paths)
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
else:
partreader = csiborgtools.read.QuijoteReader(paths)
box = None
desc = {"hid": f"Halo finder ID ({halo_finder})of the particle.",
"pos": "DM particle positions in box units.",
"vel": "DM particle velocity in km / s.",
"mass": "DM particle mass in Msun / h.",
"pid": "DM particle ID",
}
fname = paths.processed_output(nsim, simname, halo_finder)
fprint(f"loading HIDs of IC {nsim}.", verbose)
hids = partreader.read_halo_id(nsnap, nsim, halo_finder, verbose)
collect()
fprint(f"sorting HIDs of IC {nsim}.")
sort_indxs = numpy.argsort(hids)
with h5py.File(fname, "w") as f:
group = f.create_group("snapshot_final")
group.attrs["header"] = "Snapshot data at z = 0."
fprint("dumping halo IDs.", verbose)
dset = group.create_dataset("halo_ids", data=hids[sort_indxs])
dset.attrs["header"] = desc["hid"]
del hids
collect()
fprint("reading, sorting and dumping the snapshot particles.", verbose)
for kind in ["pos", "vel", "mass", "pid"]:
x = partreader.read_snapshot(nsnap, nsim, kind)[sort_indxs]
if simname == "csiborg" and kind == "vel":
x = box.box2vel(x) if simname == "csiborg" else x
if simname == "csiborg" and kind == "mass":
x = box.box2solarmass(x) if simname == "csiborg" else x
dset = f["snapshot_final"].create_dataset(kind, data=x)
dset.attrs["header"] = desc[kind]
del x
collect()
del sort_indxs
collect()
fprint(f"creating a halo map for IC {nsim}.")
with h5py.File(fname, "r") as f:
part_hids = f["snapshot_final"]["halo_ids"][:]
# We loop over the unique halo IDs and remove the 0 halo ID
unique_halo_ids = numpy.unique(part_hids)
unique_halo_ids = unique_halo_ids[unique_halo_ids != 0]
halo_map = numpy.full((unique_halo_ids.size, 3), numpy.nan,
dtype=numpy.uint64)
start_loop, niters = 0, unique_halo_ids.size
for i in trange(niters, disable=not verbose):
hid = unique_halo_ids[i]
k0, kf = minmax_halo(hid, part_hids, start_loop=start_loop)
halo_map[i, :] = hid, k0, kf
start_loop = kf
# Dump the halo mapping.
with h5py.File(fname, "r+") as f:
dset = f["snapshot_final"].create_dataset("halo_map", data=halo_map)
dset.attrs["header"] = """
Halo to particle mapping. Columns are HID, start index, end index.
"""
f.close()
del part_hids
collect()
# Add the halo finder catalogue
with h5py.File(fname, "r+") as f:
group = f.create_group("halofinder_catalogue")
group.attrs["header"] = f"Original {halo_finder} halo catalogue."
cat = partreader.read_catalogue(nsnap, nsim, halo_finder)
hid2pos = {hid: i for i, hid in enumerate(unique_halo_ids)}
for key in cat.dtype.names:
x = numpy.full(unique_halo_ids.size, numpy.nan,
dtype=cat[key].dtype)
for i in range(len(cat)):
j = hid2pos[cat["index"][i]]
x[j] = cat[key][i]
group.create_dataset(key, data=x)
f.close()
# Lastly create the halo catalogue
with h5py.File(fname, "r+") as f:
group = f.create_group("halo_catalogue")
group.attrs["header"] = f"{halo_finder} halo catalogue."
group.create_dataset("index", data=unique_halo_ids)
f.close()
def add_initial_snapshot(nsim, simname, halo_finder, verbose):
"""
Sort the initial snapshot particles according to their final snapshot and
add them to the final snapshot's HDF5 file.
"""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
fname = paths.processed_output(nsim, simname, halo_finder)
if simname == "csiborg":
partreader = csiborgtools.read.CSiBORGReader(paths)
else:
partreader = csiborgtools.read.QuijoteReader(paths)
fprint(f"processing simulation `{nsim}`.", verbose)
if simname == "csiborg":
nsnap0 = 1
elif simname == "quijote":
nsnap0 = -1
else:
raise ValueError(f"Unknown simulation `{simname}`.")
fprint("loading and sorting the initial PID.", verbose)
sort_indxs = numpy.argsort(partreader.read_snapshot(nsnap0, nsim, "pid"))
fprint("loading the final particles.", verbose)
with h5py.File(fname, "r") as f:
sort_indxs_final = f["snapshot_final/pid"][:]
f.close()
fprint("sorting the particles according to the final snapshot.", verbose)
sort_indxs_final = numpy.argsort(numpy.argsort(sort_indxs_final))
sort_indxs = sort_indxs[sort_indxs_final]
del sort_indxs_final
collect()
fprint("loading and sorting the initial particle position.", verbose)
pos = partreader.read_snapshot(nsnap0, nsim, "pos")[sort_indxs]
del sort_indxs
collect()
# In Quijote some particles are position precisely at the edge of the
# box. Move them to be just inside.
if simname == "quijote":
mask = pos >= 1
if numpy.any(mask):
spacing = numpy.spacing(pos[mask])
assert numpy.max(spacing) <= 1e-5
pos[mask] -= spacing
fprint(f"dumping particles for `{nsim}` to `{fname}`.", verbose)
with h5py.File(fname, "r+") as f:
if "snapshot_initial" in f.keys():
del f["snapshot_initial"]
group = f.create_group("snapshot_initial")
group.attrs["header"] = "Initial snapshot data."
dset = group.create_dataset("pos", data=pos)
dset.attrs["header"] = "DM particle positions in box units."
f.close()
def calculate_initial(nsim, simname, halo_finder, verbose):
"""Calculate the Lagrangian patch centre of mass and size."""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
fname = paths.processed_output(nsim, simname, halo_finder)
fprint("loading the particle information.", verbose)
f = h5py.File(fname, "r")
pos = f["snapshot_initial/pos"]
mass = f["snapshot_final/mass"]
hid = f["halo_catalogue/index"][:]
hid2map = csiborgtools.read.make_halomap_dict(
f["snapshot_final/halo_map"][:])
if simname == "csiborg":
kwargs = {"box_size": 2048, "bckg_halfsize": 512}
else:
kwargs = {"box_size": 512, "bckg_halfsize": 256}
overlapper = csiborgtools.match.ParticleOverlap(**kwargs)
lagpatch_pos = numpy.full((len(hid), 3), numpy.nan, dtype=numpy.float32)
lagpatch_size = numpy.full(len(hid), numpy.nan, dtype=numpy.float32)
lagpatch_ncells = numpy.full(len(hid), numpy.nan, dtype=numpy.int32)
for i in trange(len(hid), disable=not verbose):
h = hid[i]
# These are unasigned particles.
if h == 0:
continue
parts_pos = csiborgtools.read.load_halo_particles(h, pos, hid2map)
parts_mass = csiborgtools.read.load_halo_particles(h, mass, hid2map)
# Skip if the halo has no particles or is too small.
if parts_pos is None or parts_pos.size < 5:
continue
cm = csiborgtools.center_of_mass(parts_pos, parts_mass, boxsize=1.0)
sep = csiborgtools.periodic_distance(parts_pos, cm, boxsize=1.0)
delta = overlapper.make_delta(parts_pos, parts_mass, subbox=True)
lagpatch_pos[i] = cm
lagpatch_size[i] = numpy.percentile(sep, 99)
lagpatch_ncells[i] = csiborgtools.delta2ncells(delta)
f.close()
collect()
with h5py.File(fname, "r+") as f:
grp = f["halo_catalogue"]
dset = grp.create_dataset("lagpatch_pos", data=lagpatch_pos)
dset.attrs["header"] = "Lagrangian patch centre of mass in box units."
dset = grp.create_dataset("lagpatch_size", data=lagpatch_size)
dset.attrs["header"] = "Lagrangian patch size in box units."
dset = grp.create_dataset("lagpatch_ncells", data=lagpatch_ncells)
dset.attrs["header"] = f"Lagrangian patch number of cells on a {kwargs['box_size']}^3 grid." # noqa
f.close()
def make_phew_halo_catalogue(nsim, verbose):
"""
Process the PHEW halo catalogue for a CSiBORG simulation at all snapshots.
"""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
snapshots = paths.get_snapshots(nsim, "csiborg")
reader = csiborgtools.read.CSiBORGReader(paths)
keys_write = ["index", "x", "y", "z", "mass_cl", "parent",
"ultimate_parent", "summed_mass"]
# Create a HDF5 file to store all this.
fname = paths.processed_phew(nsim)
with h5py.File(fname, "w") as f:
f.close()
for nsnap in tqdm(snapshots, disable=not verbose, desc="Snapshot"):
try:
data = reader.read_phew_clumps(nsnap, nsim, verbose=False)
except FileExistsError:
continue
with h5py.File(fname, "r+") as f:
if str(nsnap) in f:
print(f"Group {nsnap} already exists. Deleting.", flush=True)
del f[str(nsnap)]
grp = f.create_group(str(nsnap))
for key in keys_write:
grp.create_dataset(key, data=data[key])
grp.attrs["header"] = f"CSiBORG PHEW clumps at snapshot {nsnap}."
f.close()
# Now write the redshifts
scale_factors = numpy.full(len(snapshots), numpy.nan, dtype=numpy.float32)
for i, nsnap in enumerate(snapshots):
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
scale_factors[i] = box._aexp
redshifts = scale_factors[-1] / scale_factors - 1
with h5py.File(fname, "r+") as f:
grp = f.create_group("info")
grp.create_dataset("redshift", data=redshifts)
grp.create_dataset("snapshots", data=snapshots)
grp.create_dataset("Om0", data=[box.Om0])
grp.create_dataset("boxsize", data=[box.boxsize])
f.close()
def main(nsim, args):
if args.make_final:
process_snapshot(nsim, args.simname, args.halofinder, True)
if args.make_initial:
add_initial_snapshot(nsim, args.simname, args.halofinder, True)
calculate_initial(nsim, args.simname, args.halofinder, True)
if args.make_phew:
make_phew_halo_catalogue(nsim, True)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--simname", type=str, default="csiborg",
choices=["csiborg", "quijote"],
help="Simulation name")
parser.add_argument("--nsims", type=int, nargs="+", default=None,
help="IC realisations. If `-1` processes all.")
parser.add_argument("--halofinder", type=str, help="Halo finder")
parser.add_argument("--make_final", action="store_true", default=False,
help="Process the final snapshot.")
parser.add_argument("--make_initial", action="store_true", default=False,
help="Process the initial snapshot.")
parser.add_argument("--make_phew", action="store_true", default=False,
help="Process the PHEW halo catalogue.")
args = parser.parse_args()
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = get_nsims(args, paths)
def _main(nsim):
main(nsim, args)
work_delegation(_main, nsims, MPI.COMM_WORLD)

17
scripts_independent/process_snapshot.py
View file

@ -60,16 +60,6 @@ def now():
return datetime.now()
def flip_cols(arr, col1, col2):
"""
Flip values in columns `col1` and `col2` of a structured array `arr`.
"""
if col1 not in arr.dtype.names or col2 not in arr.dtype.names:
raise ValueError(f"Both `{col1}` and `{col2}` must exist in `arr`.")
arr[col1], arr[col2] = numpy.copy(arr[col2]), numpy.copy(arr[col1])
def convert_str_to_num(s):
"""
Convert a string representation of a number to its appropriate numeric type
@ -221,11 +211,6 @@ class CSiBORG1Reader:
raise ValueError(f"Unknown kind `{kind}`. "
"Options are: `pid`, `pos`, `vel` or `mass`.")
# Because of a RAMSES bug x and z are flipped.
if kind in ["pos", "vel"]:
print(f"For kind `{kind}` flipping x and z.")
x[:, [0, 2]] = x[:, [2, 0]]
del sim
collect()
@ -273,8 +258,6 @@ class CSiBORG1Reader:
out["y"] = pos[:, 1] * h + 677.7 / 2
out["z"] = pos[:, 2] * h + 677.7 / 2
# Because of a RAMSES bug x and z are flipped.
flip_cols(out, "x", "z")
out["totpartmass"] = totmass * 1e11 * h
out["m200c"] = m200c * 1e11 * h

5
scripts_independent/run_process_snapshot.py
View file

@ -15,9 +15,8 @@
from os import system
if __name__ == "__main__":
# Quijote chains
chains = [1]
simname = "quijote"
chains = [7444]
simname = "csiborg1"
mode = 2
env = "/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"

176
scripts_plots/paper_environment.ipynb
View file

File diff suppressed because one or more lines are too long