mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-12-22 00:58:03 +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
3
.flake8
Normal file
3
.flake8
Normal file
|
@ -0,0 +1,3 @@
|
|||
[flake8]
|
||||
exclude = *.ipynb
|
||||
|
|
@ -10,6 +10,8 @@ however with little effort it can support other simulations as well.
|
|||
## TODO
|
||||
- [x] Prune old CSiBORG1 merger tree things.
|
||||
- [x] Add visualiastion of the density field.
|
||||
- [ ] Clear out `density` support.
|
||||
- [ ] Add sorting of Gadget4 initial snapshot like final snapshot.
|
||||
- [ ] Add full support for CSiBORG2 suite of simulations.
|
||||
- [ ] Add SPH field calculation from cosmotools.
|
||||
|
||||
|
|
|
@ -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 #
|
||||
###############################################################################
|
||||
|
|
|
@ -13,7 +13,8 @@
|
|||
# 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
|
||||
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
|
||||
from .utils import nside2radec, smoothen_field # 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
|
||||
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
|
||||
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
|
||||
|
||||
###############################################################################
|
||||
|
@ -820,3 +818,8 @@ def load_halo_particles(hid, particles, hid2map):
|
|||
return particles[k0:kf + 1]
|
||||
except KeyError:
|
||||
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)
|
|
@ -1,108 +0,0 @@
|
|||
# Copyright (C) 2022 Richard Stiskalek
|
||||
# This program is free software; you can redistribute it and/or modify it
|
||||
# under the terms of the GNU General Public License as published by the
|
||||
# Free Software Foundation; either version 3 of the License, or (at your
|
||||
# option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful, but
|
||||
# WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
|
||||
# Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU General Public License along
|
||||
# with this program; if not, write to the Free Software Foundation, Inc.,
|
||||
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
|
||||
"""
|
||||
Script to calculate the peculiar velocity of an observer in the centre of the
|
||||
CSiBORG box.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from distutils.util import strtobool
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
from taskmaster import work_delegation
|
||||
from tqdm import tqdm
|
||||
from utils import get_nsims
|
||||
|
||||
try:
|
||||
import csiborgtools
|
||||
except ModuleNotFoundError:
|
||||
import sys
|
||||
sys.path.append("../")
|
||||
import csiborgtools
|
||||
|
||||
|
||||
def observer_peculiar_velocity(nsim, parser_args):
|
||||
"""
|
||||
Calculate the peculiar velocity of an observer in the centre of the box
|
||||
for several smoothing scales.
|
||||
"""
|
||||
pos = numpy.array([0.5, 0.5, 0.5]).reshape(-1, 3)
|
||||
boxsize = 677.7
|
||||
smooth_scales = [0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]
|
||||
|
||||
observer_vp = numpy.full((len(smooth_scales), 3), numpy.nan,
|
||||
dtype=numpy.float32)
|
||||
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
field_path = paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, in_rsp=False)
|
||||
field0 = numpy.load(field_path)
|
||||
|
||||
for j, smooth_scale in enumerate(tqdm(smooth_scales,
|
||||
desc="Smoothing the fields",
|
||||
disable=not parser_args.verbose)):
|
||||
if smooth_scale > 0:
|
||||
field = [None, None, None]
|
||||
for k in range(3):
|
||||
field[k] = csiborgtools.field.smoothen_field(
|
||||
field0[k], smooth_scale, boxsize)
|
||||
else:
|
||||
field = field0
|
||||
|
||||
v = csiborgtools.field.evaluate_cartesian(
|
||||
field[0], field[1], field[2], pos=pos)
|
||||
observer_vp[j, 0] = v[0][0]
|
||||
observer_vp[j, 1] = v[1][0]
|
||||
observer_vp[j, 2] = v[2][0]
|
||||
|
||||
fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
|
||||
nsim)
|
||||
if parser_args.verbose:
|
||||
print(f"Saving to ... `{fout}`")
|
||||
numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
|
||||
return observer_vp
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Command line interface #
|
||||
###############################################################################
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--nsims", type=int, nargs="+", default=None,
|
||||
help="IC realisations. `-1` for all simulations.")
|
||||
parser.add_argument("--kind", type=str,
|
||||
choices=["density", "rspdensity", "velocity", "radvel",
|
||||
"potential", "environment"],
|
||||
help="What derived field to calculate?")
|
||||
parser.add_argument("--MAS", type=str,
|
||||
choices=["NGP", "CIC", "TSC", "PCS"])
|
||||
parser.add_argument("--grid", type=int, help="Grid resolution.")
|
||||
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
|
||||
help="Verbosity flag for reading in particles.")
|
||||
parser.add_argument("--simname", type=str, default="csiborg",
|
||||
help="Verbosity flag for reading in particles.")
|
||||
parser_args = parser.parse_args()
|
||||
|
||||
comm = MPI.COMM_WORLD
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsims = get_nsims(parser_args, paths)
|
||||
|
||||
def main(nsim):
|
||||
return observer_peculiar_velocity(nsim, parser_args)
|
||||
|
||||
work_delegation(main, nsims, comm, master_verbose=True)
|
|
@ -12,14 +12,9 @@
|
|||
# You should have received a copy of the GNU General Public License along
|
||||
# with this program; if not, write to the Free Software Foundation, Inc.,
|
||||
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
|
||||
"""
|
||||
MPI script to calculate density field-derived fields in the CSiBORG
|
||||
simulations' final snapshot.
|
||||
"""
|
||||
"""MPI script to calculate the various fields."""
|
||||
from argparse import ArgumentParser
|
||||
from datetime import datetime
|
||||
from distutils.util import strtobool
|
||||
from gc import collect
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
@ -29,53 +24,41 @@ import csiborgtools
|
|||
from utils import get_nsims
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Cosmotool SPH density & velocity field #
|
||||
###############################################################################
|
||||
|
||||
def cosmotool_sph(nsim, parser_args):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Density field #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def density_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the density field in the CSiBORG simulation.
|
||||
"""
|
||||
def density_field(nsim, parser_args):
|
||||
"""Calculate the density field."""
|
||||
if parser_args.MAS == "SPH":
|
||||
raise NotImplementedError("SPH is not implemented here. Use cosmotool")
|
||||
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
|
||||
nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
|
||||
|
||||
if not parser_args.in_rsp:
|
||||
# TODO I removed this function
|
||||
snap = csiborgtools.read.read_h5(fname)["snapshot_final"]
|
||||
pos = snap["pos"]
|
||||
mass = snap["mass"]
|
||||
|
||||
gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
field = gen(pos, mass, parser_args.grid, verbose=parser_args.verbose)
|
||||
# Read in the particle coordinates and masses
|
||||
if parser_args.simname == "csiborg1":
|
||||
snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
|
||||
else:
|
||||
field = numpy.load(paths.field(
|
||||
"density", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
radvel_field = numpy.load(paths.field(
|
||||
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
if parser_args.verbose:
|
||||
print(f"{datetime.now()}: converting density field to RSP.",
|
||||
flush=True)
|
||||
pos = snapshot.coordinates()
|
||||
mass = snapshot.masses()
|
||||
|
||||
field = csiborgtools.field.field2rsp(field, radvel_field, box,
|
||||
parser_args.MAS)
|
||||
# Run the field generator
|
||||
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
|
||||
gen = csiborgtools.field.DensityField(boxsize, parser_args.MAS)
|
||||
field = gen(pos, mass, parser_args.grid)
|
||||
|
||||
fout = paths.field("density", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
return field
|
||||
|
@ -86,29 +69,34 @@ def density_field(nsim, parser_args, to_save=True):
|
|||
###############################################################################
|
||||
|
||||
|
||||
def velocity_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the velocity field in a CSiBORG simulation.
|
||||
"""
|
||||
if parser_args.in_rsp:
|
||||
raise NotImplementedError("Velocity field in RSP is not implemented.")
|
||||
def velocity_field(nsim, parser_args):
|
||||
"""Calculate the velocity field."""
|
||||
if parser_args.MAS == "SPH":
|
||||
raise NotImplementedError("SPH is not implemented here. Use cosmotool")
|
||||
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
|
||||
nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
|
||||
|
||||
snap = csiborgtools.read.read_h5(fname)["snapshot_final"]
|
||||
pos = snap["pos"]
|
||||
vel = snap["vel"]
|
||||
mass = snap["mass"]
|
||||
if parser_args.simname == "csiborg1":
|
||||
snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
|
||||
field = gen(pos, vel, mass, parser_args.grid, verbose=parser_args.verbose)
|
||||
pos = snapshot.coordinates()
|
||||
vel = snapshot.velocities()
|
||||
mass = snapshot.masses()
|
||||
|
||||
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
|
||||
gen = csiborgtools.field.VelocityField(boxsize, parser_args.MAS)
|
||||
field = gen(pos, vel, mass, parser_args.grid)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, in_rsp=False)
|
||||
nsim, parser_args.simname)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
return field
|
||||
|
@ -119,125 +107,62 @@ def velocity_field(nsim, parser_args, to_save=True):
|
|||
###############################################################################
|
||||
|
||||
|
||||
def radvel_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the radial velocity field in the CSiBORG simulation.
|
||||
"""
|
||||
if parser_args.in_rsp:
|
||||
raise NotImplementedError("Radial vel. field in RSP not implemented.")
|
||||
|
||||
def radvel_field(nsim, parser_args):
|
||||
"""Calculate the radial velocity field."""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
|
||||
vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp))
|
||||
observer_velocity = csiborgtools.field.observer_vobs(vel)
|
||||
|
||||
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
|
||||
field = gen.radial_velocity(vel, observer_velocity)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
return field
|
||||
|
||||
###############################################################################
|
||||
# Potential field #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def potential_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the potential field in the CSiBORG simulation.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
|
||||
if not parser_args.in_rsp:
|
||||
rho = numpy.load(paths.field(
|
||||
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
|
||||
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
rho = density_gen.overdensity_field(rho)
|
||||
|
||||
gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
|
||||
field = gen(rho)
|
||||
if parser_args.simname == "csiborg1":
|
||||
field = csiborgtools.read.CSiBORG1Field(nsim, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
field = csiborgtools.read.QuijoteField(nsim, paths)
|
||||
else:
|
||||
field = numpy.load(paths.field(
|
||||
"potential", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
radvel_field = numpy.load(paths.field(
|
||||
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
field = csiborgtools.field.field2rsp(field, radvel_field, box,
|
||||
parser_args.MAS)
|
||||
vel = field.velocity_field(parser_args.MAS, parser_args.grid)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp)
|
||||
observer_velocity = csiborgtools.field.observer_peculiar_velocity(vel)
|
||||
radvel = csiborgtools.field.radial_velocity(vel, observer_velocity)
|
||||
|
||||
fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
numpy.save(fout, radvel)
|
||||
return field
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Environment classification #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def environment_field(nsim, parser_args, to_save=True):
|
||||
def observer_peculiar_velocity(nsim, parser_args):
|
||||
"""
|
||||
Calculate the environmental classification in the CSiBORG simulation.
|
||||
Calculate the peculiar velocity of an observer in the centre of the box
|
||||
for several smoothing scales.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
|
||||
# NOTE thevse values are hard-coded.
|
||||
smooth_scales = numpy.array([0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
|
||||
smooth_scales /= boxsize
|
||||
|
||||
rho = numpy.load(paths.field(
|
||||
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
|
||||
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
rho = density_gen.overdensity_field(rho)
|
||||
if parser_args.simname == "csiborg1":
|
||||
field = csiborgtools.read.CSiBORG1Field(nsim, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
field = csiborgtools.read.QuijoteField(nsim, paths)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
if parser_args.smooth_scale > 0.0:
|
||||
rho = csiborgtools.field.smoothen_field(
|
||||
rho, parser_args.smooth_scale, box.box2mpc(1.))
|
||||
vel = field.velocity_field(parser_args.MAS, parser_args.grid)
|
||||
|
||||
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
|
||||
field = gen(rho)
|
||||
observer_vp = csiborgtools.field.observer_peculiar_velocity(
|
||||
vel, smooth_scales)
|
||||
|
||||
del rho
|
||||
collect()
|
||||
|
||||
if parser_args.in_rsp:
|
||||
radvel_field = numpy.load(paths.field(
|
||||
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
args = (radvel_field, box, parser_args.MAS)
|
||||
|
||||
field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
|
||||
field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
|
||||
field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
|
||||
field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
|
||||
field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
|
||||
field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
|
||||
|
||||
del radvel_field
|
||||
collect()
|
||||
|
||||
eigvals = gen.tensor_field_eigvals(field)
|
||||
|
||||
del field
|
||||
collect()
|
||||
|
||||
env = gen.eigvals_to_environment(eigvals)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field("environment", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp, parser_args.smooth_scale)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, env)
|
||||
return env
|
||||
fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
print(f"Saving to ... `{fout}`")
|
||||
numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
|
||||
return observer_vp
|
||||
|
||||
|
||||
###############################################################################
|
||||
|
@ -249,39 +174,124 @@ if __name__ == "__main__":
|
|||
parser = ArgumentParser()
|
||||
parser.add_argument("--nsims", type=int, nargs="+", default=None,
|
||||
help="IC realisations. `-1` for all simulations.")
|
||||
parser.add_argument("--simname", type=str, help="Simulation name.")
|
||||
parser.add_argument("--kind", type=str,
|
||||
choices=["density", "rspdensity", "velocity", "radvel",
|
||||
"potential", "environment"],
|
||||
choices=["density", "velocity", "radvel", "observer_vp"], # noqa
|
||||
help="What derived field to calculate?")
|
||||
parser.add_argument("--MAS", type=str,
|
||||
choices=["NGP", "CIC", "TSC", "PCS"])
|
||||
choices=["NGP", "CIC", "TSC", "PCS", "SPH"],
|
||||
help="Mass assignment scheme.")
|
||||
parser.add_argument("--grid", type=int, help="Grid resolution.")
|
||||
parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
|
||||
help="Calculate in RSP?")
|
||||
parser.add_argument("--smooth_scale", type=float, default=0.0,
|
||||
help="Smoothing scale in Mpc / h. Only used for the environment field.") # noqa
|
||||
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
|
||||
help="Verbosity flag for reading in particles.")
|
||||
parser.add_argument("--simname", type=str, default="csiborg",
|
||||
choices=["csiborg", "csiborg2"],
|
||||
help="Verbosity flag for reading in particles.")
|
||||
parser_args = parser.parse_args()
|
||||
|
||||
comm = MPI.COMM_WORLD
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsims = get_nsims(parser_args, paths)
|
||||
|
||||
def main(nsim):
|
||||
if parser_args.kind == "density" or parser_args.kind == "rspdensity":
|
||||
if parser_args.kind == "density":
|
||||
density_field(nsim, parser_args)
|
||||
elif parser_args.kind == "velocity":
|
||||
velocity_field(nsim, parser_args)
|
||||
elif parser_args.kind == "radvel":
|
||||
radvel_field(nsim, parser_args)
|
||||
elif parser_args.kind == "potential":
|
||||
potential_field(nsim, parser_args)
|
||||
elif parser_args.kind == "environment":
|
||||
environment_field(nsim, parser_args)
|
||||
elif parser_args.kind == "observer_vp":
|
||||
observer_peculiar_velocity(nsim, parser_args)
|
||||
else:
|
||||
raise RuntimeError(f"Field {parser_args.kind} is not implemented.")
|
||||
|
||||
work_delegation(main, nsims, comm, master_verbose=True)
|
||||
|
||||
|
||||
# def potential_field(nsim, parser_args, to_save=True):
|
||||
# """
|
||||
# Calculate the potential field in the CSiBORG simulation.
|
||||
# """
|
||||
# paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
# nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
# box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
#
|
||||
# if not parser_args.in_rsp:
|
||||
# rho = numpy.load(paths.field(
|
||||
# "density", parser_args.MAS, parser_args.grid, nsim,
|
||||
# in_rsp=False))
|
||||
# density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
# rho = density_gen.overdensity_field(rho)
|
||||
#
|
||||
# gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
|
||||
# field = gen(rho)
|
||||
# else:
|
||||
# field = numpy.load(paths.field(
|
||||
# "potential", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
# radvel_field = numpy.load(paths.field(
|
||||
# "radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
#
|
||||
# field = csiborgtools.field.field2rsp(field, radvel_field, box,
|
||||
# parser_args.MAS)
|
||||
#
|
||||
# if to_save:
|
||||
# fout = paths.field(parser_args.kind, parser_args.MAS,
|
||||
# parser_args.grid,
|
||||
# nsim, parser_args.in_rsp)
|
||||
# print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
# numpy.save(fout, field)
|
||||
# return field
|
||||
#
|
||||
#
|
||||
# #############################################################################
|
||||
# # Environment classification #
|
||||
# #############################################################################
|
||||
#
|
||||
#
|
||||
# def environment_field(nsim, parser_args, to_save=True):
|
||||
# """
|
||||
# Calculate the environmental classification in the CSiBORG simulation.
|
||||
# """
|
||||
# paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
# nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
# box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
#
|
||||
# rho = numpy.load(paths.field(
|
||||
# "density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
|
||||
# density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
# rho = density_gen.overdensity_field(rho)
|
||||
#
|
||||
# if parser_args.smooth_scale > 0.0:
|
||||
# rho = csiborgtools.field.smoothen_field(
|
||||
# rho, parser_args.smooth_scale, box.box2mpc(1.))
|
||||
#
|
||||
# gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
|
||||
# field = gen(rho)
|
||||
#
|
||||
# del rho
|
||||
# collect()
|
||||
#
|
||||
# if parser_args.in_rsp:
|
||||
# radvel_field = numpy.load(paths.field(
|
||||
# "radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
# args = (radvel_field, box, parser_args.MAS)
|
||||
#
|
||||
# field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
|
||||
# field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
|
||||
# field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
|
||||
# field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
|
||||
# field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
|
||||
# field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
|
||||
#
|
||||
# del radvel_field
|
||||
# collect()
|
||||
#
|
||||
# eigvals = gen.tensor_field_eigvals(field)
|
||||
#
|
||||
# del field
|
||||
# collect()
|
||||
#
|
||||
# env = gen.eigvals_to_environment(eigvals)
|
||||
#
|
||||
# if to_save:
|
||||
# fout = paths.field("environment", parser_args.MAS, parser_args.grid,
|
||||
# nsim, parser_args.in_rsp,
|
||||
# parser_args.smooth_scale)
|
||||
# print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
# numpy.save(fout, env)
|
||||
# return env
|
||||
|
|
|
@ -30,7 +30,12 @@ import csiborgtools
|
|||
from utils import get_nsims
|
||||
|
||||
# TODO get rid of this.
|
||||
MPC2BOX = 1 / 677.7
|
||||
# MPC2BOX = 1 / 677.7
|
||||
SIM2BOXSIZE = {"csiborg1": 677.7,
|
||||
"csiborg2_main": None,
|
||||
"csiborg2_random": None,
|
||||
"csiborg2_varysmall": None,
|
||||
}
|
||||
|
||||
|
||||
def steps(cls, survey_name):
|
||||
|
|
|
@ -1,374 +0,0 @@
|
|||
# Copyright (C) 2022 Richard Stiskalek
|
||||
# This program is free software; you can redistribute it and/or modify it
|
||||
# under the terms of the GNU General Public License as published by the
|
||||
# Free Software Foundation; either version 3 of the License, or (at your
|
||||
# option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful, but
|
||||
# WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
|
||||
# Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU General Public License along
|
||||
# with this program; if not, write to the Free Software Foundation, Inc.,
|
||||
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
|
||||
r"""
|
||||
Script to process simulation files and create a single HDF5 file, in which
|
||||
particles are sorted by the particle halo IDs.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from gc import collect
|
||||
|
||||
import h5py
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
import csiborgtools
|
||||
from csiborgtools import fprint
|
||||
from numba import jit
|
||||
from taskmaster import work_delegation
|
||||
from tqdm import trange, tqdm
|
||||
from utils import get_nsims
|
||||
|
||||
|
||||
@jit(nopython=True, boundscheck=False)
|
||||
def minmax_halo(hid, halo_ids, start_loop=0):
|
||||
"""
|
||||
Find the start and end index of a halo in a sorted array of halo IDs.
|
||||
This is much faster than using `numpy.where` and then `numpy.min` and
|
||||
`numpy.max`.
|
||||
"""
|
||||
start = None
|
||||
end = None
|
||||
|
||||
for i in range(start_loop, halo_ids.size):
|
||||
n = halo_ids[i]
|
||||
if n == hid:
|
||||
if start is None:
|
||||
start = i
|
||||
end = i
|
||||
elif n > hid:
|
||||
break
|
||||
return start, end
|
||||
|
||||
|
||||
def process_snapshot(nsim, simname, halo_finder, verbose):
|
||||
"""
|
||||
Read in the snapshot particles, sort them by their halo ID and dump
|
||||
into a HDF5 file. Stores the first and last index of each halo in the
|
||||
particle array for fast slicing of the array to acces particles of a single
|
||||
halo.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, simname))
|
||||
|
||||
if simname == "csiborg":
|
||||
partreader = csiborgtools.read.CSiBORGReader(paths)
|
||||
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
|
||||
else:
|
||||
partreader = csiborgtools.read.QuijoteReader(paths)
|
||||
box = None
|
||||
|
||||
desc = {"hid": f"Halo finder ID ({halo_finder})of the particle.",
|
||||
"pos": "DM particle positions in box units.",
|
||||
"vel": "DM particle velocity in km / s.",
|
||||
"mass": "DM particle mass in Msun / h.",
|
||||
"pid": "DM particle ID",
|
||||
}
|
||||
|
||||
fname = paths.processed_output(nsim, simname, halo_finder)
|
||||
|
||||
fprint(f"loading HIDs of IC {nsim}.", verbose)
|
||||
hids = partreader.read_halo_id(nsnap, nsim, halo_finder, verbose)
|
||||
collect()
|
||||
|
||||
fprint(f"sorting HIDs of IC {nsim}.")
|
||||
sort_indxs = numpy.argsort(hids)
|
||||
|
||||
with h5py.File(fname, "w") as f:
|
||||
group = f.create_group("snapshot_final")
|
||||
group.attrs["header"] = "Snapshot data at z = 0."
|
||||
|
||||
fprint("dumping halo IDs.", verbose)
|
||||
dset = group.create_dataset("halo_ids", data=hids[sort_indxs])
|
||||
dset.attrs["header"] = desc["hid"]
|
||||
del hids
|
||||
collect()
|
||||
|
||||
fprint("reading, sorting and dumping the snapshot particles.", verbose)
|
||||
for kind in ["pos", "vel", "mass", "pid"]:
|
||||
x = partreader.read_snapshot(nsnap, nsim, kind)[sort_indxs]
|
||||
|
||||
if simname == "csiborg" and kind == "vel":
|
||||
x = box.box2vel(x) if simname == "csiborg" else x
|
||||
|
||||
if simname == "csiborg" and kind == "mass":
|
||||
x = box.box2solarmass(x) if simname == "csiborg" else x
|
||||
|
||||
dset = f["snapshot_final"].create_dataset(kind, data=x)
|
||||
dset.attrs["header"] = desc[kind]
|
||||
del x
|
||||
collect()
|
||||
|
||||
del sort_indxs
|
||||
collect()
|
||||
|
||||
fprint(f"creating a halo map for IC {nsim}.")
|
||||
with h5py.File(fname, "r") as f:
|
||||
part_hids = f["snapshot_final"]["halo_ids"][:]
|
||||
# We loop over the unique halo IDs and remove the 0 halo ID
|
||||
unique_halo_ids = numpy.unique(part_hids)
|
||||
unique_halo_ids = unique_halo_ids[unique_halo_ids != 0]
|
||||
halo_map = numpy.full((unique_halo_ids.size, 3), numpy.nan,
|
||||
dtype=numpy.uint64)
|
||||
start_loop, niters = 0, unique_halo_ids.size
|
||||
for i in trange(niters, disable=not verbose):
|
||||
hid = unique_halo_ids[i]
|
||||
k0, kf = minmax_halo(hid, part_hids, start_loop=start_loop)
|
||||
halo_map[i, :] = hid, k0, kf
|
||||
start_loop = kf
|
||||
|
||||
# Dump the halo mapping.
|
||||
with h5py.File(fname, "r+") as f:
|
||||
dset = f["snapshot_final"].create_dataset("halo_map", data=halo_map)
|
||||
dset.attrs["header"] = """
|
||||
Halo to particle mapping. Columns are HID, start index, end index.
|
||||
"""
|
||||
f.close()
|
||||
|
||||
del part_hids
|
||||
collect()
|
||||
|
||||
# Add the halo finder catalogue
|
||||
with h5py.File(fname, "r+") as f:
|
||||
group = f.create_group("halofinder_catalogue")
|
||||
group.attrs["header"] = f"Original {halo_finder} halo catalogue."
|
||||
cat = partreader.read_catalogue(nsnap, nsim, halo_finder)
|
||||
|
||||
hid2pos = {hid: i for i, hid in enumerate(unique_halo_ids)}
|
||||
|
||||
for key in cat.dtype.names:
|
||||
x = numpy.full(unique_halo_ids.size, numpy.nan,
|
||||
dtype=cat[key].dtype)
|
||||
for i in range(len(cat)):
|
||||
j = hid2pos[cat["index"][i]]
|
||||
x[j] = cat[key][i]
|
||||
group.create_dataset(key, data=x)
|
||||
f.close()
|
||||
|
||||
# Lastly create the halo catalogue
|
||||
with h5py.File(fname, "r+") as f:
|
||||
group = f.create_group("halo_catalogue")
|
||||
group.attrs["header"] = f"{halo_finder} halo catalogue."
|
||||
group.create_dataset("index", data=unique_halo_ids)
|
||||
f.close()
|
||||
|
||||
|
||||
def add_initial_snapshot(nsim, simname, halo_finder, verbose):
|
||||
"""
|
||||
Sort the initial snapshot particles according to their final snapshot and
|
||||
add them to the final snapshot's HDF5 file.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
fname = paths.processed_output(nsim, simname, halo_finder)
|
||||
|
||||
if simname == "csiborg":
|
||||
partreader = csiborgtools.read.CSiBORGReader(paths)
|
||||
else:
|
||||
partreader = csiborgtools.read.QuijoteReader(paths)
|
||||
|
||||
fprint(f"processing simulation `{nsim}`.", verbose)
|
||||
if simname == "csiborg":
|
||||
nsnap0 = 1
|
||||
elif simname == "quijote":
|
||||
nsnap0 = -1
|
||||
else:
|
||||
raise ValueError(f"Unknown simulation `{simname}`.")
|
||||
|
||||
fprint("loading and sorting the initial PID.", verbose)
|
||||
sort_indxs = numpy.argsort(partreader.read_snapshot(nsnap0, nsim, "pid"))
|
||||
|
||||
fprint("loading the final particles.", verbose)
|
||||
with h5py.File(fname, "r") as f:
|
||||
sort_indxs_final = f["snapshot_final/pid"][:]
|
||||
f.close()
|
||||
|
||||
fprint("sorting the particles according to the final snapshot.", verbose)
|
||||
sort_indxs_final = numpy.argsort(numpy.argsort(sort_indxs_final))
|
||||
sort_indxs = sort_indxs[sort_indxs_final]
|
||||
|
||||
del sort_indxs_final
|
||||
collect()
|
||||
|
||||
fprint("loading and sorting the initial particle position.", verbose)
|
||||
pos = partreader.read_snapshot(nsnap0, nsim, "pos")[sort_indxs]
|
||||
|
||||
del sort_indxs
|
||||
collect()
|
||||
|
||||
# In Quijote some particles are position precisely at the edge of the
|
||||
# box. Move them to be just inside.
|
||||
if simname == "quijote":
|
||||
mask = pos >= 1
|
||||
if numpy.any(mask):
|
||||
spacing = numpy.spacing(pos[mask])
|
||||
assert numpy.max(spacing) <= 1e-5
|
||||
pos[mask] -= spacing
|
||||
|
||||
fprint(f"dumping particles for `{nsim}` to `{fname}`.", verbose)
|
||||
with h5py.File(fname, "r+") as f:
|
||||
if "snapshot_initial" in f.keys():
|
||||
del f["snapshot_initial"]
|
||||
group = f.create_group("snapshot_initial")
|
||||
group.attrs["header"] = "Initial snapshot data."
|
||||
dset = group.create_dataset("pos", data=pos)
|
||||
dset.attrs["header"] = "DM particle positions in box units."
|
||||
|
||||
f.close()
|
||||
|
||||
|
||||
def calculate_initial(nsim, simname, halo_finder, verbose):
|
||||
"""Calculate the Lagrangian patch centre of mass and size."""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
|
||||
fname = paths.processed_output(nsim, simname, halo_finder)
|
||||
fprint("loading the particle information.", verbose)
|
||||
f = h5py.File(fname, "r")
|
||||
pos = f["snapshot_initial/pos"]
|
||||
mass = f["snapshot_final/mass"]
|
||||
hid = f["halo_catalogue/index"][:]
|
||||
hid2map = csiborgtools.read.make_halomap_dict(
|
||||
f["snapshot_final/halo_map"][:])
|
||||
|
||||
if simname == "csiborg":
|
||||
kwargs = {"box_size": 2048, "bckg_halfsize": 512}
|
||||
else:
|
||||
kwargs = {"box_size": 512, "bckg_halfsize": 256}
|
||||
overlapper = csiborgtools.match.ParticleOverlap(**kwargs)
|
||||
|
||||
lagpatch_pos = numpy.full((len(hid), 3), numpy.nan, dtype=numpy.float32)
|
||||
lagpatch_size = numpy.full(len(hid), numpy.nan, dtype=numpy.float32)
|
||||
lagpatch_ncells = numpy.full(len(hid), numpy.nan, dtype=numpy.int32)
|
||||
|
||||
for i in trange(len(hid), disable=not verbose):
|
||||
h = hid[i]
|
||||
# These are unasigned particles.
|
||||
if h == 0:
|
||||
continue
|
||||
|
||||
parts_pos = csiborgtools.read.load_halo_particles(h, pos, hid2map)
|
||||
parts_mass = csiborgtools.read.load_halo_particles(h, mass, hid2map)
|
||||
|
||||
# Skip if the halo has no particles or is too small.
|
||||
if parts_pos is None or parts_pos.size < 5:
|
||||
continue
|
||||
|
||||
cm = csiborgtools.center_of_mass(parts_pos, parts_mass, boxsize=1.0)
|
||||
sep = csiborgtools.periodic_distance(parts_pos, cm, boxsize=1.0)
|
||||
delta = overlapper.make_delta(parts_pos, parts_mass, subbox=True)
|
||||
|
||||
lagpatch_pos[i] = cm
|
||||
lagpatch_size[i] = numpy.percentile(sep, 99)
|
||||
lagpatch_ncells[i] = csiborgtools.delta2ncells(delta)
|
||||
|
||||
f.close()
|
||||
collect()
|
||||
|
||||
with h5py.File(fname, "r+") as f:
|
||||
grp = f["halo_catalogue"]
|
||||
dset = grp.create_dataset("lagpatch_pos", data=lagpatch_pos)
|
||||
dset.attrs["header"] = "Lagrangian patch centre of mass in box units."
|
||||
|
||||
dset = grp.create_dataset("lagpatch_size", data=lagpatch_size)
|
||||
dset.attrs["header"] = "Lagrangian patch size in box units."
|
||||
|
||||
dset = grp.create_dataset("lagpatch_ncells", data=lagpatch_ncells)
|
||||
dset.attrs["header"] = f"Lagrangian patch number of cells on a {kwargs['box_size']}^3 grid." # noqa
|
||||
|
||||
f.close()
|
||||
|
||||
|
||||
def make_phew_halo_catalogue(nsim, verbose):
|
||||
"""
|
||||
Process the PHEW halo catalogue for a CSiBORG simulation at all snapshots.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
snapshots = paths.get_snapshots(nsim, "csiborg")
|
||||
reader = csiborgtools.read.CSiBORGReader(paths)
|
||||
keys_write = ["index", "x", "y", "z", "mass_cl", "parent",
|
||||
"ultimate_parent", "summed_mass"]
|
||||
|
||||
# Create a HDF5 file to store all this.
|
||||
fname = paths.processed_phew(nsim)
|
||||
with h5py.File(fname, "w") as f:
|
||||
f.close()
|
||||
|
||||
for nsnap in tqdm(snapshots, disable=not verbose, desc="Snapshot"):
|
||||
try:
|
||||
data = reader.read_phew_clumps(nsnap, nsim, verbose=False)
|
||||
except FileExistsError:
|
||||
continue
|
||||
|
||||
with h5py.File(fname, "r+") as f:
|
||||
if str(nsnap) in f:
|
||||
print(f"Group {nsnap} already exists. Deleting.", flush=True)
|
||||
del f[str(nsnap)]
|
||||
grp = f.create_group(str(nsnap))
|
||||
for key in keys_write:
|
||||
grp.create_dataset(key, data=data[key])
|
||||
|
||||
grp.attrs["header"] = f"CSiBORG PHEW clumps at snapshot {nsnap}."
|
||||
f.close()
|
||||
|
||||
# Now write the redshifts
|
||||
scale_factors = numpy.full(len(snapshots), numpy.nan, dtype=numpy.float32)
|
||||
for i, nsnap in enumerate(snapshots):
|
||||
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
|
||||
scale_factors[i] = box._aexp
|
||||
|
||||
redshifts = scale_factors[-1] / scale_factors - 1
|
||||
|
||||
with h5py.File(fname, "r+") as f:
|
||||
grp = f.create_group("info")
|
||||
grp.create_dataset("redshift", data=redshifts)
|
||||
grp.create_dataset("snapshots", data=snapshots)
|
||||
grp.create_dataset("Om0", data=[box.Om0])
|
||||
grp.create_dataset("boxsize", data=[box.boxsize])
|
||||
f.close()
|
||||
|
||||
|
||||
def main(nsim, args):
|
||||
if args.make_final:
|
||||
process_snapshot(nsim, args.simname, args.halofinder, True)
|
||||
|
||||
if args.make_initial:
|
||||
add_initial_snapshot(nsim, args.simname, args.halofinder, True)
|
||||
calculate_initial(nsim, args.simname, args.halofinder, True)
|
||||
|
||||
if args.make_phew:
|
||||
make_phew_halo_catalogue(nsim, True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--simname", type=str, default="csiborg",
|
||||
choices=["csiborg", "quijote"],
|
||||
help="Simulation name")
|
||||
parser.add_argument("--nsims", type=int, nargs="+", default=None,
|
||||
help="IC realisations. If `-1` processes all.")
|
||||
parser.add_argument("--halofinder", type=str, help="Halo finder")
|
||||
parser.add_argument("--make_final", action="store_true", default=False,
|
||||
help="Process the final snapshot.")
|
||||
parser.add_argument("--make_initial", action="store_true", default=False,
|
||||
help="Process the initial snapshot.")
|
||||
parser.add_argument("--make_phew", action="store_true", default=False,
|
||||
help="Process the PHEW halo catalogue.")
|
||||
|
||||
args = parser.parse_args()
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsims = get_nsims(args, paths)
|
||||
|
||||
def _main(nsim):
|
||||
main(nsim, args)
|
||||
|
||||
work_delegation(_main, nsims, MPI.COMM_WORLD)
|
|
@ -60,16 +60,6 @@ def now():
|
|||
return datetime.now()
|
||||
|
||||
|
||||
def flip_cols(arr, col1, col2):
|
||||
"""
|
||||
Flip values in columns `col1` and `col2` of a structured array `arr`.
|
||||
"""
|
||||
if col1 not in arr.dtype.names or col2 not in arr.dtype.names:
|
||||
raise ValueError(f"Both `{col1}` and `{col2}` must exist in `arr`.")
|
||||
|
||||
arr[col1], arr[col2] = numpy.copy(arr[col2]), numpy.copy(arr[col1])
|
||||
|
||||
|
||||
def convert_str_to_num(s):
|
||||
"""
|
||||
Convert a string representation of a number to its appropriate numeric type
|
||||
|
@ -221,11 +211,6 @@ class CSiBORG1Reader:
|
|||
raise ValueError(f"Unknown kind `{kind}`. "
|
||||
"Options are: `pid`, `pos`, `vel` or `mass`.")
|
||||
|
||||
# Because of a RAMSES bug x and z are flipped.
|
||||
if kind in ["pos", "vel"]:
|
||||
print(f"For kind `{kind}` flipping x and z.")
|
||||
x[:, [0, 2]] = x[:, [2, 0]]
|
||||
|
||||
del sim
|
||||
collect()
|
||||
|
||||
|
@ -273,8 +258,6 @@ class CSiBORG1Reader:
|
|||
out["y"] = pos[:, 1] * h + 677.7 / 2
|
||||
out["z"] = pos[:, 2] * h + 677.7 / 2
|
||||
|
||||
# Because of a RAMSES bug x and z are flipped.
|
||||
flip_cols(out, "x", "z")
|
||||
out["totpartmass"] = totmass * 1e11 * h
|
||||
out["m200c"] = m200c * 1e11 * h
|
||||
|
||||
|
|
|
@ -15,9 +15,8 @@
|
|||
from os import system
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Quijote chains
|
||||
chains = [1]
|
||||
simname = "quijote"
|
||||
chains = [7444]
|
||||
simname = "csiborg1"
|
||||
mode = 2
|
||||
|
||||
env = "/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
|
||||
|
|
File diff suppressed because one or more lines are too long
Loading…
Reference in a new issue