Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-12-22 07:08:01 +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 ## TODO
- [x] Prune old CSiBORG1 merger tree things. - [x] Prune old CSiBORG1 merger tree things.
- [x] Add visualiastion of the density field. - [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 full support for CSiBORG2 suite of simulations.
- [ ] Add SPH field calculation from cosmotools. - [ ] 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} "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 # # 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 # 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., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from .density import (DensityField, PotentialField, TidalTensorField, # noqa from .density import (DensityField, PotentialField, TidalTensorField, # noqa
VelocityField, power_spectrum) # noqa VelocityField, radial_velocity, power_spectrum, # noqa
from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa overdensity_field) # noqa
fill_outside, make_sky, observer_peculiar_velocity) # noqa from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa
from .utils import nside2radec, smoothen_field # 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., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 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 from abc import ABC
import MAS_library as MASL import MAS_library as MASL
import Pk_library as PKL
import numpy import numpy
import Pk_library as PKL
from numba import jit from numba import jit
from tqdm import trange from tqdm import trange
from .interp import divide_nonzero from .utils import divide_nonzero, force_single_precision
from .utils import force_single_precision
class BaseField(ABC): class BaseField(ABC):
"""Base class for density field calculations.""" """Base class for density field calculations."""
_box = None
_MAS = None _MAS = None
_boxsize = None
@property @property
def box(self): def boxsize(self):
"""Simulation box information and transformations.""" """Size of the box in units matching the particle coordinates."""
return self._box return self._boxsize
@box.setter @boxsize.setter
def box(self, box): def boxsize(self, value):
try: if not isinstance(value, (int, float)):
assert box._name == "box_units" raise ValueError("`boxsize` must be an integer.")
self._box = box self._boxsize = value
except AttributeError as err:
raise TypeError from err
@property @property
def MAS(self): def MAS(self):
@ -54,7 +53,12 @@ class BaseField(ABC):
@MAS.setter @MAS.setter
def MAS(self, MAS): 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 self._MAS = MAS
@ -69,48 +73,27 @@ class DensityField(BaseField):
Parameters Parameters
---------- ----------
box : :py:class:`csiborgtools.read.CSiBORG1Box` boxsize : float
The simulation box information and transformations. Size of the periodic box.
MAS : str MAS : str
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid Mass assignment scheme. Options are: 'NGP' (nearest grid point),
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS' 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline). (piecewise cubic spline).
paths : :py:class:`csiborgtools.read.Paths`
The simulation paths.
References References
---------- ----------
[1] https://pylians3.readthedocs.io/ [1] https://pylians3.readthedocs.io/
""" """
def __init__(self, box, MAS): def __init__(self, boxsize, MAS):
self.box = box self.boxsize = boxsize
self.MAS = MAS 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): def __call__(self, pos, mass, grid, nbatch=30, verbose=True):
""" """
Calculate the density field using a Pylians routine [1, 2]. Calculate the density field using a Pylians routine [1, 2]. Iteratively
Iteratively loads the particles into memory, flips their `x` and `z` loads the particles into memory. Particle coordinates units should
coordinates. Particles are assumed to be in box units, with positions match that of `boxsize`.
in [0, 1]
Parameters Parameters
---------- ----------
@ -142,7 +125,7 @@ class DensityField(BaseField):
start = 0 start = 0
for __ in trange(nbatch + 1, disable=not verbose, 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) end = min(start + batch_size, nparts)
batch_pos = pos[start:end] batch_pos = pos[start:end]
batch_mass = mass[start:end] batch_mass = mass[start:end]
@ -150,113 +133,43 @@ class DensityField(BaseField):
batch_pos = force_single_precision(batch_pos) batch_pos = force_single_precision(batch_pos)
batch_mass = force_single_precision(batch_mass) 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: if end == nparts:
break break
start = end start = end
# Divide by the cell volume in (kpc / h)^3 # Divide by the cell volume in (kpc / h)^3
rho /= (self.box.boxsize / grid * 1e3)**3 rho /= (self.boxsize / grid * 1e3)**3
return rho return rho
# class SPHDensityVelocity(BaseField): def overdensity_field(delta, make_copy=True):
# r""" r"""
# Density field calculation. Based primarily on routines of Pylians [1]. Get the overdensity field from the density field as `rho / <rho> - 1`.
#
# Parameters Parameters
# ---------- ----------
# box : :py:class:`csiborgtools.read.CSiBORG1Box` delta : 3-dimensional array of shape `(grid, grid, grid)`
# The simulation box information and transformations. The density field.
# MAS : str make_copy : bool, optional
# Mass assignment scheme. Options are Options are: 'NGP' (nearest grid Whether to make a copy of the input array.
# point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
# (piecewise cubic spline). Returns
# paths : :py:class:`csiborgtools.read.Paths` -------
# The simulation paths. 3-dimensional array of shape `(grid, grid, grid)`.
# """
# References if make_copy:
# ---------- delta = numpy.copy(delta)
# [1] https://pylians3.readthedocs.io/
# """ delta /= delta.mean()
# delta -= 1
# def __init__(self, box, MAS):
# self.box = box return delta
# 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
############################################################################### ###############################################################################
# Velocity field calculation # # Velocity field calculation #
@ -269,11 +182,11 @@ class VelocityField(BaseField):
Parameters Parameters
---------- ----------
box : :py:class:`csiborgtools.read.CSiBORG1Box` boxsize : float
The simulation box information and transformations. Size of the periodic box.
MAS : str MAS : str
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid Mass assignment scheme. Options are: 'NGP' (nearest grid point),
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS' 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline). (piecewise cubic spline).
References References
@ -281,49 +194,11 @@ class VelocityField(BaseField):
[1] https://pylians3.readthedocs.io/ [1] https://pylians3.readthedocs.io/
""" """
def __init__(self, box, MAS): def __init__(self, boxsize, MAS):
self.box = box self.boxsize = boxsize
self.MAS = MAS self.MAS = MAS
@staticmethod def __call__(self, pos, vel, mass, grid, nbatch=30, verbose=True):
@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):
""" """
Calculate the velocity field using a Pylians routine [1, 2]. Calculate the velocity field using a Pylians routine [1, 2].
Iteratively loads the particles into memory, flips their `x` and `z` Iteratively loads the particles into memory, flips their `x` and `z`
@ -339,8 +214,6 @@ class VelocityField(BaseField):
Particle masses. Particle masses.
grid : int grid : int
Grid size. Grid size.
flip_xz : bool, optional
Whether to flip the `x` and `z` coordinates.
nbatch : int, optional nbatch : int, optional
Number of batches to split the particle loading into. Number of batches to split the particle loading into.
verbose : bool, optional verbose : bool, optional
@ -379,11 +252,12 @@ class VelocityField(BaseField):
vel *= mass.reshape(-1, 1) vel *= mass.reshape(-1, 1)
for i in range(3): 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) verbose=False)
MASL.MA(pos, cellcounts, 1., self.MAS, W=mass, MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass,
verbose=False) verbose=False)
if end == nparts: if end == nparts:
break break
start = end start = end
@ -394,6 +268,43 @@ class VelocityField(BaseField):
return numpy.stack(rho_vel) 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 # # Potential field calculation #
############################################################################### ###############################################################################
@ -405,18 +316,18 @@ class PotentialField(BaseField):
Parameters Parameters
---------- ----------
box : :py:class:`csiborgtools.read.CSiBORG1Box` boxsize : float
The simulation box information and transformations. Size of the periodic box.
MAS : str MAS : str
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid Mass assignment scheme. Options are: 'NGP' (nearest grid point),
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS' 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline). (piecewise cubic spline).
""" """
def __init__(self, box, MAS): def __init__(self, boxsize, MAS):
self.box = box self.boxsize = boxsize
self.MAS = MAS self.MAS = MAS
def __call__(self, overdensity_field): def __call__(self, overdensity_field, omega_m, aexp):
""" """
Calculate the potential field. Calculate the potential field.
@ -424,13 +335,16 @@ class PotentialField(BaseField):
---------- ----------
overdensity_field : 3-dimensional array of shape `(grid, grid, grid)` overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
The overdensity field. The overdensity field.
omega_m : float
TODO
aexp : float
TODO
Returns Returns
------- -------
3-dimensional array of shape `(grid, grid, grid)`. 3-dimensional array of shape `(grid, grid, grid)`.
""" """
return MASL.potential(overdensity_field, self.box._omega_m, return MASL.potential(overdensity_field, omega_m, aexp, self.MAS)
self.box._aexp, self.MAS)
############################################################################### ###############################################################################
@ -444,15 +358,15 @@ class TidalTensorField(BaseField):
Parameters Parameters
---------- ----------
box : :py:class:`csiborgtools.read.CSiBORG1Box` boxsize : float
The simulation box information and transformations. Size of the periodic box.
MAS : str MAS : str
Mass assignment scheme used to calculate the density field. Options Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
are: 'NGP' (nearest grid point), 'CIC' (cloud-in-cell), 'TSC' point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(triangular-shape cloud), 'PCS' (piecewise cubic spline). (piecewise cubic spline).
""" """
def __init__(self, box, MAS): def __init__(self, boxsize, MAS):
self.box = box self.boxsize = boxsize
self.MAS = MAS self.MAS = MAS
@staticmethod @staticmethod
@ -494,7 +408,7 @@ class TidalTensorField(BaseField):
""" """
return eigenvalues_to_environment(eigvals, threshold) return eigenvalues_to_environment(eigvals, threshold)
def __call__(self, overdensity_field): def __call__(self, overdensity_field, omega_m, aexp):
""" """
Calculate the tidal tensor field. Calculate the tidal tensor field.
@ -502,6 +416,10 @@ class TidalTensorField(BaseField):
---------- ----------
overdensity_field : 3-dimensional array of shape `(grid, grid, grid)` overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
The overdensity field. The overdensity field.
omega_m : float
TODO
aexp : float
TODO
Returns Returns
------- -------
@ -509,8 +427,7 @@ class TidalTensorField(BaseField):
Tidal tensor object, whose attributes `tidal_tensor.Tij` contain Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
the relevant tensor components. the relevant tensor components.
""" """
return MASL.tidal_tensor(overdensity_field, self.box._omega_m, return MASL.tidal_tensor(overdensity_field, omega_m, aexp, self.MAS)
self.box._aexp, self.MAS)
@jit(nopython=True) @jit(nopython=True)
@ -606,7 +523,9 @@ def power_spectrum(delta, boxsize, MAS, threads=1, verbose=True):
boxsize : float boxsize : float
The simulation box size in `Mpc / h`. The simulation box size in `Mpc / h`.
MAS : str 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 threads : int, optional
Number of threads to use. Number of threads to use.
verbose : bool, optional 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,)` k, Pk : 1-dimensional arrays of shape `(grid,)`
The wavenumbers and the power spectrum. 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) Pk = PKL.Pk(delta, boxsize, axis, MAS, threads, verbose)
return Pk.k3D, Pk.Pk[:, 0] 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. Tools for interpolating 3D fields at arbitrary positions.
""" """
import healpy
import MAS_library as MASL import MAS_library as MASL
import numpy import numpy
import smoothing_library as SL
from numba import jit 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 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) 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""" r"""
Make a sky map of a scalar field. The observer is in the centre of the 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, 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`. Uniformly spaced radial distances to evaluate the field in `Mpc / h`.
boxsize : float boxsize : float
Box size in `Mpc / h`. Box size in `Mpc / h`.
volume_weight : bool, optional
Whether to weight the field by the volume of the pixel.
verbose : bool, optional verbose : bool, optional
Verbosity flag. Verbosity flag.
@ -209,17 +209,30 @@ def make_sky(field, angpos, dist, boxsize, volume_weight=True, verbose=True):
dir_loop[:, 0] = dist dir_loop[:, 0] = dist
dir_loop[:, 1] = angpos[i, 0] dir_loop[:, 1] = angpos[i, 0]
dir_loop[:, 2] = angpos[i, 1] dir_loop[:, 2] = angpos[i, 1]
if volume_weight:
out[i] = numpy.sum( out[i] = numpy.sum(
dist**2 dist**2 * evaluate_sky(field, pos=dir_loop, mpc2box=1 / boxsize))
* evaluate_sky(field, pos=dir_loop, mpc2box=1 / boxsize))
else: # Assuming the field is in h^2 Msun / kpc**3, we need to convert Mpc / h
out[i] = numpy.sum( # to kpc / h and multiply by the pixel area.
evaluate_sky(field, pos=dir_loop, mpc2box=1 / boxsize)) out *= dx * 1e9 * 4 * numpy.pi / len(angpos)
out *= dx
return out 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 # # 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) @jit(nopython=True)
def fill_outside(field, fill_value, rmax, boxsize): 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: if idist2 + jdist2 + kdist2 > rmax_box2:
field[i, j, k] = fill_value field[i, j, k] = fill_value
return field 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., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 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 numpy
import smoothing_library as SL
def force_single_precision(x): def force_single_precision(x):
@ -29,28 +29,16 @@ def force_single_precision(x):
return 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", assert field0.shape == field1.shape, "Field shapes must match."
threads)
if make_copy: imax, jmax, kmax = field0.shape
field = numpy.copy(field) for i in range(imax):
for j in range(jmax):
return SL.field_smoothing(field, W_k, threads) for k in range(kmax):
if field1[i, j, k] != 0:
field0[i, j, k] /= field1[i, j, k]
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

