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