"""Module for accounting redshift space distortions."""
from collections.abc import Sequence
import numpy as np
from astropy import cosmology, units
from classy import Class
from cosmotile.cic import cloud_in_cell_los
from scipy import fft
from scipy.interpolate import RegularGridInterpolator
from .wrapper.classy_interface import compute_rms
from .wrapper.inputs import InputParameters
[docs]
def include_dvdr_in_tau21(
brightness_temp: np.ndarray,
los_velocity: np.ndarray,
redshifts: np.ndarray | float,
inputs: InputParameters,
periodic: bool,
tau_21: np.ndarray | None = None,
) -> np.ndarray:
"""Include velocity gradient corrections to the brightness temperature field.
Parameters
----------
brightness_temp
A box of the brightness temperature, without velocity gradient corrections.
los_velocity
Line-of-sight velocities for each cell of brightness_temp, in Mpc/second.
redshifts
An array of the redshifts along the los. Can also be a float (could be useful for coeval boxes)
inputs
The input parameters corresponding to the box.
tau_21
A box of the 21cm optical depth. Not required if inputs.astro_options.USE_TS_FLUCT = False.
periodic
Whether to assume periodic boundary conditions along the line-of-sight.
Returns
-------
tb_with_dvdr
A box of the brightness temperature, with velocity gradient corrections.
"""
if tau_21 is None and inputs.astro_options.USE_TS_FLUCT:
raise ValueError(
"tau_21 is not provided, but inputs.astro_options.USE_TS_FLUCT is True!"
)
if hasattr(redshifts, "__len__") and len(redshifts) != brightness_temp.shape[-1]:
raise ValueError(
"Redshifts must be a float or array with the same size as number of LoS slices"
)
if los_velocity.shape != brightness_temp.shape:
raise ValueError(
"brightness_temp must be an array with the same shape as los_velocity"
)
N = inputs.simulation_options.HII_DIM
dx_los = inputs.simulation_options.BOX_LEN / N # Mpc
# TODO: currently the gradient is applied w.r.t to the z-axis, even if the user specified a different los-axis. Needs to make it more flexible in the future
if periodic:
k_real = fft.rfftfreq(N, dx_los) * 2.0 * np.pi
k_complex = fft.fftfreq(N, dx_los) * 2.0 * np.pi
k_vector = np.stack(np.meshgrid(k_complex, k_complex, k_real, indexing="ij"))
vel_gradient = (
fft.irfftn(1j * k_vector[-1] * fft.rfftn(los_velocity), s=(N, N, N))
/ units.s
)
else:
vel_gradient = np.gradient(
los_velocity * units.Mpc / units.s,
dx_los * units.Mpc,
axis=-1,
edge_order=2,
)
H = inputs.cosmo_params.cosmo.H(redshifts)
if not inputs.astro_options.USE_TS_FLUCT:
# If we don't have spin temperature, we also assume tau_21 << 1 and make the Taylor approximation.
# Clipping is required so the brightness temperature won't diverge when the gradient term goes to small values
max_v_deriv = inputs.astro_params.MAX_DVDR * H
dvdx = np.clip(vel_gradient, -max_v_deriv, max_v_deriv)
gradient_component = np.abs(1.0 + dvdx / H)
tb_with_dvdr = brightness_temp / gradient_component.value
else:
# We have the spin temperature, and we do *not* make the Taylor approximation
tb_no_rsds = brightness_temp.copy()
tau_21 = np.float64(tau_21)
gradient_component = np.float64(np.abs(1.0 + vel_gradient / H))
with np.errstate(
divide="ignore", invalid="ignore"
): # Don't show division by 0 warnings
gradient_factor = (1.0 - np.exp(-tau_21 / gradient_component)) / (
1.0 - (np.exp(-tau_21))
)
# It doesn't really matter what the rsd_factor is when tau_21=0 because the brightness temperature is zero by definition
gradient_factor = np.float32(np.where(tau_21 < 1e-10, 1.0, gradient_factor))
tb_with_dvdr = tb_no_rsds * gradient_factor.value
return tb_with_dvdr
[docs]
def apply_rsds(
field: np.ndarray,
los_velocity: np.ndarray,
redshifts: np.ndarray | float,
inputs: InputParameters,
periodic: bool,
n_rsd_subcells: int = 4,
) -> np.ndarray:
"""Apply redshift-space distortions to a field.
Notes
-----
To ensure that we cover all the slices in the field after the velocities have
been applied, we extrapolate the densities and velocities on either end by the
maximum velocity offset in the field.
Then, to ensure we don't pick up cells with zero particles (after displacement),
we interpolate the slices onto a finer regular grid (in comoving distance) and
then displace the field on that grid, before interpolating back onto the original
slices.
Parameters
----------
field
A box of the field on which redshift-space distortions shall be applied.
los_velocity
Line-of-sight velocities for each cell of brightness_temp, in Mpc/second.
redshifts
An array of the redshifts along the los. Can also be a float (could be useful for coeval boxes)
inputs
The input parameters corresponding to the box.
periodic: bool
Whether to assume periodic boundary conditions along the line-of-sight.
n_rsd_subcells: int, optional
The number of subcells into which each cell is divided when redshift space distortions are applied. Default is 4.
Returns
-------
field_with_rsds
A box of the field, with redshift space distortions.
"""
if hasattr(redshifts, "__len__") and len(redshifts) != field.shape[-1]:
raise ValueError(
"Redshifts must be a float or array with the same size as number of LoS slices"
)
H = inputs.cosmo_params.cosmo.H(redshifts)
los_displacement = los_velocity * units.Mpc / units.s / H
equiv = units.pixel_scale(inputs.simulation_options.cell_size / units.pixel)
los_displacement = los_displacement.to(units.pixel, equivalencies=equiv)
# We transform rectilinear lightcone to be an angular-like lightcone
field_dimensions = len(field.shape)
if field_dimensions == 3:
field = field.reshape((field.shape[0] * field.shape[1], -1))
los_displacement = los_displacement.reshape(
(los_displacement.shape[0] * los_displacement.shape[1], -1)
)
# Here we move the cells along the line of sight, regardless the geometry (rectilinear or angular)
field_with_rsds = rsds_shift(
field=field.T,
los_displacement=los_displacement.T,
n_rsd_subcells=n_rsd_subcells,
periodic=periodic,
).T
# And now we transform back to a rectilinear-like lightcone
if field_dimensions == 3:
field_with_rsds = field_with_rsds.reshape(
(
int(np.sqrt(field_with_rsds.shape[0])),
int(np.sqrt(field_with_rsds.shape[0])),
-1,
)
)
return field_with_rsds
[docs]
def rsds_shift(
field: np.ndarray,
los_displacement: np.ndarray,
n_rsd_subcells: int = 4,
periodic: bool = False,
) -> np.ndarray:
"""Shift the cells of a field according to the los displacement.
Parameters
----------
field
A box of the field on which redshift-space distortions shall be applied, shape (nslices, ncoords).
los_displacement
The line-of-sight "apparent" displacement of the field, in pixel coordinates.
Equal to ``v / H(z) / cell_size``.
Positive values are towards the observer, shape ``(nslices, ncoords)``.
periodic: bool, optioanl
Whether to assume periodic boundary conditions along the line-of-sight.
n_rsd_subcells: int, optional
The number of subcells into which each cell is divided when redshift space distortions are applied. Default is 4.
Returns
-------
field_with_rsds
A box of the field, with redshift space distortions, shape (nslices, ncoords).
"""
if field.shape[0] < 2:
raise ValueError("field must have at least 2 slices")
if los_displacement.shape != field.shape:
raise ValueError(
"field must be an array with the same shape as los_displacement"
)
if not isinstance(n_rsd_subcells, int):
raise ValueError("n_rsd_subcells must be an integer")
ang_coords = np.arange(field.shape[1])
distance = np.arange(field.shape[0])
# We need to extend the grid once as we would like to interpret the displacement values
# to be associated with the *centers* of the cells, as was previously done in the C code
distance_plus = np.arange(field.shape[0] + 1)
if periodic:
# We extend the grid twice more to account for periodic boundary conditions
distance_plus_periodic = np.arange(-1, field.shape[0] + 2)
distance_grid = (distance_plus_periodic[1:] + distance_plus_periodic[:-1]) / 2
# We also extend the displacement array to have periodic boundary conditions
first_slice = los_displacement[-1, :].reshape(1, len(ang_coords))
last_slice = los_displacement[0, :].reshape(1, len(ang_coords))
los_displacement = np.concatenate(
(first_slice, los_displacement, last_slice), axis=0
)
else:
distance_grid = (distance_plus[1:] + distance_plus[:-1]) / 2
fine_field = np.repeat(field, n_rsd_subcells, axis=0) / n_rsd_subcells
# This is where we shall evaluate the subcells
distance_fine = np.linspace(
distance_plus.min(),
distance_plus.max(),
1 + n_rsd_subcells * (len(distance_plus) - 1),
)
fine_grid = (distance_fine[1:] + distance_fine[:-1]) / 2
x, y = np.meshgrid(fine_grid, ang_coords, indexing="ij")
grid = (x.flatten(), y.flatten())
fine_rsd = RegularGridInterpolator(
(distance_grid, ang_coords),
los_displacement * n_rsd_subcells,
bounds_error=False,
fill_value=None,
method="linear",
)(grid).reshape(x.shape)
fine_field = cloud_in_cell_los(fine_field, fine_rsd, periodic=periodic)
# Average the subcells back to the original grid
return np.sum(
fine_field.T.reshape(len(ang_coords), len(distance), n_rsd_subcells), axis=-1
).T
[docs]
def estimate_rsd_displacements(
classy_output: Class,
cosmo: cosmology,
redshifts: Sequence[float],
factor: float = 1.0,
) -> Sequence[units.Quantity]:
"""Estimate the rms of the redshift space distortions displacement field at given redshifts.
Parameters
----------
classy_output : classy.Class
An object containing all the information from the CLASS calculation.
cosmo: astropy.cosmology
The assumed cosmology
redshifts: list
List of the redshifts at which the rms of the displacement field is computed.
factor: float, optional
Factor to multiply the rms velocity from CLASS. Default is 1.
Returns
-------
displacements : np.array
The rms of the RSD displacement field.
"""
v_rms = (
compute_rms(classy_output=classy_output, kind="v_b", redshifts=redshifts)
* factor
)
# The multiplication by (1+z)=1/a is because CLASS returns the *proper* velocity field,
# while we work in *comoving* coordinates
return v_rms / cosmo.H(redshifts) * (1.0 + redshifts)