6
csiborgtools/read/__init__.py
View file

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

7
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, from ..utils import (cartesian_to_radec, fprint, great_circle_distance,
number_counts, periodic_distance_two_points, number_counts, periodic_distance_two_points,
real2redshift) real2redshift)
# TODO: removing these
# from .box_units import CSiBORG1Box, QuijoteBox
from .paths import Paths from .paths import Paths
############################################################################### ###############################################################################
@ -820,3 +818,8 @@ def load_halo_particles(hid, particles, hid2map):
return particles[k0:kf + 1] return particles[k0:kf + 1]
except KeyError: 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, :] cat = cat[cat[:, 12] == 0, :]
# Pre=allocate array and fillt it # Pre=allocate array and fillt it
cols = [("RA", numpy.float64), ("DEC", numpy.float64), cols = [("RA", numpy.float64), ("DEC", numpy.float64),
("Ksmag", numpy.float64), ("ZCMB", numpy.float64), ("Ksmag", numpy.float64), ("ZCMB", numpy.float64)]
("DIST", numpy.float64)]
data = cols_to_structured(cat.shape[0], cols) data = cols_to_structured(cat.shape[0], cols)
data["RA"] = cat[:, 1] data["RA"] = cat[:, 1]
data["DEC"] = cat[:, 2] data["DEC"] = cat[:, 2]

