Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-12-22 17:28:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Velocity observer (#86)

Browse source 
* Continue if r200c not defined

* Remove smooth scale

* Remove smooth scale

* Edit Max Matching plot

* Add peculiar velocity

* Add Vobs calculation

* Edit docs

* Add Vobs plot

* Improve plotting

* Edit naming convention

* Make a note

* Add new cat options

* Update density field RSP calculation

* Update field 2 rsp

* Move functions and shorten documentation

* Improve transforms and comments

* Update docs

* Update imports

* Edit calculation

* Add docs

* Remove imports

* Add Quijote flags

* Edit documentation

* Shorten documentation

* Edit func calls

* Shorten

* Docs edits

* Edit docs

* Shorten docs

* Short docs edits

* Simplify docs a little bit

* Save plotting

* Update env
This commit is contained in:
Richard Stiskalek 2023-08-30 23:27:20 +01:00 committed by GitHub
parent 8e3127f4d9
commit ae92fd9b72
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 761 additions and 788 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
csiborgtools/field/density.py
View file

@ -22,7 +22,6 @@ import numpy
from numba import jit from numba import jit
from tqdm import trange from tqdm import trange
from ..read.utils import real2redshift
from .interp import divide_nonzero from .interp import divide_nonzero
from .utils import force_single_precision from .utils import force_single_precision
@ -32,18 +31,6 @@ class BaseField(ABC):
_box = None _box = None
_MAS = None _MAS = None
@property
def boxsize(self):
"""
Box size. Particle positions are always assumed to be in box units,
therefore this is 1.
Returns
-------
boxsize : float
"""
return 1.
@property @property
def box(self): def box(self):
""" """
@ -51,7 +38,7 @@ class BaseField(ABC):
Returns Returns
------- -------
box : :py:class:`csiborgtools.units.CSiBORGBox` :py:class:`csiborgtools.units.CSiBORGBox`
""" """
return self._box return self._box
@ -70,10 +57,10 @@ class BaseField(ABC):
Returns Returns
------- -------
MAS : str str
""" """
if self._MAS is None: if self._MAS is None:
raise ValueError("`mas` is not set.") raise ValueError("`MAS` is not set.")
return self._MAS return self._MAS
@MAS.setter @MAS.setter
@ -99,6 +86,8 @@ class DensityField(BaseField):
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS' point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline). (piecewise cubic spline).
paths : :py:class:`csiborgtools.read.Paths`
The simulation paths.
References References
---------- ----------
@ -128,13 +117,12 @@ class DensityField(BaseField):
delta -= 1 delta -= 1
return delta return delta
def __call__(self, parts, grid, in_rsp, flip_xz=True, nbatch=30, def __call__(self, parts, grid, flip_xz=True, nbatch=30, verbose=True):
verbose=True):
""" """
Calculate the density field using a Pylians routine [1, 2]. Calculate the density field using a Pylians routine [1, 2].
Iteratively loads the particles into memory, flips their `x` and `z` Iteratively loads the particles into memory, flips their `x` and `z`
coordinates. Particles are assumed to be in box units, with positions coordinates. Particles are assumed to be in box units, with positions
in [0, 1] and observer in the centre of the box if RSP is applied. in [0, 1]
Parameters Parameters
---------- ----------
@ -143,8 +131,6 @@ class DensityField(BaseField):
Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`. Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`.
grid : int grid : int
Grid size. Grid size.
in_rsp : bool
Whether to calculate the density field in redshift space.
flip_xz : bool, optional flip_xz : bool, optional
Whether to flip the `x` and `z` coordinates. Whether to flip the `x` and `z` coordinates.
nbatch : int, optional nbatch : int, optional
@ -173,21 +159,14 @@ class DensityField(BaseField):
pos = parts[start:end] pos = parts[start:end]
pos, vel, mass = pos[:, :3], pos[:, 3:6], pos[:, 6] pos, vel, mass = pos[:, :3], pos[:, 3:6], pos[:, 6]
pos = force_single_precision(pos, "particle_position") pos = force_single_precision(pos)
vel = force_single_precision(vel, "particle_velocity") vel = force_single_precision(vel)
mass = force_single_precision(mass, "particle_mass") mass = force_single_precision(mass)
if flip_xz: if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]] pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]] vel[:, [0, 2]] = vel[:, [2, 0]]
if in_rsp: MASL.MA(pos, rho, 1., self.MAS, W=mass, verbose=False)
raise NotImplementedError("Redshift space needs to be fixed.")
# TODO change how called + units.
pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], self.box,
in_box_units=True, periodic_wrap=True,
make_copy=False)
MASL.MA(pos, rho, self.boxsize, self.MAS, W=mass, verbose=False)
if end == nparts: if end == nparts:
break break
start = end start = end
@ -223,7 +202,7 @@ class VelocityField(BaseField):
@staticmethod @staticmethod
@jit(nopython=True) @jit(nopython=True)
def radial_velocity(rho_vel): def radial_velocity(rho_vel, observer_velocity):
""" """
Calculate the radial velocity field around the observer in the centre Calculate the radial velocity field around the observer in the centre
of the box. of the box.
@ -232,6 +211,8 @@ class VelocityField(BaseField):
---------- ----------
rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`. rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
Velocity field along each axis. Velocity field along each axis.
observer_velocity : 3-dimensional array of shape `(3,)`
Observer velocity.
Returns Returns
------- -------
@ -240,13 +221,19 @@ class VelocityField(BaseField):
""" """
grid = rho_vel.shape[1] grid = rho_vel.shape[1]
radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32) radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
vx0, vy0, vz0 = observer_velocity
for i in range(grid): for i in range(grid):
px = i - 0.5 * (grid - 1) px = i - 0.5 * (grid - 1)
for j in range(grid): for j in range(grid):
py = j - 0.5 * (grid - 1) py = j - 0.5 * (grid - 1)
for k in range(grid): for k in range(grid):
pz = k - 0.5 * (grid - 1) pz = k - 0.5 * (grid - 1)
vx, vy, vz = rho_vel[:, i, j, k]
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) radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
/ numpy.sqrt(px**2 + py**2 + pz**2)) / numpy.sqrt(px**2 + py**2 + pz**2))
return radvel return radvel
@ -297,19 +284,19 @@ class VelocityField(BaseField):
pos = parts[start:end] pos = parts[start:end]
pos, vel, mass = pos[:, :3], pos[:, 3:6], pos[:, 6] pos, vel, mass = pos[:, :3], pos[:, 3:6], pos[:, 6]
pos = force_single_precision(pos, "particle_position") pos = force_single_precision(pos)
vel = force_single_precision(vel, "particle_velocity") vel = force_single_precision(vel)
mass = force_single_precision(mass, "particle_mass") mass = force_single_precision(mass)
if flip_xz: if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]] pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]] vel[:, [0, 2]] = vel[:, [2, 0]]
vel *= mass.reshape(-1, 1) vel *= mass.reshape(-1, 1)
for i in range(3): for i in range(3):
MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[:, i], MASL.MA(pos, rho_vel[i], 1., self.MAS, W=vel[:, i],
verbose=False) verbose=False)
MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass, MASL.MA(pos, cellcounts, 1., self.MAS, W=mass,
verbose=False) verbose=False)
if end == nparts: if end == nparts:
break break

244
csiborgtools/field/interp.py
View file

@ -20,23 +20,22 @@ import numpy
from numba import jit from numba import jit
from tqdm import trange from tqdm import trange
from ..read.utils import radec_to_cartesian, real2redshift
from .utils import force_single_precision from .utils import force_single_precision
from ..utils import periodic_wrap_grid, radec_to_cartesian
def evaluate_cartesian(*fields, pos, interp="CIC"): def evaluate_cartesian(*fields, pos, interp="CIC"):
""" """
Evaluate a scalar field at Cartesian coordinates. Evaluate a scalar field(s) at Cartesian coordinates `pos`.
Parameters Parameters
---------- ----------
field : (list of) 3-dimensional array of shape `(grid, grid, grid)` field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
Fields to be interpolated. Fields to be interpolated.
pos : 2-dimensional array of shape `(n_samples, 3)` pos : 2-dimensional array of shape `(n_samples, 3)`
Positions to evaluate the density field. Assumed to be in box Query positions in box units.
units.
interp : str, optional interp : str, optional
Interpolation method. Can be either `CIC` or `NGP`. Interpolation method, `NGP` or `CIC`.
Returns Returns
------- -------
@ -44,7 +43,7 @@ def evaluate_cartesian(*fields, pos, interp="CIC"):
""" """
assert interp in ["CIC", "NGP"] assert interp in ["CIC", "NGP"]
boxsize = 1. boxsize = 1.
pos = force_single_precision(pos, "pos") pos = force_single_precision(pos)
nsamples = pos.shape[0] nsamples = pos.shape[0]
interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32) interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32)
@ -64,60 +63,71 @@ def evaluate_cartesian(*fields, pos, interp="CIC"):
return interp_fields return interp_fields
def evaluate_sky(*fields, pos, box, isdeg=True): def evaluate_sky(*fields, pos, box):
""" """
Evaluate the scalar fields at given distance, right ascension and Evaluate a scalar field(s) at radial distance `Mpc / h`, right ascensions
declination. Assumes an observed in the centre of the box, with [0, 360) deg and declinations [-90, 90] deg.
distance being in :math:`Mpc`. Uses CIC interpolation.
Parameters Parameters
---------- ----------
fields : (list of) 3-dimensional array of shape `(grid, grid, grid)` fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
Field to be interpolated. Field to be interpolated.
pos : 2-dimensional array of shape `(n_samples, 3)` pos : 2-dimensional array of shape `(n_samples, 3)`
Spherical coordinates to evaluate the field. Columns are distance, Query spherical coordinates.
right ascension, declination, respectively.
box : :py:class:`csiborgtools.read.CSiBORGBox` box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations. The simulation box information and transformations.
isdeg : bool, optional
Whether `ra` and `dec` are in degres. By default `True`.
Returns Returns
------- -------
interp_fields : (list of) 1-dimensional array of shape `(n_samples,). interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
""" """
pos = force_single_precision(pos, "pos") pos = force_single_precision(pos)
# We first calculate convert the distance to box coordinates and then
# convert to Cartesian coordinates.
pos[:, 0] = box.mpc2box(pos[:, 0]) pos[:, 0] = box.mpc2box(pos[:, 0])
X = radec_to_cartesian(pos, isdeg) cart_pos = radec_to_cartesian(pos) + 0.5
# Then we move the origin to match the box coordinates
X += 0.5 return evaluate_cartesian(*fields, pos=cart_pos)
return evaluate_cartesian(*fields, pos=X)
def observer_vobs(velocity_field):
"""
Calculate the observer velocity from a velocity field. Assumes an observer
in the centre of the box.
Parameters
----------
velocity_field : 4-dimensional array of shape `(3, grid, grid, grid)`
Returns
-------
1-dimensional array of shape `(3,)`
"""
pos = numpy.asanyarray([0.5, 0.5, 0.5]).reshape(1, 3)
vobs = numpy.full(3, numpy.nan, dtype=numpy.float32)
for i in range(3):
vobs[i] = evaluate_cartesian(velocity_field[i, ...], pos=pos)[0]
return vobs
def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True): def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True):
r""" r"""
Make a sky map of a scalar field. The observer is in the centre of the Make a sky map of a scalar field. The observer is in the centre of the
box the field is evaluated along directions `angpos`. Along each box the field is evaluated along directions `angpos` (RA [0, 360) deg,
direction, the field is evaluated distances `dist_marg` and summed. dec [-90, 90] deg). Along each direction, the field is evaluated distances
Uses CIC interpolation. `dist` (`Mpc / h`) and summed. Uses CIC interpolation.
Parameters Parameters
---------- ----------
field : 3-dimensional array of shape `(grid, grid, grid)` field : 3-dimensional array of shape `(grid, grid, grid)`
Field to be interpolated Field to be interpolated
angpos : 2-dimensional arrays of shape `(ndir, 2)` angpos : 2-dimensional arrays of shape `(ndir, 2)`
Directions to evaluate the field. Assumed to be RA Directions to evaluate the field.
:math:`\in [0, 360]` and dec :math:`\in [-90, 90]` degrees,
respectively.
dist : 1-dimensional array dist : 1-dimensional array
Uniformly spaced radial distances to evaluate the field. Uniformly spaced radial distances to evaluate the field.
box : :py:class:`csiborgtools.read.CSiBORGBox` box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations. The simulation box information and transformations.
volume_weight : bool, optional volume_weight : bool, optional
Whether to weight the field by the volume of the pixel, i.e. a Whether to weight the field by the volume of the pixel.
:math:`r^2` correction.
verbose : bool, optional verbose : bool, optional
Verbosity flag. Verbosity flag.
@ -152,22 +162,9 @@ def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True):
@jit(nopython=True) @jit(nopython=True)
def divide_nonzero(field0, field1): def divide_nonzero(field0, field1):
""" """
Divide two fields where the second one is not zero. If the second field Perform in-place `field0 /= field1` but only where `field1 != 0`.
is zero, the first one is left unchanged. Operates in-place.
Parameters
----------
field0 : 3-dimensional array of shape `(grid, grid, grid)`
Field to be divided.
field1 : 3-dimensional array of shape `(grid, grid, grid)`
Field to divide by.
Returns
-------
field0 : 3-dimensional array of shape `(grid, grid, grid)`
Field divided by the second one.
""" """
assert field0.shape == field1.shape assert field0.shape == field1.shape, "Field shapes must match."
imax, jmax, kmax = field0.shape imax, jmax, kmax = field0.shape
for i in range(imax): for i in range(imax):
@ -177,134 +174,89 @@ def divide_nonzero(field0, field1):
field0[i, j, k] /= field1[i, j, k] field0[i, j, k] /= field1[i, j, k]
def observer_vobs(velocity_field): @jit(nopython=True)
def make_gridpos(grid_size):
"""Make a regular grid of positions and distances from the center."""
grid_pos = numpy.full((grid_size**3, 3), numpy.nan, dtype=numpy.float32)
grid_dist = numpy.full(grid_size**3, numpy.nan, dtype=numpy.float32)
n = 0
for i in range(grid_size):
px = (i - 0.5 * (grid_size - 1)) / grid_size
px2 = px**2
for j in range(grid_size):
py = (j - 0.5 * (grid_size - 1)) / grid_size
py2 = py**2
for k in range(grid_size):
pz = (k - 0.5 * (grid_size - 1)) / grid_size
pz2 = pz**2
grid_pos[n, 0] = px
grid_pos[n, 1] = py
grid_pos[n, 2] = pz
grid_dist[n] = (px2 + py2 + pz2)**0.5
n += 1
return grid_pos, grid_dist
def field2rsp(field, radvel_field, box, MAS, init_value=0.):
""" """
Calculate the observer velocity from a velocity field. Assumes the observer Forward model a real space scalar field to redshift space.
is in the centre of the box.
Parameters Parameters
---------- ----------
velocity_field : 4-dimensional array of shape `(3, grid, grid, grid)` field : 3-dimensional array of shape `(grid, grid, grid)`
Velocity field to calculate the observer velocity from. Real space field to be evolved to redshift space.
radvel_field : 3-dimensional array of shape `(grid, grid, grid)`
Returns Radial velocity field in `km / s`. Expected to account for the observer
------- velocity.
vobs : 1-dimensional array of shape `(3,)`
Observer velocity in units of `velocity_field`.
"""
pos = numpy.asanyarray([0.5, 0.5, 0.5]).reshape(1, 3)
vobs = numpy.full(3, numpy.nan, dtype=numpy.float32)
for i in range(3):
vobs[i] = evaluate_cartesian(velocity_field[i, ...], pos=pos)[0]
return vobs
def field2rsp(*fields, parts, vobs, box, nbatch=30, flip_partsxz=True,
init_value=0., verbose=True):
"""
Forward model real space scalar fields to redshift space. Attaches field
values to particles, those are then moved to redshift space and from their
positions reconstructs back the field on a grid by NGP interpolation.
Parameters
----------
fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
Real space fields to be evolved to redshift space.
parts : 2-dimensional array of shape `(n_parts, 6)`
Particle positions and velocities in real space. Must be organised as
`x, y, z, vx, vy, vz`.
vobs : 1-dimensional array of shape `(3,)`
Observer velocity in units matching `parts`.
box : :py:class:`csiborgtools.read.CSiBORGBox` box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations. The simulation box information and transformations.
nbatch : int, optional MAS : str
Number of batches to use when moving particles to redshift space. Mass assignment. Must be one of `NGP`, `CIC`, `TSC` or `PCS`.
Particles are assumed to be lazily loaded to memory.
flip_partsxz : bool, optional
Whether to flip the x and z coordinates of the particles. This is
because of a RAMSES bug.
init_value : float, optional init_value : float, optional
Initial value of the RSP field on the grid. Initial value of the RSP field on the grid.
verbose : bool, optional
Verbosity flag.
Returns Returns
------- -------
rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)` 3-dimensional array of shape `(grid, grid, grid)`
""" """
raise NotImplementedError("Figure out what to do with Vobs.") grid = field.shape[0]
nfields = len(fields) H0_inv = 1. / 100 / box.box2mpc(1.)
# Check that all fields have the same shape.
if nfields > 1:
assert all(fields[0].shape == fields[i].shape
for i in range(1, nfields))
rsp_fields = [numpy.full(field.shape, init_value, dtype=numpy.float32) # Calculate the regular grid positions and distances from the center.
for field in fields] grid_pos, grid_dist = make_gridpos(grid)
cellcounts = numpy.zeros(rsp_fields[0].shape, dtype=numpy.float32) grid_dist = grid_dist.reshape(-1, 1)
nparts = parts.shape[0] # Move the grid positions to redshift space.
batch_size = nparts // nbatch grid_pos *= (1 + H0_inv * radvel_field.reshape(-1, 1) / grid_dist)
start = 0 grid_pos += 0.5
for __ in trange(nbatch + 1) if verbose else range(nbatch + 1): grid_pos = periodic_wrap_grid(grid_pos)
# We first load the batch of particles into memory and flip x and z.
end = min(start + batch_size, nparts)
pos = parts[start:end]
pos, vel = pos[:, :3], pos[:, 3:6]
if flip_partsxz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
# Then move the particles to redshift space.
# TODO here the function is now called differently and pos assumes
# different units.
rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
in_box_units=True, periodic_wrap=True,
make_copy=True)
# ... and count the number of particles in each grid cell.
MASL.MA(rsp_pos, cellcounts, 1., "NGP")
# Now finally we evaluate the field at the particle positions in real rsp_field = numpy.full(field.shape, init_value, dtype=numpy.float32)
# space and then assign the values to the grid in redshift space. cell_counts = numpy.zeros(rsp_field.shape, dtype=numpy.float32)
for i in range(nfields):
values = evaluate_cartesian(fields[i], pos=pos)
MASL.MA(rsp_pos, rsp_fields[i], 1., "NGP", W=values)
if end == nparts:
break
start = end
# We divide by the number of particles in each cell. # Interpolate the field to the grid positions.
for i in range(len(fields)): MASL.MA(grid_pos, rsp_field, 1., MAS, W=field.reshape(-1,))
divide_nonzero(rsp_fields[i], cellcounts) MASL.MA(grid_pos, cell_counts, 1., MAS)
divide_nonzero(rsp_field, cell_counts)
if len(fields) == 1: return rsp_field
return rsp_fields[0]
return rsp_fields
@jit(nopython=True) @jit(nopython=True)
def fill_outside(field, fill_value, rmax, boxsize): def fill_outside(field, fill_value, rmax, boxsize):
""" """
Fill cells outside of a sphere of radius `rmax` with `fill_value`. Centered Fill cells outside of a sphere of radius `rmax` around the box centre with
in the middle of the box. `fill_value`.
Parameters
----------
field : 3-dimensional array of shape `(grid, grid, grid)`
Field to be filled.
fill_value : float
Value to fill the field with.
rmax : float
Radius outside of which to fill the field..
boxsize : float
Size of the box.
Returns
-------
field : 3-dimensional array of shape `(grid, grid, grid)`
""" """
imax, jmax, kmax = field.shape imax, jmax, kmax = field.shape
assert imax == jmax == kmax assert imax == jmax == kmax
N = imax N = imax
# Squared radial distance from the center of the box in box units. # Squared radial distance from the center of the box in box units.
rmax_box2 = (N * rmax / boxsize)**2 rmax_box2 = (N * rmax / boxsize)**2

49
csiborgtools/field/utils.py
View file

@ -15,31 +15,16 @@
""" """
Utility functions for the field module. Utility functions for the field module.
""" """
from warnings import warn
import healpy import healpy
import numpy import numpy
import smoothing_library as SL import smoothing_library as SL
def force_single_precision(x, name): def force_single_precision(x):
""" """
Convert `x` to float32 if it is not already. Attempt to convert an array `x` to float 32.
Parameters
----------
x : array
Array to convert.
name : str
Name of the array.
Returns
-------
x : array
Converted array.
""" """
if x.dtype != numpy.float32: if x.dtype != numpy.float32:
warn(f"Converting `{name}` to float32.", UserWarning, stacklevel=1)
x = x.astype(numpy.float32) x = x.astype(numpy.float32)
return x return x
@ -47,21 +32,6 @@ def force_single_precision(x, name):
def smoothen_field(field, smooth_scale, boxsize, threads=1): def smoothen_field(field, smooth_scale, boxsize, threads=1):
""" """
Smooth a field with a Gaussian filter. Smooth a field with a Gaussian filter.
Parameters
----------
field : 3-dimensional array of shape `(grid, grid, grid)`
Field to be smoothed.
smooth_scale : float, optional
Gaussian kernal scale to smoothen the density field, in box units.
boxsize : float
Size of the box.
threads : int, optional
Number of threads. By default 1.
Returns
-------
smoothed_field : 3-dimensional array of shape `(grid, grid, grid)`
""" """
W_k = SL.FT_filter(boxsize, smooth_scale, field.shape[0], "Gaussian", W_k = SL.FT_filter(boxsize, smooth_scale, field.shape[0], "Gaussian",
threads) threads)
@ -70,19 +40,10 @@ def smoothen_field(field, smooth_scale, boxsize, threads=1):
def nside2radec(nside): def nside2radec(nside):
""" """
Generate RA and declination for HEALPix pixel centres. Generate RA [0, 360] deg. and declination [-90, 90] deg. for HEALPix pixel
centres at a given nside.
Parameters
----------
nside : int
HEALPix nside parameter.
Returns
-------
radec : array of shape `(npix, 2)`
RA and declination in degrees.
""" """
pixs = numpy.arange(healpy.nside2npix(nside)) pixs = numpy.arange(healpy.nside2npix(nside))
theta, phi = healpy.pix2ang(nside, pixs) theta, phi = healpy.pix2ang(nside, pixs)
theta -= numpy.pi / 2 theta -= numpy.pi / 2
return numpy.rad2deg(numpy.vstack([phi, theta]).T) return 180 / numpy.pi * numpy.vstack([phi, theta]).T

43
csiborgtools/match/match.py
View file

@ -151,7 +151,7 @@ class RealisationsMatcher(BaseMatcher):
Returns Returns
------- -------
dlogmass : float float
""" """
return self._dlogmass return self._dlogmass
@ -169,7 +169,7 @@ class RealisationsMatcher(BaseMatcher):
Returns Returns
------- -------
mass_kind : str str
""" """
return self._mass_kind return self._mass_kind
@ -186,7 +186,7 @@ class RealisationsMatcher(BaseMatcher):
Returns Returns
------- -------
overlapper : :py:class:`csiborgtools.match.ParticleOverlap` :py:class:`csiborgtools.match.ParticleOverlap`
""" """
return self._overlapper return self._overlapper
@ -661,20 +661,7 @@ class ParticleOverlap(BaseMatcher):
def pos2cell(pos, ncells): def pos2cell(pos, ncells):
""" """
Convert position to cell number. If `pos` is in Convert position to cell number if there are `ncells` cells along the axis.
`numpy.typecodes["AllInteger"]` assumes it to already be the cell
number.
Parameters
----------
pos : 1-dimensional array
Array of positions along an axis in the box.
ncells : int
Number of cells along the axis.
Returns
-------
cells : 1-dimensional array
""" """
if pos.dtype.char in numpy.typecodes["AllInteger"]: if pos.dtype.char in numpy.typecodes["AllInteger"]:
return pos return pos
@ -684,17 +671,7 @@ def pos2cell(pos, ncells):
def read_nshift(smooth_kwargs): def read_nshift(smooth_kwargs):
""" """
Determine the number of cells to pad the density field if smoothing is Determine the number of cells to pad the density field if smoothing is
applied. It defaults to the ceiling of three times the smoothing scale. applied. Defaults to the ceiling of three times the smoothing scale.
Parameters
----------
smooth_kwargs : dict or None
Arguments to be passed to :py:func:`scipy.ndimage.gaussian_filter`.
If `None`, no smoothing is applied.
Returns
-------
nshift : int
""" """
return 0 if smooth_kwargs is None else ceil(3 * smooth_kwargs["sigma"]) return 0 if smooth_kwargs is None else ceil(3 * smooth_kwargs["sigma"])
@ -702,8 +679,7 @@ def read_nshift(smooth_kwargs):
@jit(nopython=True) @jit(nopython=True)
def fill_delta(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights): def fill_delta(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights):
""" """
Fill array `delta` by adding `weights` to the specified cells. This is a Fill array `delta` by adding `weights` to the specified cells.
JIT implementation.
Parameters Parameters
---------- ----------
@ -730,8 +706,8 @@ def fill_delta(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights):
@jit(nopython=True) @jit(nopython=True)
def fill_delta_indxs(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights): def fill_delta_indxs(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights):
""" """
Fill array `delta` by adding `weights` to the specified cells and return Fill array `delta` by adding `weights` to the specified cells. Returns
indices where `delta` was assigned a value. This is a JIT implementation. the indices where `delta` was assigned a value.
Parameters Parameters
---------- ----------
@ -1161,6 +1137,9 @@ def matching_max(cat0, catx, mass_kind, mult, periodic, overlap=None,
out[i]["hid0"] = hid0 out[i]["hid0"] = hid0
out[i]["mass0"] = 10**mass0[i] out[i]["mass0"] = 10**mass0[i]
if not numpy.isfinite(rad0[i]):
continue
neigh_dists, neigh_inds = knnx.radius_neighbors(pos0[i].reshape(1, -1), neigh_dists, neigh_inds = knnx.radius_neighbors(pos0[i].reshape(1, -1),
mult * rad0[i]) mult * rad0[i])
neigh_dists, neigh_inds = neigh_dists[0], neigh_inds[0] neigh_dists, neigh_inds = neigh_dists[0], neigh_inds[0]

3
csiborgtools/read/__init__.py
View file

@ -26,5 +26,4 @@ from .pk_summary import PKReader # noqa
from .readsim import (MmainReader, CSiBORGReader, QuijoteReader, halfwidth_mask, # noqa from .readsim import (MmainReader, CSiBORGReader, QuijoteReader, halfwidth_mask, # noqa
load_halo_particles) # noqa load_halo_particles) # noqa
from .tpcf_summary import TPCFReader # noqa from .tpcf_summary import TPCFReader # noqa
from .utils import (cartesian_to_radec, cols_to_structured, radec_to_cartesian, # noqa from .utils import (cols_to_structured, read_h5) # noqa
read_h5, real2redshift) # noqa

20
csiborgtools/read/halo_cat.py
View file

@ -31,9 +31,9 @@ from sklearn.neighbors import NearestNeighbors
from .box_units import CSiBORGBox, QuijoteBox from .box_units import CSiBORGBox, QuijoteBox
from .paths import Paths from .paths import Paths
from .readsim import CSiBORGReader from .readsim import CSiBORGReader
from .utils import (add_columns, cartesian_to_radec, cols_to_structured, from .utils import add_columns, cols_to_structured, flip_cols
flip_cols, radec_to_cartesian, real2redshift) from ..utils import (periodic_distance_two_points, real2redshift,
from ..utils import periodic_distance_two_points cartesian_to_radec, radec_to_cartesian)
class BaseCatalogue(ABC): class BaseCatalogue(ABC):
@ -631,6 +631,9 @@ class QuijoteHaloCatalogue(BaseCatalogue):
Load initial positions from `fit_init.py`. Load initial positions from `fit_init.py`.
with_lagpatch : bool, optional with_lagpatch : bool, optional
Load halos with a resolved Lagrangian patch. Load halos with a resolved Lagrangian patch.
load_backup : bool, optional
Load halos from the backup catalogue that do not have corresponding
snapshots.
""" """
_nsnap = None _nsnap = None
_origin = None _origin = None
@ -638,14 +641,15 @@ class QuijoteHaloCatalogue(BaseCatalogue):
def __init__(self, nsim, paths, nsnap, def __init__(self, nsim, paths, nsnap,
observer_location=[500., 500., 500.], observer_location=[500., 500., 500.],
bounds=None, load_fitted=True, load_initial=True, bounds=None, load_fitted=True, load_initial=True,
with_lagpatch=False): with_lagpatch=False, load_backup=False):
self.nsim = nsim self.nsim = nsim
self.paths = paths self.paths = paths
self.nsnap = nsnap self.nsnap = nsnap
self.observer_location = observer_location self.observer_location = observer_location
self._box = QuijoteBox(nsnap, nsim, paths) # NOTE watch out about here setting nsim = 0
self._box = QuijoteBox(nsnap, 0, paths)
fpath = self.paths.fof_cat(nsim, "quijote") fpath = self.paths.fof_cat(nsim, "quijote", load_backup)
fof = FoF_catalog(fpath, self.nsnap, long_ids=False, swap=False, fof = FoF_catalog(fpath, self.nsnap, long_ids=False, swap=False,
SFR=False, read_IDs=False) SFR=False, read_IDs=False)
@ -667,6 +671,10 @@ class QuijoteHaloCatalogue(BaseCatalogue):
# particles unassigned to any FoF group. # particles unassigned to any FoF group.
data["index"] = 1 + numpy.arange(data.size, dtype=numpy.int32) data["index"] = 1 + numpy.arange(data.size, dtype=numpy.int32)
if load_backup and (load_initial or load_fitted):
raise ValueError("Cannot load initial/fitted data for the backup "
"catalogues.")
if load_initial: if load_initial:
data = self.load_initial(data, paths, "quijote") data = self.load_initial(data, paths, "quijote")
if load_fitted: if load_fitted:

63
csiborgtools/read/paths.py
View file

@ -214,7 +214,7 @@ class Paths:
fout = fout.replace(".npy", "_sorted.npy") fout = fout.replace(".npy", "_sorted.npy")
return fout return fout
def fof_cat(self, nsim, simname): def fof_cat(self, nsim, simname, from_quijote_backup=False):
r""" r"""
Path to the :math:`z = 0` FoF halo catalogue. Path to the :math:`z = 0` FoF halo catalogue.
@ -224,12 +224,23 @@ class Paths:
IC realisation index. IC realisation index.
simname : str simname : str
Simulation name. Must be one of `csiborg` or `quijote`. Simulation name. Must be one of `csiborg` or `quijote`.
from_quijote_backup : bool, optional
Whether to return the path to the Quijote FoF catalogue from the
backup.
Returns
-------
path : str
""" """
if simname == "csiborg": if simname == "csiborg":
fdir = join(self.postdir, "FoF_membership", ) fdir = join(self.postdir, "FoF_membership", )
try_create_directory(fdir) try_create_directory(fdir)
return join(fdir, f"halo_catalog_{nsim}_FOF.txt") return join(fdir, f"halo_catalog_{nsim}_FOF.txt")
elif simname == "quijote": elif simname == "quijote":
if from_quijote_backup:
return join(self.quijote_dir, "halos_backup", str(nsim))
else:
return join(self.quijote_dir, "Halos_fiducial", str(nsim)) return join(self.quijote_dir, "Halos_fiducial", str(nsim))
else: else:
raise ValueError(f"Unknown simulation name `{simname}`.") raise ValueError(f"Unknown simulation name `{simname}`.")
@ -285,7 +296,7 @@ class Paths:
try_create_directory(fdir) try_create_directory(fdir)
return join(fdir, f"{kind}_{str(nsim).zfill(5)}.{ftype}") return join(fdir, f"{kind}_{str(nsim).zfill(5)}.{ftype}")
def get_ics(self, simname): def get_ics(self, simname, from_quijote_backup=False):
""" """
Get available IC realisation IDs for either the CSiBORG or Quijote Get available IC realisation IDs for either the CSiBORG or Quijote
simulation suite. simulation suite.
@ -294,6 +305,8 @@ class Paths:
---------- ----------
simname : str simname : str
Simulation name. Must be `csiborg` or `quijote`. Simulation name. Must be `csiborg` or `quijote`.
from_quijote_backup : bool, optional
Whether to return the ICs from the Quijote backup.
Returns Returns
------- -------
@ -311,10 +324,13 @@ class Paths:
except ValueError: except ValueError:
pass pass
elif simname == "quijote" or simname == "quijote_full": elif simname == "quijote" or simname == "quijote_full":
files = glob( if from_quijote_backup:
"/mnt/extraspace/rstiskalek/Quijote/Snapshots_fiducial/*") files = glob(join(self.quijote_dir, "halos_backup", "*"))
else:
warn(("Taking only the snapshots that also have "
"a FoF catalogue!"))
files = glob(join(self.quijote_dir, "Snapshots_fiducial", "*"))
files = [int(f.split("/")[-1]) for f in files] files = [int(f.split("/")[-1]) for f in files]
warn("Taking only the snapshots that also have a FoF catalogue!")
files = [f for f in files if f < 100] files = [f for f in files if f < 100]
else: else:
raise ValueError(f"Unknown simulation name `{simname}`.") raise ValueError(f"Unknown simulation name `{simname}`.")
@ -545,7 +561,7 @@ class Paths:
return join(fdir, fname) return join(fdir, fname)
def field(self, kind, MAS, grid, nsim, in_rsp, smooth_scale=None): def field(self, kind, MAS, grid, nsim, in_rsp):
r""" r"""
Path to the files containing the calculated density fields in CSiBORG. Path to the files containing the calculated density fields in CSiBORG.
@ -562,8 +578,6 @@ class Paths:
IC realisation index. IC realisation index.
in_rsp : bool in_rsp : bool
Whether the calculation is performed in redshift space. Whether the calculation is performed in redshift space.
smooth_scale : float, optional
Smoothing scale in :math:`\mathrm{Mpc}/h`
Returns Returns
------- -------
@ -579,12 +593,32 @@ class Paths:
kind = kind + "_rsp" kind = kind + "_rsp"
fname = f"{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npy" fname = f"{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npy"
if smooth_scale is not None and smooth_scale > 0:
smooth_scale = float(smooth_scale)
fname = fname.replace(".npy", f"smooth{smooth_scale}.npy")
return join(fdir, fname) return join(fdir, fname)
def halo_counts(self, simname, nsim): def observer_peculiar_velocity(self, MAS, grid, nsim):
"""
Path to the files containing the observer peculiar velocity.
Parameters
----------
MAS : str
Mass-assignment scheme.
grid : int
Grid size.
nsim : int
IC realisation index.
Returns
-------
path : str
"""
fdir = join(self.postdir, "environment")
try_create_directory(fdir)
fname = f"obs_vp_{MAS}_{str(nsim).zfill(5)}_{grid}.npz"
return join(fdir, fname)
def halo_counts(self, simname, nsim, from_quijote_backup=False):
""" """
Path to the files containing the binned halo counts. Path to the files containing the binned halo counts.
@ -594,6 +628,9 @@ class Paths:
Simulation name. Must be `csiborg`, `quijote` or `quijote_full`. Simulation name. Must be `csiborg`, `quijote` or `quijote_full`.
nsim : int nsim : int
IC realisation index. IC realisation index.
from_quijote_backup : bool, optional
Whether to return the path to the Quijote halo counts from the
backup catalogues.
Returns Returns
------- -------
@ -602,6 +639,8 @@ class Paths:
fdir = join(self.postdir, "HMF") fdir = join(self.postdir, "HMF")
try_create_directory(fdir) try_create_directory(fdir)
fname = f"halo_counts_{simname}_{str(nsim).zfill(5)}.npz" fname = f"halo_counts_{simname}_{str(nsim).zfill(5)}.npz"
if from_quijote_backup:
fname = fname.replace("halo_counts", "halo_counts_backup")
return join(fdir, fname) return join(fdir, fname)
def cross_nearest(self, simname, run, kind, nsim=None, nobs=None): def cross_nearest(self, simname, run, kind, nsim=None, nobs=None):

164
csiborgtools/read/utils.py
View file

@ -12,126 +12,11 @@
# You should have received a copy of the GNU General Public License along # You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Various coordinate transformations.
"""
from os.path import isfile from os.path import isfile
import numpy import numpy
from h5py import File from h5py import File
###############################################################################
# Coordinate transforms #
###############################################################################
def cartesian_to_radec(X, indeg=True):
"""
Calculate the radial distance, RA, dec from Cartesian coordinates. Note,
RA is in range [0, 360) degrees and dec in range [-90, 90] degrees.
Parameters
----------
X : 2-dimensional array `(nsamples, 3)`
Cartesian coordinates.
indeg : bool, optional
Whether to return RA and DEC in degrees.
Returns
-------
out : 2-dimensional array `(nsamples, 3)`
Radial distance, RA and dec.
"""
x, y, z = X[:, 0], X[:, 1], X[:, 2]
dist = numpy.linalg.norm(X, axis=1)
dec = numpy.arcsin(z / dist)
ra = numpy.arctan2(y, x)
ra[ra < 0] += 2 * numpy.pi # Wrap RA to [0, 2pi)
if indeg:
ra = numpy.rad2deg(ra)
dec = numpy.rad2deg(dec)
return numpy.vstack([dist, ra, dec]).T
def radec_to_cartesian(X, isdeg=True):
"""
Calculate Cartesian coordinates from radial distance, RA, dec. Note, RA is
expected in range [0, 360) degrees and dec in range [-90, 90] degrees.
Parameters
----------
X : 2-dimensional array `(nsamples, 3)`
Radial distance, RA and dec.
isdeg : bool, optional
Whether to return RA and DEC in degrees.
Returns
-------
out : 2-dimensional array `(nsamples, 3)`
Cartesian coordinates.
"""
dist, ra, dec = X[:, 0], X[:, 1], X[:, 2]
if isdeg:
ra = numpy.deg2rad(ra)
dec = numpy.deg2rad(dec)
x = dist * numpy.cos(dec) * numpy.cos(ra)
y = dist * numpy.cos(dec) * numpy.sin(ra)
z = dist * numpy.sin(dec)
return numpy.vstack([x, y, z]).T
def real2redshift(pos, vel, observer_location, box, periodic_wrap=True,
subtract_observer=False, make_copy=True):
r"""
Convert real-space position to redshift space position.
Parameters
----------
pos : 2-dimensional array `(nsamples, 3)`
Real-space Cartesian components in :math:`\mathrm{cMpc} / h`.
vel : 2-dimensional array `(nsamples, 3)`
Cartesian velocity in :math:`\mathrm{km} \mathrm{s}^{-1}`.
observer_location: 1-dimensional array `(3,)`
Observer location in :math:`\mathrm{cMpc} / h`.
box : py:class:`csiborg.read.CSiBORGBox`
Box units.
periodic_wrap : bool, optional
Whether to wrap around the box, particles may be outside the default
bounds once RSD is applied.
subtract_observer : bool, optional
If True, subtract the observer's location from the returned
positions.
make_copy : bool, optional
Whether to make a copy of `pos` before modifying it.
Returns
-------
pos : 2-dimensional array `(nsamples, 3)`
Redshift-space Cartesian position in :math:`\mathrm{cMpc} / h`.
"""
if make_copy:
pos = numpy.copy(pos)
# Place the observer at the origin
pos -= observer_location
# Dot product of position vector and velocity
vr_dot = numpy.sum(pos * vel, axis=1)
# Compute the norm squared of the displacement
norm2 = numpy.sum(pos**2, axis=1)
pos *= (1 + box._aexp / box.H0 * vr_dot / norm2).reshape(-1, 1)
# Place the observer back at the original location
if not subtract_observer:
pos += observer_location
if periodic_wrap:
boxsize = box.box2mpc(1.)
# Wrap around the box.
pos = numpy.where(pos > boxsize, pos - boxsize, pos)
pos = numpy.where(pos < 0, pos + boxsize, pos)
return pos
############################################################################### ###############################################################################
# Array manipulation # # Array manipulation #
@ -140,19 +25,7 @@ def real2redshift(pos, vel, observer_location, box, periodic_wrap=True,
def cols_to_structured(N, cols): def cols_to_structured(N, cols):
""" """
Allocate a structured array from `cols`. Allocate a structured array from `cols`, a list of (name, dtype) tuples.
Parameters
----------
N : int
Structured array size.
cols: list of tuples
Column names and dtypes. Each tuple must be written as `(name, dtype)`.
Returns
-------
out : structured array
Initialized structured array.
""" """
if not (isinstance(cols, list) if not (isinstance(cols, list)
and all(isinstance(c, tuple) and len(c) == 2 for c in cols)): and all(isinstance(c, tuple) and len(c) == 2 for c in cols)):
@ -167,19 +40,6 @@ def cols_to_structured(N, cols):
def add_columns(arr, X, cols): def add_columns(arr, X, cols):
""" """
Add new columns `X` to a record array `arr`. Creates a new array. Add new columns `X` to a record array `arr`. Creates a new array.
Parameters
----------
arr : structured array
Structured array to add columns to.
X : (list of) 1-dimensional array(s)
Columns to be added.
cols : str or list of str
Column names to be added.
Returns
-------
out : structured array
""" """
cols = [cols] if isinstance(cols, str) else cols cols = [cols] if isinstance(cols, str) else cols
@ -214,17 +74,6 @@ def add_columns(arr, X, cols):
def rm_columns(arr, cols): def rm_columns(arr, cols):
""" """
Remove columns `cols` from a structured array `arr`. Allocates a new array. Remove columns `cols` from a structured array `arr`. Allocates a new array.
Parameters
----------
arr : structured array
Structured array to remove columns from.
cols : str or list of str
Column names to be removed.
Returns
-------
out : structured array
""" """
# Ensure cols is a list # Ensure cols is a list
cols = [cols] if isinstance(cols, str) else cols cols = [cols] if isinstance(cols, str) else cols
@ -247,16 +96,7 @@ def rm_columns(arr, cols):
def flip_cols(arr, col1, col2): def flip_cols(arr, col1, col2):
""" """
Flip values in columns `col1` and `col2`. `arr` is modified in place. Flip values in columns `col1` and `col2` of a structured array `arr`.
Parameters
----------
arr : structured array
Array whose columns are to be flipped.
col1 : str
First column name.
col2 : str
Second column name.
""" """
if col1 not in arr.dtype.names or col2 not in arr.dtype.names: 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`.") raise ValueError(f"Both `{col1}` and `{col2}` must exist in `arr`.")

