From ced54dc969505e275d1704aef3dc0b10bc94d9f7 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Mon, 6 Jul 2026 20:34:30 +0200 Subject: [PATCH 1/5] Add tomography to the theory prediction. --- cs_util/cosmo.py | 42 ++++++++++++++++++++++++++++++------------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/cs_util/cosmo.py b/cs_util/cosmo.py index 421a5a8..ff7f9f2 100644 --- a/cs_util/cosmo.py +++ b/cs_util/cosmo.py @@ -15,9 +15,11 @@ :Authors: Martin Kilbinger Axel Guinot Cail Daley - + Sacha Guerrini """ +import itertools + import camb import numpy as np @@ -641,7 +643,7 @@ def get_theo_c_ell( z : array Redshifts for n(z) distribution nz : array - n(z) redshift distribution + n(z) redshift distribution. If nz.shape[1] > 1, assumes multiple tomographic bins. backend : str, default="ccl" Backend to use: "ccl" or "camb" cosmo : ccl.Cosmology, optional @@ -677,12 +679,21 @@ def get_theo_c_ell( wa=wa, ) + n_tomo_bins = nz.shape[1] if len(nz.shape) > 1 else 1 + tomo_bin_pairs = list(itertools.combinations_with_replacement(range(1, n_tomo_bins + 1), 2)) + cl = {} + if backend == "ccl": + tracers = {} + # Create lensing tracer - lens = ccl.WeakLensingTracer(cosmo, dndz=(z, nz)) + for bin_key in range(1, n_tomo_bins + 1): + tracers[f"W{bin_key}"]= ccl.WeakLensingTracer(cosmo, dndz=(z, nz[:, bin_key])) - # Calculate power spectrum - cl = ccl.angular_cl(cosmo, lens, lens, ell) + for bin_key1, bin_key2 in tomo_bin_pairs: + cl[f"W{bin_key1}xW{bin_key2}"] = ccl.angular_cl( + cosmo, tracers[f"W{bin_key1}"], tracers[f"W{bin_key2}"], ell + ) elif backend == "camb": # Convert CCL cosmology to CAMB parameters @@ -705,19 +716,26 @@ def get_theo_c_ell( # Set up lensing source window pars.min_l = ell.min() pars.set_for_lmax(ell.max()) - pars.SourceWindows = [ - camb.sources.SplinedSourceWindow(z=z, W=nz, source_type="lensing") - ] + if len(nz.shape) == 1: + pars.SourceWindows = [ + camb.sources.SplinedSourceWindow(z=z, W=nz, source_type="lensing") + ] + else: + pars.SourceWindows = [ + camb.sources.SplinedSourceWindow(z=z, W=nz[:, i], source_type="lensing") + for i in range(nz.shape[1]) + ] # Calculate power spectrum results = camb.get_results(pars) theory_cls = results.get_source_cls_dict(lmax=ell.max(), raw_cl=True) - cl_full = theory_cls["W1xW1"] # Interpolate to match input ell array # CAMB returns C_ell for ell = 0, 1, 2, ..., lmax - ell_camb = np.arange(len(cl_full)) - cl = np.interp(ell, ell_camb, cl_full) + ell_camb = np.arange(len(theory_cls["W1xW1"])) + for bin_key1, bin_key2 in tomo_bin_pairs: + cl_full = theory_cls[f"W{bin_key1}xW{bin_key2}"] + cl[f"W{bin_key1}xW{bin_key2}"] = np.interp(ell, ell_camb, cl_full) else: raise ValueError(f"Unknown backend: {backend}. Must be 'ccl' or 'camb'") @@ -821,4 +839,4 @@ def get_theo_xi( cl = get_theo_c_ell(ell, z, nz, backend=backend, cosmo=cosmo) # Convert to xi - return c_ell_to_xi(cosmo, theta, ell, cl) + return {k: c_ell_to_xi(cosmo, theta, ell, v) for k, v in cl.items()} From 28cf0765f535ccfdce56b8717a1048e85074c860 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Mon, 6 Jul 2026 20:37:21 +0200 Subject: [PATCH 2/5] Update docstrings --- cs_util/cosmo.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cs_util/cosmo.py b/cs_util/cosmo.py index ff7f9f2..9e2464b 100644 --- a/cs_util/cosmo.py +++ b/cs_util/cosmo.py @@ -665,8 +665,8 @@ def get_theo_c_ell( Returns ------- - cl : array - Angular power spectrum + cl : dict + Angular power spectrum in dictionnary indexed by the tomographic bin keys. """ if cosmo is None: cosmo = get_cosmo( @@ -799,7 +799,7 @@ def get_theo_xi( z : array Redshift array nz : array - n(z) redshift distribution + n(z) redshift distribution. If nz.shape[1] > 1, assumes multiple tomographic bins. Omega_m : float, default=None Matter density parameter (defaults to Planck 2018) h : float, default=None @@ -823,8 +823,8 @@ def get_theo_xi( Returns ------- - xip, xim : arrays - Theoretical xi+ and xi- correlation functions + dict + Theoretical xi+ and xi- correlation functions per tomographic bin combinations """ # Create ell array for C_ell calculation ell = np.geomspace(ell_min, ell_max, n_ell) From bbe676d9d5038de641b92f64b7cce36d50a6a5f6 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Mon, 6 Jul 2026 20:57:36 +0200 Subject: [PATCH 3/5] (fix) Update the tests to the existing API. To do: add minimal unit tests for the tomographic case. --- cs_util/cosmo.py | 5 ++++- cs_util/tests/test_cosmology.py | 28 ++++++++++++++-------------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/cs_util/cosmo.py b/cs_util/cosmo.py index 9e2464b..1c9e571 100644 --- a/cs_util/cosmo.py +++ b/cs_util/cosmo.py @@ -680,6 +680,9 @@ def get_theo_c_ell( ) n_tomo_bins = nz.shape[1] if len(nz.shape) > 1 else 1 + # Add a new axis to nz if it's a single tomographic bin for consistent indexing + if n_tomo_bins == 1 and len(nz.shape) == 1: + nz = nz[:, np.newaxis] tomo_bin_pairs = list(itertools.combinations_with_replacement(range(1, n_tomo_bins + 1), 2)) cl = {} @@ -688,7 +691,7 @@ def get_theo_c_ell( # Create lensing tracer for bin_key in range(1, n_tomo_bins + 1): - tracers[f"W{bin_key}"]= ccl.WeakLensingTracer(cosmo, dndz=(z, nz[:, bin_key])) + tracers[f"W{bin_key}"]= ccl.WeakLensingTracer(cosmo, dndz=(z, nz[:, bin_key - 1])) for bin_key1, bin_key2 in tomo_bin_pairs: cl[f"W{bin_key1}xW{bin_key2}"] = ccl.angular_cl( diff --git a/cs_util/tests/test_cosmology.py b/cs_util/tests/test_cosmology.py index 019f055..af47243 100644 --- a/cs_util/tests/test_cosmology.py +++ b/cs_util/tests/test_cosmology.py @@ -275,7 +275,7 @@ def test_ccl_backend_basic(self, fast_ell_array, fast_redshift_data): z, nz = fast_redshift_data cosmo = cosmology.get_cosmo() - cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl", cosmo=cosmo) + cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl", cosmo=cosmo)["W1xW1"] # Check output shape and positivity assert len(cl) == len(fast_ell_array) @@ -286,7 +286,7 @@ def test_ccl_backend_with_params(self, fast_ell_array, fast_redshift_data): z, nz = fast_redshift_data cl = cosmology.get_theo_c_ell( fast_ell_array, z, nz, backend="ccl", **TEST_COSMOLOGY - ) + )["W1xW1"] # Check output shape and positivity assert len(cl) == len(fast_ell_array) @@ -295,7 +295,7 @@ def test_ccl_backend_with_params(self, fast_ell_array, fast_redshift_data): def test_ccl_backend_default_params(self, fast_ell_array, fast_redshift_data): """Test C_ell calculation with CCL backend using default parameters.""" z, nz = fast_redshift_data - cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl") + cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl")["W1xW1"] # Check output shape and positivity assert len(cl) == len(fast_ell_array) @@ -305,7 +305,7 @@ def test_camb_backend(self, fast_ell_array, fast_redshift_data): """Test C_ell calculation with CAMB backend.""" z, nz = fast_redshift_data try: - cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="camb") + cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="camb")["W1xW1"] # Check output shape and positivity assert len(cl) == len(fast_ell_array) @@ -323,7 +323,7 @@ def test_small_ell_array(self, fast_redshift_data): """Test behavior with small ell array.""" z, nz = fast_redshift_data small_ell = np.array([10, 100]) - cl = cosmology.get_theo_c_ell(small_ell, z, nz, backend="ccl") + cl = cosmology.get_theo_c_ell(small_ell, z, nz, backend="ccl")["W1xW1"] assert len(cl) == 2 assert np.all(cl > 0) @@ -334,8 +334,8 @@ def test_backend_consistency_fast(self, fast_ell_array, fast_redshift_data): z, nz = fast_redshift_data try: # Calculate with both backends using fast parameters - cl_ccl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl") - cl_camb = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="camb") + cl_ccl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl")["W1xW1"] + cl_camb = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="camb")["W1xW1"] # Check that both have same shape assert len(cl_ccl) == len(cl_camb) == len(fast_ell_array) @@ -374,7 +374,7 @@ def test_backwards_compatibility(self, fast_theta_array, fast_redshift_data): ell_max=FAST_ELL_MAX, n_ell=FAST_N_ELL, backend="ccl", - ) + )["W1xW1"] # Check output shapes assert len(xip) == len(fast_theta_array) @@ -388,7 +388,7 @@ def test_default_parameters(self, fast_theta_array, fast_redshift_data): z, nz = fast_redshift_data xip, xim = cosmology.get_theo_xi( fast_theta_array, z, nz, ell_min=10, ell_max=FAST_ELL_MAX, n_ell=FAST_N_ELL - ) + )["W1xW1"] # Check output shapes assert len(xip) == len(fast_theta_array) @@ -402,12 +402,12 @@ def test_different_ell_ranges(self, fast_theta_array, fast_redshift_data): # Coarse integration (lower ell_max) xip1, xim1 = cosmology.get_theo_xi( fast_theta_array, z, nz, ell_min=10, ell_max=1000, n_ell=50 - ) + )["W1xW1"] # Fine integration (higher ell_max) xip2, xim2 = cosmology.get_theo_xi( fast_theta_array, z, nz, ell_min=10, ell_max=FAST_ELL_MAX, n_ell=FAST_N_ELL - ) + )["W1xW1"] # Scale-dependent tolerance: # Xi+ converges well at all scales (~4% max differences) @@ -451,7 +451,7 @@ def test_end_to_end_pipeline_fast(self, fast_redshift_data): # Calculate C_ell with fast parameters ell = np.logspace(1, 3, 15) # Reduced from 30 points - cl = cosmology.get_theo_c_ell(ell, z, nz, cosmo=cosmo) + cl = cosmology.get_theo_c_ell(ell, z, nz, cosmo=cosmo)["W1xW1"] # Calculate xi theta = np.array([5.0, 10.0, 20.0]) @@ -473,7 +473,7 @@ def test_realistic_precision_pipeline(self, realistic_redshift_data): # Calculate C_ell with realistic precision ell = np.logspace(1, np.log10(REALISTIC_ELL_MAX), REALISTIC_N_ELL) - cl = cosmology.get_theo_c_ell(ell, z, nz, cosmo=cosmo) + cl = cosmology.get_theo_c_ell(ell, z, nz, cosmo=cosmo)["W1xW1"] # Test xi calculation with realistic theta range theta = np.geomspace(1, 300, 20) # 1 to 300 arcmin, 20 points @@ -485,7 +485,7 @@ def test_realistic_precision_pipeline(self, realistic_redshift_data): ell_max=REALISTIC_ELL_MAX, n_ell=REALISTIC_N_ELL, **TEST_COSMOLOGY, - ) + )["W1xW1"] # Comprehensive checks for realistic test assert len(cl) == len(ell) From 8e63fe625caab3ffd48ef18fd2533cb71979b5ac Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Tue, 7 Jul 2026 09:28:49 +0200 Subject: [PATCH 4/5] Add unit tests for the tomographic theory prediction --- cs_util/tests/test_cosmology.py | 248 ++++++++++++++++++++++++++++++++ 1 file changed, 248 insertions(+) diff --git a/cs_util/tests/test_cosmology.py b/cs_util/tests/test_cosmology.py index af47243..9538974 100644 --- a/cs_util/tests/test_cosmology.py +++ b/cs_util/tests/test_cosmology.py @@ -251,6 +251,16 @@ def fast_redshift_data(): nz /= np.trapezoid(nz, z) # Normalize return z, nz +@pytest.fixture +def fast_two_bins_data(): + """Fast two bin redshift distribution for more tests.""" + z = np.linspace(0.1, 2.0, FAST_Z_POINTS) + nz = np.exp(-(((z - 0.8) / 0.2) ** 2)) # Gaussian around z=0.8 + nz /= np.trapezoid(nz, z) # Normalize + nz2 = np.exp(-(((z - 1.2) / 0.2) ** 2)) # Gaussian around z=1.2 + nz2 /= np.trapezoid(nz2, z) + nz = np.vstack([nz, nz2]) # Stack to create two bins + return z, nz.T @pytest.fixture def fast_ell_array(): @@ -266,6 +276,17 @@ def realistic_redshift_data(): nz /= np.trapezoid(nz, z) return z, nz +@pytest.fixture +def realistic_two_bins_data(): + """Realistic two-bin redshift distribution for integration tests.""" + z = np.linspace(0.1, 2.0, 50) + nz1 = np.exp(-(((z - 0.8) / 0.2) ** 2)) + nz1 /= np.trapezoid(nz1, z) + nz2 = np.exp(-(((z - 1.2) / 0.2) ** 2)) + nz2 /= np.trapezoid(nz2, z) + nz = np.vstack([nz1, nz2]) # Stack to create two bins + return z, nz.T + class TestGetTheoCell: """Test get_theo_c_ell function.""" @@ -281,6 +302,22 @@ def test_ccl_backend_basic(self, fast_ell_array, fast_redshift_data): assert len(cl) == len(fast_ell_array) assert np.all(cl > 0) + def test_ccl_backend_tomo(self, fast_ell_array, fast_two_bins_data): + """Test C_ell calculation with CCL backend for two tomographic bins.""" + z, nz = fast_two_bins_data + cosmo = cosmology.get_cosmo() + + cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl", cosmo=cosmo) + + # Check that output is a dictionary with correct keys + assert isinstance(cl, dict) + assert "W1xW1" in cl and "W1xW2" in cl and "W2xW2" in cl + + # Check output shapes and positivity + for key in cl: + assert len(cl[key]) == len(fast_ell_array) + assert np.all(cl[key] > 0) + def test_ccl_backend_with_params(self, fast_ell_array, fast_redshift_data): """Test C_ell calculation with CCL backend using individual parameters.""" z, nz = fast_redshift_data @@ -292,6 +329,18 @@ def test_ccl_backend_with_params(self, fast_ell_array, fast_redshift_data): assert len(cl) == len(fast_ell_array) assert np.all(cl > 0) + def test_ccl_backend_with_params_tomo(self, fast_ell_array, fast_two_bins_data): + """Test C_ell calculation with CCL backend using individual parameters for tomographic bins.""" + z, nz = fast_two_bins_data + cl = cosmology.get_theo_c_ell( + fast_ell_array, z, nz, backend="ccl", **TEST_COSMOLOGY + ) + + # Check output shape and positivity + for key in cl: + assert len(cl[key]) == len(fast_ell_array) + assert np.all(cl[key] > 0) + def test_ccl_backend_default_params(self, fast_ell_array, fast_redshift_data): """Test C_ell calculation with CCL backend using default parameters.""" z, nz = fast_redshift_data @@ -301,6 +350,16 @@ def test_ccl_backend_default_params(self, fast_ell_array, fast_redshift_data): assert len(cl) == len(fast_ell_array) assert np.all(cl > 0) + def test_ccl_backend_default_params_tomo(self, fast_ell_array, fast_two_bins_data): + """Test C_ell calculation with CCL backend using default parameters for tomographic bins.""" + z, nz = fast_two_bins_data + cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl") + + # Check output shape and positivity + for key in cl: + assert len(cl[key]) == len(fast_ell_array) + assert np.all(cl[key] > 0) + def test_camb_backend(self, fast_ell_array, fast_redshift_data): """Test C_ell calculation with CAMB backend.""" z, nz = fast_redshift_data @@ -313,6 +372,21 @@ def test_camb_backend(self, fast_ell_array, fast_redshift_data): except ImportError: pytest.skip("CAMB not available") + def test_camb_backend_tomo(self, fast_ell_array, fast_two_bins_data): + """Test C_ell calculation with CAMB backend for two tomographic bins.""" + z, nz = fast_two_bins_data + try: + cl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="camb") + + assert isinstance(cl, dict) + assert "W1xW1" in cl and "W1xW2" in cl and "W2xW2" in cl + # Check output shape and positivity + for key in cl: + assert len(cl[key]) == len(fast_ell_array) + assert np.all(cl[key] > 0) + except ImportError: + pytest.skip("CAMB not available") + def test_invalid_backend(self, fast_ell_array, fast_redshift_data): """Test error with invalid backend.""" z, nz = fast_redshift_data @@ -352,6 +426,33 @@ def test_backend_consistency_fast(self, fast_ell_array, fast_redshift_data): except ImportError: pytest.skip("CAMB not available") + @pytest.mark.slow + def test_backend_consistency_fast_tomo(self, fast_ell_array, fast_two_bins_data): + """Test that CCL and CAMB backends give consistent results (fast version) for tomographic bins.""" + z, nz = fast_two_bins_data + try: + # Calculate with both backends using fast parameters + cl_ccl = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="ccl") + cl_camb = cosmology.get_theo_c_ell(fast_ell_array, z, nz, backend="camb") + + # Check that both backend have the same keys + assert cl_ccl.keys() == cl_camb.keys() + + # Check that both have same shape + for key in cl_ccl: + assert len(cl_ccl[key]) == len(cl_camb[key]) == len(fast_ell_array) + + # Check relative differences between cosmology codes + # CCL and CAMB should agree well since they compute the same physics + rel_diff = np.abs(cl_camb[key] - cl_ccl[key]) / cl_ccl[key] + max_rel_diff = rel_diff.max() + + assert max_rel_diff < 0.08, ( + f"CCL and CAMB backends differ by >8%: {max_rel_diff:.1%}" + ) + except ImportError: + pytest.skip("CAMB not available") + @pytest.fixture def fast_theta_array(): @@ -383,6 +484,32 @@ def test_backwards_compatibility(self, fast_theta_array, fast_redshift_data): # Check that xi+ is positive for typical scales assert np.any(xip > 0) + def test_backwards_compatibility_tomo(self, fast_theta_array, fast_two_bins_data): + """Test backwards compatible parameter interface with fast parameters.""" + z, nz = fast_two_bins_data + xis = cosmology.get_theo_xi( + fast_theta_array, + z, + nz, + **TEST_COSMOLOGY, + ell_min=10, + ell_max=FAST_ELL_MAX, + n_ell=FAST_N_ELL, + backend="ccl", + ) + + assert isinstance(xis, dict) + assert "W1xW1" in xis and "W1xW2" in xis and "W2xW2" in xis + + for key in xis: + xip, xim = xis[key] + # Check output shapes + assert len(xip) == len(fast_theta_array) + assert len(xim) == len(fast_theta_array) + + # Check that xi+ is positive for typical scales + assert np.any(xip > 0) + def test_default_parameters(self, fast_theta_array, fast_redshift_data): """Test with default cosmological parameters using fast settings.""" z, nz = fast_redshift_data @@ -394,6 +521,22 @@ def test_default_parameters(self, fast_theta_array, fast_redshift_data): assert len(xip) == len(fast_theta_array) assert len(xim) == len(fast_theta_array) + def test_default_parameters_tomo(self, fast_theta_array, fast_two_bins_data): + """Test with default cosmological parameters using fast settings.""" + z, nz = fast_two_bins_data + xis = cosmology.get_theo_xi( + fast_theta_array, z, nz, ell_min=10, ell_max=FAST_ELL_MAX, n_ell=FAST_N_ELL + ) + + assert isinstance(xis, dict) + assert "W1xW1" in xis and "W1xW2" in xis and "W2xW2" in xis + + for key in xis: + xip, xim = xis[key] + # Check output shapes + assert len(xip) == len(fast_theta_array) + assert len(xim) == len(fast_theta_array) + @pytest.mark.slow def test_different_ell_ranges(self, fast_theta_array, fast_redshift_data): """Test that different ell ranges give consistent results.""" @@ -434,6 +577,37 @@ def test_different_ell_ranges(self, fast_theta_array, fast_redshift_data): xim1[large_scale_mask], xim2[large_scale_mask], rtol=0.05 ) # 5% tolerance for xi- at large scales + @pytest.mark.slow + def test_different_ell_ranges_tomo(self, fast_theta_array, fast_two_bins_data): + """Test that different ell ranges give consistent results.""" + z, nz = fast_two_bins_data + + # Coarse integration (lower ell_max) + xis1 = cosmology.get_theo_xi( + fast_theta_array, z, nz, ell_min=10, ell_max=1000, n_ell=50 + ) + + # Fine integration (higher ell_max) + xis2 = cosmology.get_theo_xi( + fast_theta_array, z, nz, ell_min=10, ell_max=FAST_ELL_MAX, n_ell=FAST_N_ELL + ) + + # Check that all keys are consistent + for key in xis1: + assert key in xis2 + xip1, xim1 = xis1[key] + xip2, xim2 = xis2[key] + + npt.assert_allclose(xip1, xip2, rtol=0.06) + + small_scale_mask = fast_theta_array <= 10.0 # arcmin + large_scale_mask = fast_theta_array > 10.0 # arcmin + + if np.any(small_scale_mask): + npt.assert_allclose(xim1[small_scale_mask], xim2[small_scale_mask], rtol=0.18) + + if np.any(large_scale_mask): + npt.assert_allclose(xim1[large_scale_mask], xim2[large_scale_mask], rtol=0.18) class TestCosmologyIntegration: """Integration tests for cosmology functions.""" @@ -463,6 +637,35 @@ def test_end_to_end_pipeline_fast(self, fast_redshift_data): assert len(xim) == len(theta) assert np.all(cl > 0) + def test_end_to_end_pipeline_fast_tomo(self, fast_two_bins_data): + """Test complete pipeline from parameters to xi with tomography (fast version)""" + z, nz = fast_two_bins_data + + # Create cosmology using TEST_CAMB parameters + camb_subset = { + k: v + for k, v in TEST_CAMB.items() + if k in ["H0", "ombh2", "omch2", "ns", "sigma8"] + } + cosmo = cosmology.get_cosmo(camb_params=camb_subset) + + # Calculate C_ell with fast parameters + ell = np.logspace(1, 3, 15) # Reduced from 30 points + cl = cosmology.get_theo_c_ell(ell, z, nz, cosmo=cosmo) + + # Calculate xi + theta = np.array([5.0, 10.0, 20.0]) + xis = {k: cosmology.c_ell_to_xi(cosmo, theta, ell, v) for k, v in cl.items()} + + for key in cl: + xip, xim = xis[key] + cl_ = cl[key] + # Sanity checks + assert len(cl_) == len(ell) + assert len(xip) == len(theta) + assert len(xim) == len(theta) + assert np.all(cl_ > 0) + @pytest.mark.slow def test_realistic_precision_pipeline(self, realistic_redshift_data): """Test complete pipeline with realistic precision for production use.""" @@ -501,6 +704,51 @@ def test_realistic_precision_pipeline(self, realistic_redshift_data): max_xip = np.max(np.abs(xip)) assert 1e-6 < max_xip < 1e-2, f"xi+ amplitude unrealistic: {max_xip}" + @pytest.mark.slow + def test_realistic_precision_pipeline_tomo(self, realistic_two_bins_data): + """Test complete pipeline with realistic precision for production use.""" + z, nz = realistic_two_bins_data + + # Create cosmology with realistic parameters + cosmo = cosmology.get_cosmo(**TEST_COSMOLOGY) + + # Calculate C_ell with realistic precision + ell = np.logspace(1, np.log10(REALISTIC_ELL_MAX), REALISTIC_N_ELL) + cl = cosmology.get_theo_c_ell(ell, z, nz, cosmo=cosmo) + + # Test xi calculation with realistic theta range + theta = np.geomspace(1, 300, 20) # 1 to 300 arcmin, 20 points + xis = cosmology.get_theo_xi( + theta, + z, + nz, + ell_min=10, + ell_max=REALISTIC_ELL_MAX, + n_ell=REALISTIC_N_ELL, + **TEST_COSMOLOGY, + ) + + for key in cl: + assert key in xis + + xip, xim = xis[key] + cl_ = cl[key] + + # Comprehensive checks for realistic test + assert len(cl_) == len(ell) + assert len(xip) == len(theta) + assert len(xim) == len(theta) + assert np.all(cl_ > 0) + + # Check that xi+ is positive at small scales (expected for cosmic shear) + small_scale_mask = theta < 10 # arcmin + assert np.any(xip[small_scale_mask] > 0) + + # Check that xi+ amplitude is reasonable (order of magnitude check) + max_xip = np.max(np.abs(xip)) + assert 1e-6 < max_xip < 1e-2, f"xi+ amplitude unrealistic: {max_xip}" + + def test_parameter_format_consistency(self): """Test that different parameter formats give consistent results.""" # Individual parameters From e10fd31c6ae7f048bcbd07f19bd73bfec3a35b30 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Tue, 7 Jul 2026 10:40:51 +0200 Subject: [PATCH 5/5] Account for comments from Cail --- cs_util/cosmo.py | 30 +++++++++++++----------------- cs_util/tests/test_cosmology.py | 2 +- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/cs_util/cosmo.py b/cs_util/cosmo.py index 1c9e571..501207a 100644 --- a/cs_util/cosmo.py +++ b/cs_util/cosmo.py @@ -643,7 +643,7 @@ def get_theo_c_ell( z : array Redshifts for n(z) distribution nz : array - n(z) redshift distribution. If nz.shape[1] > 1, assumes multiple tomographic bins. + n(z) redshift distribution. If nz.shape[1] > 1, assumes multiple tomographic bins. Input shape (n_z, n_tomo_bins). backend : str, default="ccl" Backend to use: "ccl" or "camb" cosmo : ccl.Cosmology, optional @@ -666,7 +666,7 @@ def get_theo_c_ell( Returns ------- cl : dict - Angular power spectrum in dictionnary indexed by the tomographic bin keys. + Angular power spectrum in dictionary indexed by the tomographic bin keys. """ if cosmo is None: cosmo = get_cosmo( @@ -679,6 +679,7 @@ def get_theo_c_ell( wa=wa, ) + assert nz.shape[0] == len(z), f"nz first axis ({nz.shape[0]}) must match z ({len(z)})" n_tomo_bins = nz.shape[1] if len(nz.shape) > 1 else 1 # Add a new axis to nz if it's a single tomographic bin for consistent indexing if n_tomo_bins == 1 and len(nz.shape) == 1: @@ -687,11 +688,10 @@ def get_theo_c_ell( cl = {} if backend == "ccl": - tracers = {} - - # Create lensing tracer - for bin_key in range(1, n_tomo_bins + 1): - tracers[f"W{bin_key}"]= ccl.WeakLensingTracer(cosmo, dndz=(z, nz[:, bin_key - 1])) + tracers = { + f"W{bin_key}": ccl.WeakLensingTracer(cosmo, dndz=(z, nz[:, bin_key - 1])) + for bin_key in range(1, n_tomo_bins + 1) + } for bin_key1, bin_key2 in tomo_bin_pairs: cl[f"W{bin_key1}xW{bin_key2}"] = ccl.angular_cl( @@ -719,15 +719,11 @@ def get_theo_c_ell( # Set up lensing source window pars.min_l = ell.min() pars.set_for_lmax(ell.max()) - if len(nz.shape) == 1: - pars.SourceWindows = [ - camb.sources.SplinedSourceWindow(z=z, W=nz, source_type="lensing") - ] - else: - pars.SourceWindows = [ - camb.sources.SplinedSourceWindow(z=z, W=nz[:, i], source_type="lensing") - for i in range(nz.shape[1]) - ] + + pars.SourceWindows = [ + camb.sources.SplinedSourceWindow(z=z, W=nz[:, i], source_type="lensing") + for i in range(nz.shape[1]) + ] # Calculate power spectrum results = camb.get_results(pars) @@ -802,7 +798,7 @@ def get_theo_xi( z : array Redshift array nz : array - n(z) redshift distribution. If nz.shape[1] > 1, assumes multiple tomographic bins. + n(z) redshift distribution. If nz.shape[1] > 1, assumes multiple tomographic bins. Input shape: (n_z, n_tomo_bins) Omega_m : float, default=None Matter density parameter (defaults to Planck 2018) h : float, default=None diff --git a/cs_util/tests/test_cosmology.py b/cs_util/tests/test_cosmology.py index 9538974..7116037 100644 --- a/cs_util/tests/test_cosmology.py +++ b/cs_util/tests/test_cosmology.py @@ -607,7 +607,7 @@ def test_different_ell_ranges_tomo(self, fast_theta_array, fast_two_bins_data): npt.assert_allclose(xim1[small_scale_mask], xim2[small_scale_mask], rtol=0.18) if np.any(large_scale_mask): - npt.assert_allclose(xim1[large_scale_mask], xim2[large_scale_mask], rtol=0.18) + npt.assert_allclose(xim1[large_scale_mask], xim2[large_scale_mask], rtol=0.05) class TestCosmologyIntegration: """Integration tests for cosmology functions."""