def test_gaussian_log_likelihood():
n_ell = 5
ell = jnp.logspace(1, 3, n_ell)
nz1 = smail_nz(1.0, 2.0, 1.0)
nz2 = smail_nz(1.0, 2.0, 0.5)
n_cls = 3
P = [probes.NumberCounts([nz1, nz2], constant_linear_bias(1.0))]
cosmo = Planck15()
mu, cov_sparse = gaussian_cl_covariance_and_mean(cosmo,
ell,
P,
sparse=True)
cov_dense = to_dense(cov_sparse)
data = 1.1 * mu
for include_logdet in (True, False):
loglike_sparse = gaussian_log_likelihood(data,
mu,
cov_sparse,
include_logdet=include_logdet)
for method in "inverse", "cholesky":
loglike_dense = gaussian_log_likelihood(
data,
mu,
cov_dense,
include_logdet=include_logdet,
inverse_method=method,
)
assert_allclose(loglike_sparse, loglike_dense, rtol=1e-6)
def test_clustering_cl():
# 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="halofit",
)
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,
)
# Define a redshift distribution
nz = smail_nz(1.0, 2.0, 1.0)
# And a bias model
bias = constant_linear_bias(1.0)
z = np.linspace(0, 5.0, 1024)
tracer_ccl = ccl.NumberCountsTracer(cosmo_ccl,
has_rsd=False,
dndz=(z, nz(z)),
bias=(z, bias(cosmo_jax, z)))
tracer_jax = probes.NumberCounts([nz], bias)
# Get an ell range for the cls
ell = np.logspace(0.1, 4)
# Compute the cls
cl_ccl = ccl.angular_cl(cosmo_ccl, tracer_ccl, tracer_ccl, ell)
cl_jax = angular_cl(cosmo_jax, ell, [tracer_jax])
assert_allclose(cl_ccl, cl_jax[0], rtol=0.5e-2)
def test_sparse_cov():
n_ell = 25
ell = jnp.logspace(1, 3, n_ell)
nz1 = smail_nz(1.0, 2.0, 1.0)
nz2 = smail_nz(1.0, 2.0, 0.5)
n_cls = 3
P = [probes.NumberCounts([nz1, nz2], constant_linear_bias(1.0))]
cl_signal = jnp.ones((n_cls, n_ell))
cl_noise = jnp.ones_like(cl_signal)
cov_dense = gaussian_cl_covariance(ell,
P,
cl_signal,
cl_noise,
sparse=False)
cov_sparse = gaussian_cl_covariance(ell,
P,
cl_signal,
cl_noise,
sparse=True)
assert cov_sparse.shape == (n_cls, n_cls, n_ell)
assert_array_equal(to_dense(cov_sparse), cov_dense)