206
csiborgtools/utils.py
View file

@ -16,34 +16,25 @@
import numpy import numpy
from numba import jit from numba import jit
###############################################################################
# Positions #
###############################################################################
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def center_of_mass(points, mass, boxsize): def center_of_mass(particle_positions, particles_mass, boxsize):
""" """
Calculate the center of mass of a halo while assuming periodic boundary Calculate the center of mass of a halo while assuming periodic boundary
conditions of a cubical box. Assuming that particle positions are in conditions of a cubical box.
`[0, boxsize)` range. This is a JIT implementation.
Parameters
----------
points : 2-dimensional array of shape (n_particles, 3)
Particle position array.
mass : 1-dimensional array of shape `(n_particles, )`
Particle mass array.
boxsize : float
Box size in the same units as `parts` coordinates.
Returns
-------
cm : 1-dimensional array of shape `(3, )`
""" """
cm = numpy.zeros(3, dtype=points.dtype) cm = numpy.zeros(3, dtype=particle_positions.dtype)
totmass = sum(mass) totmass = sum(particles_mass)
# Convert positions to unit circle coordinates in the complex plane, # Convert positions to unit circle coordinates in the complex plane,
# calculate the weighted average and convert it back to box coordinates. # calculate the weighted average and convert it back to box coordinates.
for i in range(3): for i in range(3):
cm_i = sum(mass * numpy.exp(2j * numpy.pi * points[:, i] / boxsize)) cm_i = sum(particles_mass * numpy.exp(
2j * numpy.pi * particle_positions[:, i] / boxsize))
cm_i /= totmass cm_i /= totmass
cm_i = numpy.arctan2(cm_i.imag, cm_i.real) * boxsize / (2 * numpy.pi) cm_i = numpy.arctan2(cm_i.imag, cm_i.real) * boxsize / (2 * numpy.pi)
@ -55,24 +46,11 @@ def center_of_mass(points, mass, boxsize):
return cm return cm
@jit(nopython=True) @jit(nopython=True, fastmath=True, boundscheck=False)
def periodic_distance(points, reference, boxsize): def periodic_distance(points, reference_point, boxsize):
""" """
Compute the 3D distance between multiple points and a reference point using Compute the 3D distance between multiple points and a reference point using
periodic boundary conditions. This is an optimized JIT implementation. periodic boundary conditions.
Parameters
----------
points : 2-dimensional array of shape `(n_points, 3)`
Points to calculate the distance from the reference point.
reference : 1-dimensional array of shape `(3, )`
Reference point.
boxsize : float
Box size.
Returns
-------
dist : 1-dimensional array of shape `(n_points, )`
""" """
npoints = len(points) npoints = len(points)
half_box = boxsize / 2 half_box = boxsize / 2
@ -80,7 +58,7 @@ def periodic_distance(points, reference, boxsize):
dist = numpy.zeros(npoints, dtype=points.dtype) dist = numpy.zeros(npoints, dtype=points.dtype)
for i in range(npoints): for i in range(npoints):
for j in range(3): for j in range(3):
dist_1d = abs(points[i, j] - reference[j]) dist_1d = abs(points[i, j] - reference_point[j])
if dist_1d > (half_box): if dist_1d > (half_box):
dist_1d = boxsize - dist_1d dist_1d = boxsize - dist_1d
@ -94,23 +72,7 @@ def periodic_distance(points, reference, boxsize):
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def periodic_distance_two_points(p1, p2, boxsize): def periodic_distance_two_points(p1, p2, boxsize):
""" """Compute the 3D distance between two points in a periodic box."""
Compute the 3D distance between two points using periodic boundary
conditions. This is an optimized JIT implementation.
Parameters
----------
p1 : 1-dimensional array of shape `(3, )`
First point.
p2 : 1-dimensional array of shape `(3, )`
Second point.
boxsize : float
Box size.
Returns
-------
dist : 1-dimensional array of shape `(n_points, )`
"""
half_box = boxsize / 2 half_box = boxsize / 2
dist = 0 dist = 0
@ -125,47 +87,133 @@ def periodic_distance_two_points(p1, p2, boxsize):
return dist**0.5 return dist**0.5
@jit(nopython=True)
def periodic_wrap_grid(pos, boxsize=1):
"""Wrap positions in a periodic box."""
for n in range(pos.shape[0]):
for i in range(3):
if pos[n, i] > boxsize:
pos[n, i] -= boxsize
elif pos[n, i] < 0:
pos[n, i] += boxsize
return pos
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def delta2ncells(delta): def delta2ncells(field):
""" """
Calculate the number of cells in `delta` that are non-zero. Calculate the number of cells in `field` that are non-zero.
Parameters
----------
delta : 3-dimensional array
Halo density field.
Returns
-------
ncells : int
Number of non-zero cells.
""" """
tot = 0 tot = 0
imax, jmax, kmax = delta.shape imax, jmax, kmax = field.shape
for i in range(imax): for i in range(imax):
for j in range(jmax): for j in range(jmax):
for k in range(kmax): for k in range(kmax):
if delta[i, j, k] > 0: if field[i, j, k] > 0:
tot += 1 tot += 1
return tot return tot
def cartesian_to_radec(X):
"""
Calculate the radial distance, RA [0, 360) deg and dec [-90, 90] deg.
"""
x, y, z = X[:, 0], X[:, 1], X[:, 2]
dist = numpy.linalg.norm(X, axis=1)
dec = numpy.arcsin(z / dist)
ra = numpy.arctan2(y, x)
ra[ra < 0] += 2 * numpy.pi
ra *= 180 / numpy.pi
dec *= 180 / numpy.pi
return numpy.vstack([dist, ra, dec]).T
def radec_to_cartesian(X):
"""
Calculate Cartesian coordinates from radial distance, RA [0, 360) deg and
dec [-90, 90] deg.
"""
dist, ra, dec = X[:, 0], X[:, 1], X[:, 2]
ra *= numpy.pi / 180
dec *= numpy.pi / 180
cdec = numpy.cos(dec)
return numpy.vstack([
dist * cdec * numpy.cos(ra),
dist * cdec * numpy.sin(ra),
dist * numpy.sin(dec)
]).T
def real2redshift(pos, vel, observer_location, observer_velocity, box,
periodic_wrap=True, make_copy=True):
r"""
Convert real-space position to redshift space position.
Parameters
----------
pos : 2-dimensional array `(nsamples, 3)`
Real-space Cartesian components in `Mpc / h`.
vel : 2-dimensional array `(nsamples, 3)`
Cartesian velocity in `km / s`.
observer_location: 1-dimensional array `(3,)`
Observer location in `Mpc / h`.
observer_velocity: 1-dimensional array `(3,)`
Observer velocity in `km / s`.
box : py:class:`csiborg.read.CSiBORGBox`
Box units.
periodic_wrap : bool, optional
Whether to wrap around the box, particles may be outside the default
bounds once RSD is applied.
make_copy : bool, optional
Whether to make a copy of `pos` before modifying it.
Returns
-------
pos : 2-dimensional array `(nsamples, 3)`
Redshift-space Cartesian position in `Mpc / h`.
"""
if make_copy:
pos = numpy.copy(pos)
vel = numpy.copy(vel)
H0_inv = 1. / 100
# Place the observer at the origin
pos -= observer_location
vel -= observer_velocity
vr_dot = numpy.einsum('ij,ij->i', pos, vel)
norm2 = numpy.einsum('ij,ij->i', pos, pos)
pos *= (1 + H0_inv * vr_dot / norm2).reshape(-1, 1)
# Place the observer back
pos += observer_location
if not make_copy:
vel += observer_velocity
if periodic_wrap:
boxsize = box.box2mpc(1.)
pos = periodic_wrap_grid(pos, boxsize)
return pos
###############################################################################
# Statistics #
###############################################################################
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def number_counts(x, bin_edges): def number_counts(x, bin_edges):
""" """
Calculate counts of samples in bins. Calculate counts of samples in bins.
Parameters
----------
x : 1-dimensional array
Samples' values.
bin_edges : 1-dimensional array
Bin edges.
Returns
-------
counts : 1-dimensional array
Bin counts.
""" """
out = numpy.full(bin_edges.size - 1, numpy.nan, dtype=numpy.float32) out = numpy.full(bin_edges.size - 1, numpy.nan, dtype=numpy.float32)
for i in range(bin_edges.size - 1): for i in range(bin_edges.size - 1):

