diff --git a/cloelib/cosmology/camb_cosmology.py b/cloelib/cosmology/camb_cosmology.py index 8d54bf4c..8a6849ca 100644 --- a/cloelib/cosmology/camb_cosmology.py +++ b/cloelib/cosmology/camb_cosmology.py @@ -297,6 +297,11 @@ def rdrag(self) -> float: """Sound horizon radius at last scattering in Mpc.""" return self.results.get_derived_params()["rdrag"] + @property + def z_star(self) -> float: + """Redshift of photon decoupling.""" + return self.results.get_derived_params()["zstar"] + class CAMBLinearPerturbations: """A wrapper for CAMB linear perturbation calculations.""" diff --git a/cloelib/cosmology/class_cosmology.py b/cloelib/cosmology/class_cosmology.py index 232d0a28..d38e0c66 100644 --- a/cloelib/cosmology/class_cosmology.py +++ b/cloelib/cosmology/class_cosmology.py @@ -281,6 +281,11 @@ def rdrag(self) -> float: """Sound horizon radius at last scattering in Mpc.""" return self.results.rs_drag() + @property + def z_star(self) -> float: + """Redshift of photon decoupling.""" + return self.results.get_current_derived_parameters(["z_star"])["z_star"] + class CLASSLinearPerturbations: """Class for perturbations cosmology using CLASS, inheriting from Perturbations parent class.""" diff --git a/cloelib/cosmology/cosmology.py b/cloelib/cosmology/cosmology.py index 595b7d21..35fd00eb 100644 --- a/cloelib/cosmology/cosmology.py +++ b/cloelib/cosmology/cosmology.py @@ -130,6 +130,11 @@ def rdrag(self) -> float: """Sound horizon radius at last scattering in Mpc.""" ... + @property + def z_star(self) -> float: + """Redshift of photon decoupling.""" + ... + @runtime_checkable class Perturbations(Protocol): diff --git a/cloelib/cosmology/derived_cosmology.py b/cloelib/cosmology/derived_cosmology.py index b806c507..a1913267 100644 --- a/cloelib/cosmology/derived_cosmology.py +++ b/cloelib/cosmology/derived_cosmology.py @@ -79,3 +79,21 @@ def rdrag_fitting_function(background, neff=3.046): / (omega_cb**0.2436 * omega_b**0.128876 * (1 + (neff - 3.046) / 30.6)) ) return r_d + + +def z_star_fitting_function(background): + r"""Compute the redshift of photon decoupling. + + Assumes a cosmology-independent z_star. + + Parameters + ---------- + background: Background + Background class containing cosmology + + Returns + ------- + z_star: float + redshift of photon decoupling. + """ + return 1090.0 diff --git a/cloelib/cosmology/jax_cosmology.py b/cloelib/cosmology/jax_cosmology.py index e718f892..df0241e2 100644 --- a/cloelib/cosmology/jax_cosmology.py +++ b/cloelib/cosmology/jax_cosmology.py @@ -8,7 +8,10 @@ # cloelib imports from cloelib.auxiliary.units import SPEED_OF_LIGHT from cloelib.cosmology.cosmology import Background -from cloelib.cosmology.derived_cosmology import rdrag_fitting_function +from cloelib.cosmology.derived_cosmology import ( + rdrag_fitting_function, + z_star_fitting_function, +) # General imports import jax.numpy as jnp @@ -347,6 +350,11 @@ def rdrag(self) -> float: """Sound horizon radius at last scattering.""" return rdrag_fitting_function(self) + @property + def z_star(self) -> float: + """Redshift of photon decoupling.""" + return z_star_fitting_function(self) + class JAXLinearPerturbations: """A wrapper for JAX linear perturbation calculations.""" diff --git a/cloelib/observables/cmb.py b/cloelib/observables/cmb.py new file mode 100644 index 00000000..bb9978df --- /dev/null +++ b/cloelib/observables/cmb.py @@ -0,0 +1,84 @@ +""" +Module implementing CMB lensing. + +This class is compatible with the Tracer protocol. +""" + +# cloelib imports +from cloelib.auxiliary.units import SPEED_OF_LIGHT +from cloelib.cosmology.cosmology import Perturbations + +# General imports +import jax.numpy as np # type: ignore + + +# UNITS +c_0 = SPEED_OF_LIGHT / 1000 # Convert to km/s + + +class CMBLensingTracer: + """Class for the kernel for CMB Lensing convergence.""" + + def __init__( + self, + perturbations: Perturbations, + z: np.ndarray, + ): + r""" + Initialize the class instance. + + Parameters + ---------- + perturbations : object + An object from NonLinearPerturbations class + z : np.ndarray + A 1-dimensional array used to perform line-of-sight integration. + """ + if 0.0 in z: + raise ValueError( + "One of the z array elements is equal to zero, breaking Limber integration." + ) + self.perturbations = perturbations + self.background = self.perturbations.background + self.z = z + self.n_z_bins = 1 + # This is to add the necessary prefactor to shear + self.prefact_toggle = 0 + + def get_window(self, z): + r"""Compute the Window. + + Computes CMB lensing window function + + .. math:: + W^{\kappa}(\ell, z, k) = + \frac{3}{2}\left ( \frac{H_0}{c}\right )^2 + \Omega_{{\rm m},0} (1 + z) + f_K\left[\tilde{r}(z)\right] + \frac{f_K\left[\tilde{r}(z_*) - \tilde{r}(z)\right]} + {f_K\left[\tilde{r}(z_*)\right]}\\ + + Parameters + ---------- + z: float + Redshift at which window kernel is being evaluated + + Returns + ------- + window: np.ndarray + """ + Omega_m0 = self.background.Omega_m(0.0) + factor = ( + 3 + / 2 + * (self.background.H0 / c_0) ** 2 + * Omega_m0 + * (1 + z) + * self.background.comoving_distance(z) + ) + rz = self.background.comoving_distance(z) + z_star = self.background.z_star + rz_star = self.background.comoving_distance(z_star) + efficiency = 1 - rz / rz_star + result = factor * efficiency + return np.expand_dims(result, 0) diff --git a/cloelib/summary_statistics/angular_two_point.py b/cloelib/summary_statistics/angular_two_point.py index 56aa4fc2..83cf7df6 100644 --- a/cloelib/summary_statistics/angular_two_point.py +++ b/cloelib/summary_statistics/angular_two_point.py @@ -4,6 +4,7 @@ from cloelib.observables.tracer import Tracer from cloelib.observables.photo import PositionsTracer from cloelib.observables.photo import ShearTracer +from cloelib.observables.cmb import CMBLensingTracer from cloelib.auxiliary.units import SPEED_OF_LIGHT from cloelib.auxiliary.math_utils import simpsons_weights_jit from cloelib.profiling import profile_function @@ -291,7 +292,8 @@ def get_Cl(self, ells, nl, ks) -> dict: C_ell_calc = C_ell_calc * prefactor_cell[:, None, None] self.C_ell_calc = C_ell_calc - n_bin = self.tracer1.n_z_bins + n_bin1 = self.tracer1.n_z_bins + n_bin2 = self.tracer2.n_z_bins C_ell_out = {} # Prepare dictionary with tuples as keys for the output @@ -309,6 +311,11 @@ def get_Cl(self, ells, nl, ks) -> dict: # Why? Because it expects the cross-correlation between positions and shear. # so the second dimension is filled with zeros. # POS - POS returns an array of shape (len(ells)) + # CMBL - SHE returns an array of shape (2, len(ells)) + # Why? Because it expects the cross-correlation between CMB lensing and shear. + # so the second dimension is filled with zeros. + # CMBL - POS returns an array of shape (len(ells)) + # CMBL - CMBL returns an array of shape (len(ells)) def pos_pos_rule(C, i, j): return {("POS", "POS", i, j): C[:, i - 1, j - 1]} @@ -328,10 +335,25 @@ def she_she_rule(C, i, j): arr = arr.at[0, 0, :].set(block) return {("SHE", "SHE", i, j): arr} + def cmbl_cmbl_rule(C, i, j): + return {("CMBL", "CMBL", i, j): C[:, i - 1, j - 1]} + + def cmbl_pos_rule(C, i, j): + a, b = sorted((i, j)) + return {("CMBL", "POS", a, b): C[:, i - 1, j - 1]} + + def cmbl_she_rule(C, i, j): + block = C[:, i - 1, j - 1] + a, b = sorted((i, j)) + return {("CMBL", "SHE", a, b): np.stack([block, np.zeros_like(block)])} + tracer_rules = { (PositionsTracer, PositionsTracer): pos_pos_rule, (PositionsTracer, ShearTracer): pos_she_rule, (ShearTracer, ShearTracer): she_she_rule, + (CMBLensingTracer, PositionsTracer): cmbl_pos_rule, + (CMBLensingTracer, ShearTracer): cmbl_she_rule, + (CMBLensingTracer, CMBLensingTracer): cmbl_cmbl_rule, } # normalize the key so (A, B) and (B, A) are both supported @@ -346,11 +368,16 @@ def she_she_rule(C, i, j): ) # Vectorized update of C_ell_out using dictionary comprehensions + a, b = sorted((n_bin1, n_bin2)) C_ell_out = { k: v - for i in range(1, n_bin + 1) - for j in range(i, n_bin + 1) - for k, v in rule_fn(C_ell_calc, i, j).items() + for i in range(1, a + 1) + for j in range(i, b + 1) + for k, v in ( + rule_fn(C_ell_calc, i, j) + if n_bin1 <= n_bin2 + else rule_fn(C_ell_calc, j, i) + ).items() } # Use dictionary comprehension for cosmolib_Cls creation diff --git a/docs/code_structure/background.md b/docs/code_structure/background.md index 185870b1..b1897f09 100644 --- a/docs/code_structure/background.md +++ b/docs/code_structure/background.md @@ -39,6 +39,7 @@ Every Background implementation must provide: - **`gamma_MG`**: Modified gravity parameter - **`rdrag`**: Sound horizon radius at last scattering (Mpc) - **`interface_args`**: Dictionary storing interface-specific parameters +- **`z_star`**: Redshift of photon decoupling. ### Required Methods diff --git a/docs/code_structure/observables.md b/docs/code_structure/observables.md index 71d3e242..289e5973 100644 --- a/docs/code_structure/observables.md +++ b/docs/code_structure/observables.md @@ -10,7 +10,7 @@ This module computes survey-specific quantities including selection functions, w This module addresses: -- Window functions for weak lensing surveys +- Window functions for weak lensing surveys and CMB lensing - Galaxy bias modeling and corrections - Redshift-space power spectra P(k, μ) @@ -166,6 +166,30 @@ tracer = PositionsTracer( window = tracer.get_window(z) ``` +#### CMBLensingTracer + +For CMB weak gravitational lensing (convergence) measurements. + +**Location**: `cloelib/observables/cmb.py` + +**What it does**: + +- Computes lensing window function W^κ(z) + +**Example**: + +```python +from cloelib.observables.cmb import CMBLensingTracer + +z = np.arange(1e-3,pert.background.z_star,0.01) +tracer = CMBLensingTracer( + perturbations=pert, + z=z, +) + +window = tracer.get_window(z) +``` + ### Adding Your Own Tracer To add a new type of photometric observable, follow these steps. diff --git a/tests/test_cmb.py b/tests/test_cmb.py new file mode 100644 index 00000000..a7eafeaf --- /dev/null +++ b/tests/test_cmb.py @@ -0,0 +1,139 @@ +""" +Unit tests to verify cmb tracers run correctly. +""" + +import pytest +import numpy as np + +from cloelib.cosmology.camb_cosmology import ( + CAMBBackground, + CAMBNonLinearPerturbations, +) +from cloelib.observables.cmb import CMBLensingTracer +from cloelib.observables.photo import ShearTracer, PositionsTracer +from cloelib.summary_statistics.angular_two_point import AngularTwoPoint + + +@pytest.fixture +def camb_cmb_setup(): + """Create CAMB perturbations and position/shear tracers.""" + # Create background (using same params as test_camb_cosmology.py) + H0 = 67.7 + h = H0 / 100.0 + omch2 = 0.12 + Omega_cdm0 = omch2 / h**2 + ombh2 = 0.022 + Omega_b0 = ombh2 / h**2 + + background = CAMBBackground( + H0=H0, + Omega_b0=Omega_b0, + Omega_cdm0=Omega_cdm0, + Omega_k0=0.0, + As=2e-9, + ns=0.96, + mnu=0.06, + w0=-1.0, + wa=0.0, + gamma_MG=0.0, + N_mnu=1, + ) + + z_auto = np.linspace(0.01, 1100.0, 100) + z_cross = np.linspace(0.2, 2.0, 40) + perturbations = CAMBNonLinearPerturbations(background, z_auto) + n_z_bins = 1 + dndz = np.ones((n_z_bins, len(z_cross))) + dndz /= np.trapezoid(dndz, z_cross, axis=1)[:, None] + return perturbations, z_auto, z_cross, dndz + + +def test_cmb_lensing_window_cl(camb_cmb_setup): + """ + Test that cmb lensing window function auto and cross power spectra can be computed without error. + + It further verifies that: + 1. CMB lensing window function has the expected shape. + 2. CMB lensing power spectrum and cross-correlations have the expected shapes. + """ + perturbations, z_auto, z_cross, dndz = camb_cmb_setup + + # Create nuisance parameters + nuisance_pos = { + "b1_photo_bin1": 1.0, + "dz_pos_1": 0.0, + "width_pos_1": 1.0, + "magnification_bias_1": 0.0, + } + nuisance_shear = { + "AIA": 0.0, + "CIA": 0.0, + "EtaIA": 0.0, + "multiplicative_bias_1": 0.0, + "dz_shear_1": 0.0, + "width_shear_1": 1.0, + } + + # Create CMB lensing tracers + cmblens_cross = CMBLensingTracer(perturbations=perturbations, z=z_cross) + cmblens_auto = CMBLensingTracer(perturbations=perturbations, z=z_auto) + + window_cross = cmblens_cross.get_window(z_cross) + window_auto = cmblens_auto.get_window(z_auto) + + # Check that the window functions have the expected shape + assert window_cross.shape == (1, len(z_cross)) + assert window_auto.shape == (1, len(z_auto)) + + # Create position tracer + pos_tracer = PositionsTracer( + perturbations=perturbations, + dndz=dndz, + z=z_cross, + galaxy_bias_model="per_bin", + nuisance_params=nuisance_pos, + ) + + # Create shear tracer + shear_tracer = ShearTracer( + perturbations=perturbations, + dndz=dndz, + z=z_cross, + nuisance_params=nuisance_shear, + ) + + nl = 10 + ells = np.logspace(1.0, np.log10(100), nl) + + twopoint_kpos = AngularTwoPoint(cmblens_cross, pos_tracer) + twopoint_kshe = AngularTwoPoint(cmblens_cross, shear_tracer) + twopoint_kk = AngularTwoPoint(cmblens_auto, cmblens_auto) + + cells = { + **twopoint_kpos.get_Cl(ells, 0, perturbations.k), + **twopoint_kshe.get_Cl(ells, 0, perturbations.k), + **twopoint_kk.get_Cl(ells, 0, perturbations.k), + } + + # Check that cells has the right number of elements and right keys + assert len(cells.keys()) == 3 + assert ("CMBL", "POS", 1, 1) in cells.keys() + assert ("CMBL", "SHE", 1, 1) in cells.keys() + assert ("CMBL", "CMBL", 1, 1) in cells.keys() + assert cells[("CMBL", "POS", 1, 1)].shape == (nl,) + assert cells[("CMBL", "SHE", 1, 1)].shape == (2, nl) + assert cells[("CMBL", "CMBL", 1, 1)].shape == (nl,) + + # We also check everything works well if we reverse the order of the probes + twopoint_posk = AngularTwoPoint(pos_tracer, cmblens_cross) + twopoint_shek = AngularTwoPoint(shear_tracer, cmblens_cross) + cells_reverse = { + **twopoint_posk.get_Cl(ells, 0, perturbations.k), + **twopoint_shek.get_Cl(ells, 0, perturbations.k), + } + + assert len(cells_reverse.keys()) == 2 + assert ("CMBL", "POS", 1, 1) in cells_reverse.keys() + assert ("CMBL", "SHE", 1, 1) in cells_reverse.keys() + assert cells_reverse[("CMBL", "POS", 1, 1)].shape == (nl,) + assert cells_reverse[("CMBL", "SHE", 1, 1)].shape == (2, nl) diff --git a/tests/test_cosmology_background.py b/tests/test_cosmology_background.py index 4233816f..64994b58 100644 --- a/tests/test_cosmology_background.py +++ b/tests/test_cosmology_background.py @@ -53,6 +53,7 @@ def test_background_required_attributes(): "N_eff", "interface_args", "rdrag", + "z_star", } assert attributes_required == attributes_found