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 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
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 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)
def test_distances_flat():
# We first define equivalent CCL and jax_cosmo cosmologies
cosmo_ccl = ccl.Cosmology(
Omega_c=0.3,
Omega_b=0.05,
h=0.7,
sigma8=0.8,
n_s=0.96,
Neff=0,
transfer_function="eisenstein_hu",
matter_power_spectrum="linear",
)
cosmo_jax = Cosmology(
Omega_c=0.3,
Omega_b=0.05,
h=0.7,
sigma8=0.8,
n_s=0.96,
Omega_k=0.0,
w0=-1.0,
wa=0.0,
)
# Test array of scale factors
a = np.linspace(0.01, 1.0)
chi_ccl = ccl.comoving_radial_distance(cosmo_ccl, a)
chi_jax = bkgrd.radial_comoving_distance(cosmo_jax, a) / cosmo_jax.h
assert_allclose(chi_ccl, chi_jax, rtol=0.5e-2)
chi_ccl = ccl.comoving_angular_distance(cosmo_ccl, a)
chi_jax = bkgrd.transverse_comoving_distance(cosmo_jax, a) / cosmo_jax.h
assert_allclose(chi_ccl, chi_jax, rtol=0.5e-2)
chi_ccl = ccl.angular_diameter_distance(cosmo_ccl, a)
chi_jax = bkgrd.angular_diameter_distance(cosmo_jax, a) / cosmo_jax.h
assert_allclose(chi_ccl, chi_jax, rtol=0.5e-2)
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)