155
csiborgtools/read/paths.py
View file

@ -41,23 +41,26 @@ class Paths:
Parameters Parameters
---------- ----------
# HERE EDIT EVERYTHING csiborg1_srcdir : str
srcdir : str, optional Path to the CSiBORG1 simulation directory.
Path to the folder where the RAMSES outputs are stored. csiborg2_main_srcdir : str
postdir: str, optional Path to the CSiBORG2 main simulation directory.
Path to the folder where post-processed files are stored. csiborg2_random_srcdir : str
borg_dir : str, optional Path to the CSiBORG2 random simulation directory.
Path to the folder where BORG MCMC chains are stored. csiborg2_varysmall_srcdir : str
quiote_dir : str, optional Path to the CSiBORG2 varysmall simulation directory.
Path to the folder where Quijote simulations are stored. postdir : str
Path to the CSiBORG post-processing directory.
quijote_dir : str
Path to the Quijote simulation directory.
""" """
def __init__(self, def __init__(self,
csiborg1_srcdir=None, csiborg1_srcdir,
csiborg2_main_srcdir=None, csiborg2_main_srcdir,
csiborg2_random_srcdir=None, csiborg2_random_srcdir,
csiborg2_varysmall_srcdir=None, csiborg2_varysmall_srcdir,
postdir=None, postdir,
quijote_dir=None quijote_dir,
): ):
self.csiborg1_srcdir = csiborg1_srcdir self.csiborg1_srcdir = csiborg1_srcdir
self.csiborg2_main_srcdir = csiborg2_main_srcdir self.csiborg2_main_srcdir = csiborg2_main_srcdir
@ -117,8 +120,6 @@ class Paths:
------- -------
snapshots : 1-dimensional array snapshots : 1-dimensional array
""" """
# simpath = self.snapshots(nsim, simname, tonew=False)
if simname == "csiborg1": if simname == "csiborg1":
snaps = glob(join(self.csiborg1_srcdir, f"chain_{nsim}", snaps = glob(join(self.csiborg1_srcdir, f"chain_{nsim}",
"snapshot_*")) "snapshot_*"))
@ -176,14 +177,14 @@ class Paths:
return join(self.csiborg1_srcdir, f"chain_{nsim}", return join(self.csiborg1_srcdir, f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(5)}.hdf5") f"snapshot_{str(nsnap).zfill(5)}.hdf5")
elif simname == "csiborg2_main": elif simname == "csiborg2_main":
return join(self.csiborg2_main_srcdir, f"chain_{nsim}", return join(self.csiborg2_main_srcdir, f"chain_{nsim}", "output",
f"snapshot_{str(nsnap).zfill(3)}.hdf5") f"snapshot_{str(nsnap).zfill(3)}_full.hdf5")
elif simname == "csiborg2_random": elif simname == "csiborg2_random":
return join(self.csiborg2_random_srcdir, f"chain_{nsim}", return join(self.csiborg2_random_srcdir, f"chain_{nsim}", "output",
f"snapshot_{str(nsnap).zfill(3)}.hdf5") f"snapshot_{str(nsnap).zfill(3)}_full.hdf5")
elif simname == "csiborg2_varysmall": elif simname == "csiborg2_varysmall":
return join(self.csiborg2_varysmall_srcdir, f"chain_{nsim}", 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": elif simname == "quijote":
return join(self.quijote_dir, "fiducial_processed", return join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}", f"chain_{nsim}",
@ -191,6 +192,43 @@ class Paths:
else: else:
raise ValueError(f"Unknown simulation name `{simname}`.") 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): def overlap(self, simname, nsim0, nsimx, min_logmass, smoothed):
""" """
Path to the overlap files between two CSiBORG simulations. Path to the overlap files between two CSiBORG simulations.
@ -274,48 +312,77 @@ class Paths:
return join(fdir, fname) 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""" r"""
Path to the files containing the calculated fields in CSiBORG. Path to the files containing the calculated fields in CSiBORG.
Parameters Parameters
---------- ----------
kind : str kind : str
Field type. Must be one of: `density`, `velocity`, `potential`, Field type.
`radvel`, `environment`.
MAS : str MAS : str
Mass-assignment scheme. Mass-assignment scheme.
grid : int grid : int
Grid size. Grid size.
nsim : int nsim : int
IC realisation index. IC realisation index.
in_rsp : bool simname : str
Whether the calculation is performed in redshift space. Simulation name.
smooth_scale : float, optional
Smoothing scale in Mpc/h.
Returns Returns
------- -------
str str
""" """
assert kind in ["density", "velocity", "potential", "radvel", if MAS == "SPH":
"environment"] if kind not in ["density", "velocity"]:
fdir = join(self.postdir, "environment") 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) try_create_directory(fdir)
if in_rsp: fname = f"{kind}_{simname}_{MAS}_{str(nsim).zfill(5)}_{grid}.npy"
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")
return join(fdir, fname) return join(fdir, fname)
def field_interpolated(self, survey, kind, MAS, grid, nsim, in_rsp, def observer_peculiar_velocity(self, MAS, grid, nsim, simname):
smooth_scale=None): """
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 Path to the files containing the CSiBORG interpolated field for a given
survey. survey.
@ -335,13 +402,12 @@ class Paths:
IC realisation index. IC realisation index.
in_rsp : bool in_rsp : bool
Whether the calculation is performed in redshift space. Whether the calculation is performed in redshift space.
smooth_scale : float, optional
Smoothing scale in Mpc/h.
Returns Returns
------- -------
str str
""" """
raise NotImplementedError("This function is not implemented yet.")
assert kind in ["density", "velocity", "potential", "radvel", assert kind in ["density", "velocity", "potential", "radvel",
"environment"] "environment"]
fdir = join(self.postdir, "environment_interpolated") fdir = join(self.postdir, "environment_interpolated")
@ -353,9 +419,6 @@ class Paths:
fname = f"{survey}_{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npz" 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) return join(fdir, fname)
def cross_nearest(self, simname, run, kind, nsim=None, nobs=None): 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 # 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., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""" """MPI script to calculate the various fields."""
MPI script to calculate density field-derived fields in the CSiBORG
simulations' final snapshot.
"""
from argparse import ArgumentParser from argparse import ArgumentParser
from datetime import datetime from datetime import datetime
from distutils.util import strtobool
from gc import collect
import numpy import numpy
from mpi4py import MPI from mpi4py import MPI
@ -29,55 +24,43 @@ import csiborgtools
from utils import get_nsims from utils import get_nsims
###############################################################################
# Cosmotool SPH density & velocity field #
###############################################################################
def cosmotool_sph(nsim, parser_args):
pass
############################################################################### ###############################################################################
# Density field # # Density field #
############################################################################### ###############################################################################
def density_field(nsim, parser_args, to_save=True): def density_field(nsim, parser_args):
""" """Calculate the density field."""
Calculate the density field in the CSiBORG simulation. if parser_args.MAS == "SPH":
""" raise NotImplementedError("SPH is not implemented here. Use cosmotool")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg")) nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
if not parser_args.in_rsp: # Read in the particle coordinates and masses
# TODO I removed this function if parser_args.simname == "csiborg1":
snap = csiborgtools.read.read_h5(fname)["snapshot_final"] snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
pos = snap["pos"] elif "csiborg2" in parser_args.simname:
mass = snap["mass"] kind = parser_args.simname.split("_")[-1]
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
gen = csiborgtools.field.DensityField(box, parser_args.MAS) elif parser_args.simname == "quijote":
field = gen(pos, mass, parser_args.grid, verbose=parser_args.verbose) snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
else: else:
field = numpy.load(paths.field( raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
"density", parser_args.MAS, parser_args.grid, nsim, False))
radvel_field = numpy.load(paths.field(
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
if parser_args.verbose: pos = snapshot.coordinates()
print(f"{datetime.now()}: converting density field to RSP.", mass = snapshot.masses()
flush=True)
field = csiborgtools.field.field2rsp(field, radvel_field, box, # Run the field generator
parser_args.MAS) 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("density", parser_args.MAS, parser_args.grid,
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid, nsim, parser_args.simname)
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.") print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field) numpy.save(fout, field)
return 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): def velocity_field(nsim, parser_args):
""" """Calculate the velocity field."""
Calculate the velocity field in a CSiBORG simulation. if parser_args.MAS == "SPH":
""" raise NotImplementedError("SPH is not implemented here. Use cosmotool")
if parser_args.in_rsp:
raise NotImplementedError("Velocity field in RSP is not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg")) nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
snap = csiborgtools.read.read_h5(fname)["snapshot_final"] if parser_args.simname == "csiborg1":
pos = snap["pos"] snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
vel = snap["vel"] elif "csiborg2" in parser_args.simname:
mass = snap["mass"] 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) pos = snapshot.coordinates()
field = gen(pos, vel, mass, parser_args.grid, verbose=parser_args.verbose) vel = snapshot.velocities()
mass = snapshot.masses()
if to_save: boxsize = csiborgtools.simname2boxsize(parser_args.simname)
fout = paths.field("velocity", parser_args.MAS, parser_args.grid, gen = csiborgtools.field.VelocityField(boxsize, parser_args.MAS)
nsim, in_rsp=False) field = gen(pos, vel, mass, parser_args.grid)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field) 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 return field
@ -119,125 +107,62 @@ def velocity_field(nsim, parser_args, to_save=True):
############################################################################### ###############################################################################
def radvel_field(nsim, parser_args, to_save=True): def radvel_field(nsim, parser_args):
""" """Calculate the radial velocity field."""
Calculate the radial velocity field in the CSiBORG simulation.
"""
if parser_args.in_rsp:
raise NotImplementedError("Radial vel. field in RSP not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) 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, if parser_args.simname == "csiborg1":
nsim, parser_args.in_rsp)) field = csiborgtools.read.CSiBORG1Field(nsim, paths)
observer_velocity = csiborgtools.field.observer_vobs(vel) elif "csiborg2" in parser_args.simname:
kind = parser_args.simname.split("_")[-1]
gen = csiborgtools.field.VelocityField(box, parser_args.MAS) field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
field = gen.radial_velocity(vel, observer_velocity) elif parser_args.simname == "quijote":
field = csiborgtools.read.QuijoteField(nsim, paths)
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)
else: else:
field = numpy.load(paths.field( raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
"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, vel = field.velocity_field(parser_args.MAS, parser_args.grid)
parser_args.MAS)
if to_save: observer_velocity = csiborgtools.field.observer_peculiar_velocity(vel)
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid, radvel = csiborgtools.field.radial_velocity(vel, observer_velocity)
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.") fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
numpy.save(fout, field) nsim, parser_args.simname)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, radvel)
return field return field
############################################################################### def observer_peculiar_velocity(nsim, parser_args):
# Environment classification #
###############################################################################
def environment_field(nsim, parser_args, to_save=True):
""" """
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) boxsize = csiborgtools.simname2boxsize(parser_args.simname)
nsnap = max(paths.get_snapshots(nsim, "csiborg")) # NOTE thevse values are hard-coded.
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths) 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( if parser_args.simname == "csiborg1":
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False)) field = csiborgtools.read.CSiBORG1Field(nsim, paths)
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS) elif "csiborg2" in parser_args.simname:
rho = density_gen.overdensity_field(rho) 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: vel = field.velocity_field(parser_args.MAS, parser_args.grid)
rho = csiborgtools.field.smoothen_field(
rho, parser_args.smooth_scale, box.box2mpc(1.))
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS) observer_vp = csiborgtools.field.observer_peculiar_velocity(
field = gen(rho) vel, smooth_scales)
del rho fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
collect() nsim, parser_args.simname)
print(f"Saving to ... `{fout}`")
if parser_args.in_rsp: numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
radvel_field = numpy.load(paths.field( return observer_vp
"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
############################################################################### ###############################################################################
@ -249,39 +174,124 @@ if __name__ == "__main__":
parser = ArgumentParser() parser = ArgumentParser()
parser.add_argument("--nsims", type=int, nargs="+", default=None, parser.add_argument("--nsims", type=int, nargs="+", default=None,
help="IC realisations. `-1` for all simulations.") help="IC realisations. `-1` for all simulations.")
parser.add_argument("--simname", type=str, help="Simulation name.")
parser.add_argument("--kind", type=str, parser.add_argument("--kind", type=str,
choices=["density", "rspdensity", "velocity", "radvel", choices=["density", "velocity", "radvel", "observer_vp"], # noqa
"potential", "environment"],
help="What derived field to calculate?") help="What derived field to calculate?")
parser.add_argument("--MAS", type=str, 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("--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() parser_args = parser.parse_args()
comm = MPI.COMM_WORLD comm = MPI.COMM_WORLD
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = get_nsims(parser_args, paths) nsims = get_nsims(parser_args, paths)
def main(nsim): def main(nsim):
if parser_args.kind == "density" or parser_args.kind == "rspdensity": if parser_args.kind == "density":
density_field(nsim, parser_args) density_field(nsim, parser_args)
elif parser_args.kind == "velocity": elif parser_args.kind == "velocity":
velocity_field(nsim, parser_args) velocity_field(nsim, parser_args)
elif parser_args.kind == "radvel": elif parser_args.kind == "radvel":
radvel_field(nsim, parser_args) radvel_field(nsim, parser_args)
elif parser_args.kind == "potential": elif parser_args.kind == "observer_vp":
potential_field(nsim, parser_args) observer_peculiar_velocity(nsim, parser_args)
elif parser_args.kind == "environment":
environment_field(nsim, parser_args)
else: else:
raise RuntimeError(f"Field {parser_args.kind} is not implemented.") raise RuntimeError(f"Field {parser_args.kind} is not implemented.")
work_delegation(main, nsims, comm, master_verbose=True) 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 from utils import get_nsims
# TODO get rid of this. # 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): 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() 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): def convert_str_to_num(s):
""" """
Convert a string representation of a number to its appropriate numeric type Convert a string representation of a number to its appropriate numeric type
@ -221,11 +211,6 @@ class CSiBORG1Reader:
raise ValueError(f"Unknown kind `{kind}`. " raise ValueError(f"Unknown kind `{kind}`. "
"Options are: `pid`, `pos`, `vel` or `mass`.") "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 del sim
collect() collect()
@ -273,8 +258,6 @@ class CSiBORG1Reader:
out["y"] = pos[:, 1] * h + 677.7 / 2 out["y"] = pos[:, 1] * h + 677.7 / 2
out["z"] = pos[:, 2] * 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["totpartmass"] = totmass * 1e11 * h
out["m200c"] = m200c * 1e11 * h out["m200c"] = m200c * 1e11 * h

5
scripts_independent/run_process_snapshot.py
View file

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