108
scripts/field_obs_vp.py Normal file
View file

@ -0,0 +1,108 @@
# 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)

246
scripts/field_prop.py
View file

@ -43,44 +43,29 @@ from utils import get_nsims
def density_field(nsim, parser_args, to_save=True): def density_field(nsim, parser_args, to_save=True):
""" """
Calculate the density field in the CSiBORG simulation. Calculate the density field in the CSiBORG simulation.
Parameters
----------
nsim : int
Simulation index.
parser_args : argparse.Namespace
Parsed arguments.
to_save : bool, optional
Whether to save the output to disk.
Returns
-------
field : 3-dimensional array
""" """
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg")) nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths) box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
if not parser_args.in_rsp:
parts = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg")) parts = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg"))
parts = parts["particles"] parts = parts["particles"]
gen = csiborgtools.field.DensityField(box, parser_args.MAS) gen = csiborgtools.field.DensityField(box, parser_args.MAS)
field = gen(parts, parser_args.grid, verbose=parser_args.verbose)
if parser_args.kind == "density":
field = gen(parts, parser_args.grid, in_rsp=False,
verbose=parser_args.verbose)
if parser_args.in_rsp:
field = csiborgtools.field.field2rsp(*field, parts=parts, box=box,
verbose=parser_args.verbose)
else: else:
field = gen(parts, parser_args.grid, in_rsp=parser_args.in_rsp, field = numpy.load(paths.field(
verbose=parser_args.verbose) "density", parser_args.MAS, parser_args.grid, nsim, False))
radvel_field = numpy.load(paths.field(
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
if parser_args.smooth_scale > 0: field = csiborgtools.field.field2rsp(field, radvel_field, box,
field = csiborgtools.field.smoothen_field( parser_args.MAS)
field, parser_args.smooth_scale, box.boxsize * box.h, threads=1)
if to_save: if to_save:
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid, fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp, parser_args.smooth_scale) nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.") print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field) numpy.save(fout, field)
return field return field
@ -93,35 +78,20 @@ def density_field(nsim, parser_args, to_save=True):
def velocity_field(nsim, parser_args, to_save=True): def velocity_field(nsim, parser_args, to_save=True):
""" """
Calculate the velocity field in the CSiBORG simulation. Calculate the velocity field in a CSiBORG simulation.
Parameters
----------
nsim : int
Simulation index.
parser_args : argparse.Namespace
Parsed arguments.
to_save : bool, optional
Whether to save the output to disk.
Returns
-------
velfield : 4-dimensional array
""" """
if parser_args.in_rsp: if parser_args.in_rsp:
raise NotImplementedError("Velocity field in RSP is not implemented.") raise NotImplementedError("Velocity field in RSP is not implemented.")
if parser_args.smooth_scale > 0:
raise NotImplementedError(
"Smoothed velocity field is not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
mpart = 1.1641532e-10 # Particle mass in CSiBORG simulations.
nsnap = max(paths.get_snapshots(nsim, "csiborg")) nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths) box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
parts = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg")) parts = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg"))
parts = parts["particles"] parts = parts["particles"]
gen = csiborgtools.field.VelocityField(box, parser_args.MAS) gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
field = gen(parts, parser_args.grid, mpart, verbose=parser_args.verbose) field = gen(parts, parser_args.grid, verbose=parser_args.verbose)
if to_save: if to_save:
fout = paths.field("velocity", parser_args.MAS, parser_args.grid, fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
@ -131,57 +101,6 @@ def velocity_field(nsim, parser_args, to_save=True):
return field return field
###############################################################################
# Potential field #
###############################################################################
def potential_field(nsim, parser_args, to_save=True):
"""
Calculate the potential field in the CSiBORG simulation.
Parameters
----------
nsim : int
Simulation index.
parser_args : argparse.Namespace
Parsed arguments.
to_save : bool, optional
Whether to save the output to disk.
Returns
-------
potential : 3-dimensional array
"""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
# Load the real space overdensity field
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False))
if parser_args.smooth_scale > 0:
rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
box.boxsize * box.h, threads=1)
rho = density_gen.overdensity_field(rho)
# Calculate the real space potentiel field
gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
field = gen(rho)
if parser_args.in_rsp:
parts = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg"))
parts = parts["particles"]
field = csiborgtools.field.field2rsp(*field, parts=parts, box=box,
verbose=parser_args.verbose)
if to_save:
fout = paths.field(parser_args.kind, 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, field)
return field
############################################################################### ###############################################################################
# Radial velocity field # # Radial velocity field #
############################################################################### ###############################################################################
@ -190,33 +109,21 @@ def potential_field(nsim, parser_args, to_save=True):
def radvel_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. Calculate the radial velocity field in the CSiBORG simulation.
Parameters
----------
nsim : int
Simulation index.
parser_args : argparse.Namespace
Parsed arguments.
to_save : bool, optional
Whether to save the output to disk.
Returns
-------
radvel : 3-dimensional array
""" """
if parser_args.in_rsp: if parser_args.in_rsp:
raise NotImplementedError("Radial vel. field in RSP not implemented.") raise NotImplementedError("Radial vel. field in RSP not implemented.")
if parser_args.smooth_scale > 0:
raise NotImplementedError(
"Smoothed radial vel. field not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg")) nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths) box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid, vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)) nsim, parser_args.in_rsp))
observer_velocity = csiborgtools.field.observer_vobs(vel)
gen = csiborgtools.field.VelocityField(box, parser_args.MAS) gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
field = gen.radial_velocity(vel) field = gen.radial_velocity(vel, observer_velocity)
if to_save: if to_save:
fout = paths.field("radvel", parser_args.MAS, parser_args.grid, fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp) nsim, parser_args.in_rsp)
@ -224,6 +131,43 @@ def radvel_field(nsim, parser_args, to_save=True):
numpy.save(fout, field) numpy.save(fout, field)
return 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.CSiBORGBox(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 # # Environment classification #
@ -233,77 +177,47 @@ def radvel_field(nsim, parser_args, to_save=True):
def environment_field(nsim, parser_args, to_save=True): def environment_field(nsim, parser_args, to_save=True):
""" """
Calculate the environmental classification in the CSiBORG simulation. Calculate the environmental classification in the CSiBORG simulation.
Parameters
----------
nsim : int
Simulation index.
parser_args : argparse.Namespace
Parsed arguments.
to_save : bool, optional
Whether to save the output to disk.
Returns
-------
env : 3-dimensional array
""" """
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim, "csiborg")) nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths) box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
# Load the real space overdensity field rho = numpy.load(paths.field(
if parser_args.verbose: "density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
print(f"{datetime.now()}: loading density field.") density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False))
if parser_args.smooth_scale > 0:
rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
box.boxsize * box.h, threads=1)
rho = density_gen.overdensity_field(rho) rho = density_gen.overdensity_field(rho)
# Calculate the real space tidal tensor field, delete overdensity.
if parser_args.verbose: gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
print(f"{datetime.now()}: calculating tidal tensor field.") field = gen(rho)
tensor_field = gen(rho)
del rho del rho
collect() collect()
# Optionally drag the field to RSP.
if parser_args.in_rsp: if parser_args.in_rsp:
parts = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg")) radvel_field = numpy.load(paths.field(
parts = parts["particles"] "radvel", parser_args.MAS, parser_args.grid, nsim, False))
fields = (tensor_field.T00, tensor_field.T11, tensor_field.T22, args = (radvel_field, box, parser_args.MAS)
tensor_field.T01, tensor_field.T02, tensor_field.T12)
T00, T11, T22, T01, T02, T12 = csiborgtools.field.field2rsp( field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
*fields, parts=parts, box=box, verbose=parser_args.verbose) field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
tensor_field.T00[...] = T00 field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
tensor_field.T11[...] = T11 field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
tensor_field.T22[...] = T22 field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
tensor_field.T01[...] = T01 field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
tensor_field.T02[...] = T02
tensor_field.T12[...] = T12 del radvel_field
del T00, T11, T22, T01, T02, T12
collect() collect()
# Calculate the eigenvalues of the tidal tensor field, delete tensor field. eigvals = gen.tensor_field_eigvals(field)
if parser_args.verbose:
print(f"{datetime.now()}: calculating eigenvalues.") del field
eigvals = gen.tensor_field_eigvals(tensor_field)
del tensor_field
collect() collect()
# Classify the environment based on the eigenvalues.
if parser_args.verbose:
print(f"{datetime.now()}: classifying environment.")
env = gen.eigvals_to_environment(eigvals) env = gen.eigvals_to_environment(eigvals)
del eigvals
collect()
if to_save: if to_save:
fout = paths.field("environment", parser_args.MAS, parser_args.grid, fout = paths.field("environment", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp, parser_args.smooth_scale) nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.") print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, env) numpy.save(fout, env)
return env return env
@ -327,8 +241,6 @@ if __name__ == "__main__":
parser.add_argument("--grid", type=int, help="Grid resolution.") parser.add_argument("--grid", type=int, help="Grid resolution.")
parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)), parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
help="Calculate in RSP?") help="Calculate in RSP?")
parser.add_argument("--smooth_scale", type=float, default=0,
help="Smoothing scale in Mpc/h.")
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)), parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
help="Verbosity flag for reading in particles.") help="Verbosity flag for reading in particles.")
parser.add_argument("--simname", type=str, default="csiborg", parser.add_argument("--simname", type=str, default="csiborg",

23
scripts/fit_hmf.py
View file

@ -37,21 +37,6 @@ except ModuleNotFoundError:
def get_counts(nsim, bins, paths, parser_args): def get_counts(nsim, bins, paths, parser_args):
""" """
Calculate and save the number of haloes in each mass bin. Calculate and save the number of haloes in each mass bin.
Parameters
----------
nsim : int
Simulation index.
bins : 1-dimensional array
Array of bin edges (in log10 mass).
paths : csiborgtools.read.Paths
Paths object.
parser_args : argparse.Namespace
Parsed command-line arguments.
Returns
-------
None
""" """
simname = parser_args.simname simname = parser_args.simname
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
@ -75,13 +60,14 @@ def get_counts(nsim, bins, paths, parser_args):
counts[nobs, :] = csiborgtools.number_counts(logmass, bins) counts[nobs, :] = csiborgtools.number_counts(logmass, bins)
elif simname == "quijote_full": elif simname == "quijote_full":
cat = csiborgtools.read.QuijoteHaloCatalogue( cat = csiborgtools.read.QuijoteHaloCatalogue(
nsim, paths, nsnap=4, load_fitted=False, load_initial=False) nsim, paths, nsnap=4, load_fitted=False, load_initial=False,
load_backup=parser_args.from_quijote_backup)
logmass = numpy.log10(cat["group_mass"]) logmass = numpy.log10(cat["group_mass"])
counts = csiborgtools.number_counts(logmass, bins) counts = csiborgtools.number_counts(logmass, bins)
else: else:
raise ValueError(f"Unknown simulation name `{simname}`.") raise ValueError(f"Unknown simulation name `{simname}`.")
fout = paths.halo_counts(simname, nsim) fout = paths.halo_counts(simname, nsim, parser_args.from_quijote_backup)
if parser_args.verbose: if parser_args.verbose:
print(f"{datetime.now()}: saving halo counts to `{fout}`.") print(f"{datetime.now()}: saving halo counts to `{fout}`.")
numpy.savez(fout, counts=counts, bins=bins, rmax=parser_args.Rmax) numpy.savez(fout, counts=counts, bins=bins, rmax=parser_args.Rmax)
@ -97,6 +83,9 @@ if __name__ == "__main__":
parser.add_argument( parser.add_argument(
"--Rmax", type=float, default=155, "--Rmax", type=float, default=155,
help="High-res region radius in Mpc / h. Ignored for `quijote_full`.") help="High-res region radius in Mpc / h. Ignored for `quijote_full`.")
parser.add_argument("--from_quijote_backup",
type=lambda x: bool(strtobool(x)), default=False,
help="Flag to indicate Quijote backup data.")
parser.add_argument("--lims", type=float, nargs="+", default=[11., 16.], parser.add_argument("--lims", type=float, nargs="+", default=[11., 16.],
help="Mass limits in Msun / h.") help="Mass limits in Msun / h.")
parser.add_argument("--bw", type=float, default=0.2, parser.add_argument("--bw", type=float, default=0.2,

9
scripts/fit_init.py
View file

@ -41,15 +41,6 @@ def _main(nsim, simname, verbose):
""" """
Calculate the Lagrangian halo centre of mass and Lagrangian patch size in Calculate the Lagrangian halo centre of mass and Lagrangian patch size in
the initial snapshot. the initial snapshot.
Parameters
----------
nsim : int
IC realisation index.
simname : str
Simulation name.
verbose : bool
Verbosity flag.
""" """
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
cols = [("index", numpy.int32), cols = [("index", numpy.int32),

31
scripts/match_all.py
View file

@ -29,16 +29,6 @@ def get_combs(simname):
""" """
Get the list of all pairs of IC indices and permute them with a fixed Get the list of all pairs of IC indices and permute them with a fixed
seed. seed.
Parameters
----------
simname : str
Simulation name.
Returns
-------
combs : list
List of pairs of simulations.
""" """
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
combs = list(combinations(paths.get_ics(simname), 2)) combs = list(combinations(paths.get_ics(simname), 2))
@ -50,27 +40,6 @@ def get_combs(simname):
def main(comb, kind, simname, min_logmass, sigma, mult, verbose): def main(comb, kind, simname, min_logmass, sigma, mult, verbose):
""" """
Match a pair of simulations. Match a pair of simulations.
Parameters
----------
comb : tuple
Pair of simulation IC indices.
kind : str
Kind of matching.
simname : str
Simulation name.
min_logmass : float
Minimum log halo mass.
sigma : float
Smoothing scale in number of grid cells.
mult : float
Multiplicative factor for search radius.
verbose : bool
Verbosity flag.
Returns
-------
None
""" """
nsim0, nsimx = comb nsim0, nsimx = comb
if kind == "overlap": if kind == "overlap":

11
scripts/utils.py
View file

@ -52,7 +52,7 @@ def get_nsims(args, paths):
Simulation indices. Simulation indices.
""" """
if args.nsims is None or args.nsims[0] == -1: if args.nsims is None or args.nsims[0] == -1:
nsims = paths.get_ics(args.simname) nsims = paths.get_ics(args.simname, args.from_quijote_backup)
else: else:
nsims = args.nsims nsims = args.nsims
return list(nsims) return list(nsims)
@ -93,9 +93,14 @@ def read_single_catalogue(args, config, nsim, run, rmax, paths, nobs=None):
nsim, paths, load_fitted=True, load_inital=True, nsim, paths, load_fitted=True, load_inital=True,
with_lagpatch=False) with_lagpatch=False)
else: else:
if args.from_quijote_backup:
load_fitted = False
load_initial = False
cat = csiborgtools.read.QuijoteHaloCatalogue( cat = csiborgtools.read.QuijoteHaloCatalogue(
nsim, paths, nsnap=4, load_fitted=True, load_initial=True, nsim, paths, nsnap=4, load_fitted=load_fitted,
with_lagpatch=False) load_initial=load_initial, with_lagpatch=False,
load_backup=args.from_quijote_backup)
if nobs is not None: if nobs is not None:
# We may optionally already here pick a fiducial observer. # We may optionally already here pick a fiducial observer.
cat = cat.pick_fiducial_observer(nobs, args.Rmax) cat = cat.pick_fiducial_observer(nobs, args.Rmax)

36
scripts_plots/plot_data.py
View file

@ -305,7 +305,7 @@ def plot_hmf_quijote_full(pdf=False):
plt.close() plt.close()
def load_field(kind, nsim, grid, MAS, in_rsp=False, smooth_scale=None): def load_field(kind, nsim, grid, MAS, in_rsp=False):
r""" r"""
Load a single field. Load a single field.
@ -321,16 +321,13 @@ def load_field(kind, nsim, grid, MAS, in_rsp=False, smooth_scale=None):
Mass assignment scheme. Mass assignment scheme.
in_rsp : bool, optional in_rsp : bool, optional
Whether to load the field in redshift space. Whether to load the field in redshift space.
smooth_scale : float, optional
Smoothing scale in :math:`\mathrm{Mpc} / h`.
Returns Returns
------- -------
field : n-dimensional array field : n-dimensional array
""" """
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
return numpy.load(paths.field(kind, MAS, grid, nsim, in_rsp=in_rsp, return numpy.load(paths.field(kind, MAS, grid, nsim, in_rsp=in_rsp))
smooth_scale=smooth_scale))
############################################################################### ###############################################################################
@ -377,16 +374,16 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths) box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
if kind == "overdensity": if kind == "overdensity":
field = load_field("density", nsim, grid, MAS=MAS, in_rsp=in_rsp, field = load_field("density", nsim, grid, MAS=MAS, in_rsp=in_rsp)
smooth_scale=smooth_scale)
density_gen = csiborgtools.field.DensityField(box, MAS) density_gen = csiborgtools.field.DensityField(box, MAS)
field = density_gen.overdensity_field(field) + 1 field = density_gen.overdensity_field(field) + 2
field = numpy.log10(field)
elif kind == "borg_density": elif kind == "borg_density":
field = File(paths.borg_mcmc(nsim), 'r') field = File(paths.borg_mcmc(nsim), 'r')
field = field["scalars"]["BORG_final_density"][...] field = field["scalars"]["BORG_final_density"][...]
else: else:
field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp, field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp)
smooth_scale=smooth_scale)
if kind == "velocity": if kind == "velocity":
field = field[vel_component, ...] field = field[vel_component, ...]
@ -423,6 +420,9 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
else: else:
ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap) ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap=cmap)
k = img.shape[0] // 2
ax[i].scatter(k, k, marker="x", s=5, zorder=2, c="red")
frad = 155.5 / 677.7 frad = 155.5 / 677.7
R = 155.5 / 677.7 * grid R = 155.5 / 677.7 * grid
if slice_find is None: if slice_find is None:
@ -487,6 +487,7 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
else: else:
fig.colorbar(im, cax=cbar_ax, label=clabel) fig.colorbar(im, cax=cbar_ax, label=clabel)
fig.tight_layout(h_pad=0, w_pad=0) fig.tight_layout(h_pad=0, w_pad=0)
for ext in ["png"] if pdf is False else ["png", "pdf"]: for ext in ["png"] if pdf is False else ["png", "pdf"]:
fout = join( fout = join(
@ -652,7 +653,7 @@ if __name__ == "__main__":
if False: if False:
plot_mass_vs_ncells(7444, pdf=False) plot_mass_vs_ncells(7444, pdf=False)
if True: if False:
plot_hmf(pdf=False) plot_hmf(pdf=False)
if False: if False:
@ -665,16 +666,17 @@ if __name__ == "__main__":
plot_groups=False, dmin=45, dmax=60, plot_groups=False, dmin=45, dmax=60,
plot_halos=5e13, volume_weight=True) plot_halos=5e13, volume_weight=True)
if False: if True:
kind = "overdensity" kind = "potential"
grid = 256 grid = 512
smooth_scale = 0 smooth_scale = 0
# plot_projected_field("overdensity", 7444, grid, in_rsp=True, # plot_projected_field("overdensity", 7444, grid, in_rsp=True,
# highres_only=False) # highres_only=False)
plot_projected_field(kind, 7444, grid, in_rsp=False, # for in_rsp in [True, False]:
for in_rsp in [True, False]:
plot_projected_field(kind, 7444, grid, in_rsp=in_rsp,
smooth_scale=smooth_scale, slice_find=0.5, smooth_scale=smooth_scale, slice_find=0.5,
MAS="PCS", MAS="PCS", highres_only=True)
highres_only=True)
if False: if False:
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)

