def cl(ell):
def integrand(a):
# Step 1: retrieve the associated comoving distance
chi = bkgrd.radial_comoving_distance(cosmo, a)
# Step 2: get the power spectrum for this combination of chi and a
k = (ell + 0.5) / np.clip(chi, 1.0)
# pk should have shape [na]
pk = power.nonlinear_matter_power(cosmo, k, a, transfer_fn,
nonlinear_fn)
# Compute the kernels for all probes
kernels = np.vstack([p.kernel(cosmo, a2z(a), ell) for p in probes])
# Define an ordering for the blocks of the signal vector
cl_index = np.array(_get_cl_ordering(probes))
# Compute all combinations of tracers
def combine_kernels(inds):
return kernels[inds[0]] * kernels[inds[1]]
# Now kernels has shape [ncls, na]
kernels = lax.map(combine_kernels, cl_index)
result = pk * kernels * bkgrd.dchioverda(cosmo, a) / np.clip(
chi**2, 1.0)
# We transpose the result just to make sure that na is first
return result.T
return simps(integrand, z2a(zmax), 1.0, 512) / const.c**2
def angular_cl(cosmo, ell, tracer_fn1, tracer_fn2, amin=0.002):
"""
Computes angular Cls for the provided trace functions
All using the Limber approximation
TODO: support generic power spectra
"""
@vmap
def integrand(a):
# Step 1: retrieve the associated comoving distance
chi = bkgrd.radial_comoving_distance(cosmo, a)
# Step 2: get the powers pectrum for this combination of chi and a
k = (ell + 0.5) / np.clip(chi, 1.)
pk = power.linear_matter_power(cosmo, k, a)
# Step 3: Get the kernels evaluated at given (ell, z)
kernel = tracer_fn1(cosmo, ell, a2z(a)) * tracer_fn2(
cosmo, ell, a2z(a))
return np.squeeze(pk * kernel * bkgrd.dchioverda(cosmo, a) / a**2)
# Integral over a.... would probably be better in log scale... idk
return simps(integrand, amin, 1., 512)
def weak_lensing_kernel(cosmo, pzs, z, ell):
"""
Returns a weak lensing kernel
"""
z = np.atleast_1d(z)
zmax = max([pz.zmax for pz in pzs])
# Retrieve comoving distance corresponding to z
chi = bkgrd.radial_comoving_distance(cosmo, z2a(z))
@vmap
def integrand(z_prime):
chi_prime = bkgrd.radial_comoving_distance(cosmo, z2a(z_prime))
# Stack the dndz of all redshift bins
dndz = np.stack([pz(z_prime) for pz in pzs], axis=0)
return dndz * np.clip(chi_prime - chi, 0) / np.clip(chi_prime, 1.0)
# Computes the radial weak lensing kernel
radial_kernel = np.squeeze(
simps(integrand, z, zmax, 256) * (1.0 + z) * chi)
# Constant term
constant_factor = 3.0 * const.H0**2 * cosmo.Omega_m / 2.0 / const.c
# Ell dependent factor
ell_factor = np.sqrt(
(ell - 1) * (ell) * (ell + 1) * (ell + 2)) / (ell + 0.5)**2
return constant_factor * ell_factor * radial_kernel
def __call__(self, z):
""" Computes the normalized n(z)
"""
if self._norm is None:
self._norm = simps(lambda t: self.pz_fn(t), 0.0,
self.config["zmax"], 256)
return self.pz_fn(z) / self._norm
def __call__(self, z):
"""
Computes the normalized n(z)
"""
if self._norm is None:
self._norm = simps(lambda t: self.pz_fn(t, **(self._params)), 0.,
self._zmax, 256)
return self.pz_fn(z, **(self._params)) / self._norm
def lensing_kernel(cosmo, z, nz):
"""
Computes the lensing kernel W_L, for a flat universe
"""
chi = bkgrd.radial_comoving_distance(cosmo, z2a(z))
def integrand(z_prime):
chi_prime = bkgrd.radial_comoving_distance(cosmo, z2a(z_prime))
return nz(z_prime) * np.clip(chi_prime - chi, 0) / (chi_prime + 1e-5)
return np.squeeze(simps(integrand, z, nz.zmax, 256))
def _halofit_parameters(cosmo, a, transfer_fn):
r""" Computes the non linear scale,
effective spectral index,
spectral curvature
"""
# Step 1: Finding the non linear scale for which sigma(R)=1
# That's our search range for the non linear scale
r = np.logspace(-3, 1, 256)
@jax.vmap
def R_nl(a):
def int_sigma(logk):
k = np.exp(logk)
y = np.outer(k, r)
pk = linear_matter_power(cosmo, k, transfer_fn=transfer_fn)
g = bkgrd.growth_factor(cosmo, np.atleast_1d(a))
return (
np.expand_dims(pk * k ** 3, axis=1)
* np.exp(-(y ** 2))
/ (2.0 * np.pi ** 2)
* g ** 2
)
sigma = simps(int_sigma, np.log(1e-4), np.log(1e4), 256)
root = interp(np.atleast_1d(1.0), sigma, r)
return root
# Compute non linear scale
k_nl = 1.0 / R_nl(np.atleast_1d(a)).squeeze()
# Step 2: Retrieve the spectral index and spectral curvature
def integrand(logk):
k = np.exp(logk)
y = np.outer(k, 1.0 / k_nl)
pk = linear_matter_power(cosmo, k, transfer_fn=transfer_fn)
g = np.expand_dims(bkgrd.growth_factor(cosmo, np.atleast_1d(a)), 0)
res = (
np.expand_dims(pk * k ** 3, axis=1)
* np.exp(-(y ** 2))
* g ** 2
/ (2.0 * np.pi ** 2)
)
dneff_dlogk = 2 * res * y ** 2
dC_dlogk = 4 * res * (y ** 2 - y ** 4)
return np.stack([dneff_dlogk, dC_dlogk], axis=1)
res = simps(integrand, np.log(1e-4), np.log(1e4), 256)
n_eff = res[0] - 3.0
C = res[0] ** 2 + res[1]
return k_nl, n_eff, C
def R_nl(a):
def int_sigma(logk):
k = np.exp(logk)
y = np.outer(k, r)
pk = linear_matter_power(cosmo, k, transfer_fn=transfer_fn)
g = bkgrd.growth_factor(cosmo, np.atleast_1d(a))
return (
np.expand_dims(pk * k ** 3, axis=1)
* np.exp(-(y ** 2))
/ (2.0 * np.pi ** 2)
* g ** 2
)
sigma = simps(int_sigma, np.log(1e-4), np.log(1e4), 256)
root = interp(np.atleast_1d(1.0), sigma, r)
return root