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 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 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 nla_kernel(cosmo, pzs, bias, z, ell):
"""
Computes the NLA IA kernel
"""
# stack the dndz of all redshift bins
dndz = np.stack([pz(z) for pz in pzs], axis=0)
# Compute radial NLA kernel: same as clustering
if isinstance(bias, list):
# This is to handle the case where we get a bin-dependent bias
b = np.stack([b(cosmo, z) for b in bias], axis=0)
else:
b = bias(cosmo, z)
radial_kernel = dndz * b * bkgrd.H(cosmo, z2a(z))
# Apply common A_IA normalization to the kernel
# Joachimi et al. (2011), arXiv: 1008.3491, Eq. 6.
radial_kernel *= (-(5e-14 * const.rhocrit) * cosmo.Omega_m /
bkgrd.growth_factor(cosmo, z2a(z)))
# Constant factor
constant_factor = 1.0
# 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 density_kernel(cosmo, pzs, bias, z, ell):
"""
Computes the number counts density kernel
"""
# stack the dndz of all redshift bins
dndz = np.stack([pz(z) for pz in pzs], axis=0)
# Compute radial NLA kernel: same as clustering
if isinstance(bias, list):
# This is to handle the case where we get a bin-dependent bias
b = np.stack([b(cosmo, z) for b in bias], axis=0)
else:
b = bias(cosmo, z)
radial_kernel = dndz * b * bkgrd.H(cosmo, z2a(z))
# Normalization,
constant_factor = 1.0
# Ell dependent factor
ell_factor = 1.0
return constant_factor * ell_factor * radial_kernel
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)
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)
def __call__(self, cosmo, z):
b = self.params[0]
return b / bkgrd.growth_factor(cosmo, z2a(z))