129
scripts_plots/plot_match.py
View file

@ -191,7 +191,7 @@ def get_mtot_vs_maxpairoverlap(nsim0, simname, mass_kind, min_logmass,
def get_max(y_): def get_max(y_):
if len(y_) == 0: if len(y_) == 0:
return 0 return 0
return numpy.max(y_) return numpy.nanmax(y_)
reader = csiborgtools.read.NPairsOverlap(cat0, catxs, paths, min_logmass) reader = csiborgtools.read.NPairsOverlap(cat0, catxs, paths, min_logmass)
@ -218,7 +218,6 @@ def mtot_vs_maxpairoverlap(nsim0, simname, mass_kind, min_logmass, smoothed,
x, y, xbins = get_mtot_vs_maxpairoverlap(nsim0, simname, mass_kind, x, y, xbins = get_mtot_vs_maxpairoverlap(nsim0, simname, mass_kind,
min_logmass, smoothed, nbins) min_logmass, smoothed, nbins)
plt.close("all")
with plt.style.context(plt_utils.mplstyle): with plt.style.context(plt_utils.mplstyle):
plt.figure() plt.figure()
plt.hexbin(x, y, mincnt=1, gridsize=50, bins="log") plt.hexbin(x, y, mincnt=1, gridsize=50, bins="log")
@ -252,6 +251,87 @@ def mtot_vs_maxpairoverlap(nsim0, simname, mass_kind, min_logmass, smoothed,
plt.close() plt.close()
# --------------------------------------------------------------------------- #
###############################################################################
# Total DM halo mass vs maximum pair overlap consistency #
###############################################################################
# --------------------------------------------------------------------------- #
@cache_to_disk(120)
def get_mtot_vs_maxpairoverlap_consistency(nsim0, simname, mass_kind,
min_logmass, smoothed):
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsimxs = csiborgtools.read.get_cross_sims(simname, nsim0, paths,
min_logmass, smoothed=smoothed)
cat0 = open_cat(nsim0, simname)
catxs = open_cats(nsimxs, simname)
reader = csiborgtools.read.NPairsOverlap(cat0, catxs, paths, min_logmass)
x = numpy.log10(cat0[mass_kind])
mask = x > min_logmass
x = x[mask]
nhalos = len(x)
y = numpy.full((len(catxs), nhalos), numpy.nan)
for i in trange(len(catxs), desc="Stacking catalogues"):
overlaps = reader[i].overlap(smoothed)
for j in range(nhalos):
# if len(overlaps[j]) > 0:
y[i, j] = numpy.sum(overlaps[j])
return x, y
def mtot_vs_maxpairoverlap_consistency(nsim0, simname, mass_kind, min_logmass,
smoothed, ext="png"):
left_edges = numpy.arange(min_logmass, 15, 0.1)
# delete_disk_caches_for_function("get_mtot_vs_maxpairoverlap_consistency")
x, y0 = get_mtot_vs_maxpairoverlap_consistency(nsim0, simname, mass_kind,
min_logmass, smoothed)
nsims, nhalos = y0.shape
x_2, y0_2 = get_mtot_vs_maxpairoverlap_consistency(
0, "quijote", "group_mass", min_logmass, smoothed)
nsims2, nhalos = y0_2.shape
with plt.style.context(plt_utils.mplstyle):
plt.figure()
yplot = numpy.full(len(left_edges), numpy.nan)
yplot_2 = numpy.full(len(left_edges), numpy.nan)
for ymin in [0.3]:
y = numpy.sum(y0 > ymin, axis=0) / nsims
y_2 = numpy.sum(y0_2 > ymin, axis=0) / nsims2
for i, left_edge in enumerate(left_edges):
mask = x > left_edge
yplot[i] = numpy.mean(y[mask]) #/ nsims
mask = x_2 > left_edge
yplot_2[i] = numpy.mean(y_2[mask]) #/ nsims
plt.plot(left_edges, yplot, label="CSiBORG")
plt.plot(left_edges, yplot_2, label="Quijote")
plt.legend()
# y2 = numpy.concatenate(y0)
# y2 = y2[y2 > 0]
# m = y0 > 0
# plt.hist(y0[m], bins=30, density=True, histtype="step")
# m = y0_2 > 0
# plt.hist(y0_2[m], bins=30, density=True, histtype="step")
# plt.yscale("log")
plt.tight_layout()
fout = join(
plt_utils.fout,
f"mass_vs_max_pair_overlap_consistency_{simname}_{nsim0}.{ext}")
print(f"Saving to `{fout}`.")
plt.savefig(fout, dpi=plt_utils.dpi, bbox_inches="tight")
plt.close()
# --------------------------------------------------------------------------- # # --------------------------------------------------------------------------- #
############################################################################### ###############################################################################
# Total DM halo mass vs summed pair overlaps # # Total DM halo mass vs summed pair overlaps #
@ -876,8 +956,7 @@ def get_matching_max_vs_overlap(simname, nsim0, min_logmass, mult):
def matching_max_vs_overlap(simname, nsim0, min_logmass): def matching_max_vs_overlap(simname, nsim0, min_logmass):
left_edges = numpy.arange(min_logmass, 15, 0.1) left_edges = numpy.arange(min_logmass, 15, 0.1)
nsims = 100 if simname == "csiborg" else 9
delete_disk_caches_for_function("get_matching_max_vs_overlap")
with plt.style.context("science"): with plt.style.context("science"):
fig, axs = plt.subplots(ncols=2, figsize=(3.5 * 2, 2.625)) fig, axs = plt.subplots(ncols=2, figsize=(3.5 * 2, 2.625))
@ -891,34 +970,36 @@ def matching_max_vs_overlap(simname, nsim0, min_logmass):
success = x["success"] success = x["success"]
nbins = len(left_edges) nbins = len(left_edges)
y = numpy.full((nbins, 100), numpy.nan) y = numpy.full((nbins, nsims), numpy.nan)
y2 = numpy.full((nbins, 100), numpy.nan) y2 = numpy.full(nbins, numpy.nan)
y2err = numpy.full(nbins, numpy.nan)
for i in range(nbins): for i in range(nbins):
m = mass0 > left_edges[i] m = mass0 > left_edges[i]
for j in range(100): for j in range(nsims):
y[i, j] = numpy.sum( y[i, j] = numpy.sum(
max_overlap[m, j] == match_overlap[m, j]) max_overlap[m, j] == match_overlap[m, j])
y[i, j] /= numpy.sum(success[m, j]) y[i, j] /= numpy.sum(success[m, j])
y2[i, j] = numpy.sum(success[m, j]) / numpy.sum(m)
offset = numpy.random.normal(0, 0.01) y2[i] = numpy.mean(numpy.sum(success[m, :], axis=1) / nsims)
y2err[i] = numpy.std(numpy.sum(success[m, :], axis=1) / nsims)
offset = numpy.random.normal(0, 0.015)
ysummary = numpy.percentile(y, [16, 50, 84], axis=1) ysummary = numpy.percentile(y, [16, 50, 84], axis=1)
axs[0].errorbar( axs[0].errorbar(
left_edges + offset, ysummary[1], left_edges + offset, ysummary[1],
yerr=[ysummary[1] - ysummary[0], ysummary[2] - ysummary[1]], yerr=[ysummary[1] - ysummary[0], ysummary[2] - ysummary[1]],
capsize=3, c=cols[n], capsize=4, c=cols[n], ls="dashed",
label=r"$\leq {}~R_{{\rm 200c}}$".format(mult), errorevery=2) label=r"$\leq {}~R_{{\rm 200c}}$".format(mult), errorevery=2)
ysummary = numpy.percentile(y2, [16, 50, 84], axis=1)
axs[1].errorbar( axs[1].errorbar(left_edges + offset, y2, yerr=y2err,
left_edges + offset, ysummary[1], ls="--", capsize=4, errorevery=2, c=cols[n], ls="dashed")
yerr=[ysummary[1] - ysummary[0], ysummary[2] - ysummary[1]],
capsize=3, c=cols[n], errorevery=2)
axs[0].legend(ncols=2, fontsize="small") axs[0].legend(ncols=2, fontsize="small")
for i in range(2): for i in range(2):
axs[i].set_xlabel(r"$\log M_{\rm tot, min} ~ [M_\odot / h]$") axs[i].set_xlabel(r"$\log M_{\rm tot, min} ~ [M_\odot / h]$")
axs[1].set_ylim(0)
axs[0].set_ylabel(r"$f_{\rm agreement}$") axs[0].set_ylabel(r"$f_{\rm agreement}$")
axs[1].set_ylabel(r"$f_{\rm match}$") axs[1].set_ylabel(r"$f_{\rm match}$")
@ -1048,7 +1129,7 @@ if __name__ == "__main__":
smoothed = True smoothed = True
nbins = 10 nbins = 10
ext = "png" ext = "png"
plot_quijote = False plot_quijote = True
min_maxoverlap = 0. min_maxoverlap = 0.
funcs = [ funcs = [
@ -1128,5 +1209,19 @@ if __name__ == "__main__":
mtot_vs_maxoverlap_property(0, "quijote", min_logmass, key, mtot_vs_maxoverlap_property(0, "quijote", min_logmass, key,
min_maxoverlap, smoothed) min_maxoverlap, smoothed)
if False:
matching_max_vs_overlap("csiborg", 7444, min_logmass)
if plot_quijote:
matching_max_vs_overlap("quijote", 0, min_logmass)
if True: if True:
matching_max_vs_overlap(7444, min_logmass) mtot_vs_maxpairoverlap_consistency(
7444, "csiborg", "fof_totpartmass", min_logmass, smoothed,
ext="png")
# if plot_quijote:
# mtot_vs_maxpairoverlap_consistency(
# 0, "quijote", "group_mass", min_logmass, smoothed,
# ext="png")

89
scripts_plots/plot_vobs.py Normal file
View file

@ -0,0 +1,89 @@
# 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.
from os.path import join
import matplotlib.pyplot as plt
import numpy
import scienceplots # noqa
from tqdm import tqdm
import csiborgtools
import plt_utils
def observer_peculiar_velocity(MAS, grid, ext="png"):
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = paths.get_ics("csiborg")
for n, nsim in enumerate(nsims):
fpath = paths.observer_peculiar_velocity(MAS, grid, nsim)
f = numpy.load(fpath)
if n == 0:
data = numpy.full((len(nsims), *f["observer_vp"].shape), numpy.nan)
smooth_scales = f["smooth_scales"]
data[n] = f["observer_vp"]
for n, smooth_scale in enumerate(tqdm(smooth_scales,
desc="Plotting smooth scales")):
with plt.style.context(plt_utils.mplstyle):
fig, axs = plt.subplots(ncols=3, figsize=(3.5 * 3, 2.625),
sharey=True)
fig.subplots_adjust(wspace=0)
for i, ax in enumerate(axs):
ax.hist(data[:, n, i], bins="auto", histtype="step")
mu, sigma = numpy.mean(data[:, n, i]), numpy.std(data[:, n, i])
ax.set_title(r"$V_{{\rm obs, i}} = {:.3f} \pm {:.3f} ~ \mathrm{{km}} / \mathrm{{s}}$".format(mu, sigma)) # noqa
axs[0].set_xlabel(r"$V_{\rm obs, x} ~ [\mathrm{km} / \mathrm{s}]$")
axs[1].set_xlabel(r"$V_{\rm obs, y} ~ [\mathrm{km} / \mathrm{s}]$")
axs[2].set_xlabel(r"$V_{\rm obs, z} ~ [\mathrm{km} / \mathrm{s}]$")
axs[0].set_ylabel(r"Counts")
fig.suptitle(r"$N_{{\rm grid}} = {}$, $\sigma_{{\rm smooth}} = {:.2f} ~ [\mathrm{{Mpc}} / h]$".format(grid, smooth_scale)) # noqa
fig.tight_layout(w_pad=0)
fout = join(plt_utils.fout,
f"observer_vp_{grid}_{smooth_scale}.{ext}")
fig.savefig(fout, dpi=plt_utils.dpi, bbox_inches="tight")
plt.close()
with plt.style.context(plt_utils.mplstyle):
fig, axs = plt.subplots(ncols=3, figsize=(3.5 * 3, 2.625))
for i, ax in enumerate(axs):
ymu = numpy.mean(data[..., i], axis=0)
ystd = numpy.std(data[..., i], axis=0)
ylow, yupp = ymu - ystd, ymu + ystd
ax.plot(smooth_scales, ymu, c="k")
ax.fill_between(smooth_scales, ylow, yupp, color="k", alpha=0.2)
ax.set_xlabel(r"$\sigma_{{\rm smooth}} ~ [\mathrm{{Mpc}} / h]$")
axs[0].set_ylabel(r"$V_{\rm obs, x} ~ [\mathrm{km} / \mathrm{s}]$")
axs[1].set_ylabel(r"$V_{\rm obs, y} ~ [\mathrm{km} / \mathrm{s}]$")
axs[2].set_ylabel(r"$V_{\rm obs, z} ~ [\mathrm{km} / \mathrm{s}]$")
fig.suptitle(r"$N_{{\rm grid}} = {}$".format(grid))
fig.tight_layout(w_pad=0)
fout = join(plt_utils.fout, f"observer_vp_summary_{grid}.{ext}")
fig.savefig(fout, dpi=plt_utils.dpi, bbox_inches="tight")
plt.close()
if __name__ == "__main__":
observer_peculiar_velocity("PCS", 512)