Clean density calculation (#97)

* 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

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 #
###############################################################################

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

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]

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)

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]

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

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 #
###############################################################################

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]

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):

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)