gh-268: including CMB lensing

cloelib/cosmology/camb_cosmology.py

@@ -287,6 +287,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."""

cloelib/cosmology/class_cosmology.py

@@ -267,6 +267,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."""

cloelib/cosmology/cosmology.py

@@ -126,6 +126,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):

cloelib/cosmology/derived_cosmology.py

@@ -79,3 +79,23 @@ 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, neff=3.046):
+    r"""Compute the redshift of photon decoupling.
+
+    Assumes a cosmology-independent z_star.
+
+    Parameters
+    ----------
+    background: Background
+        Background class containing cosmology
+    neff: float
+        Effective number of neutrinos.
+
+    Returns
+    -------
+    z_star: float
+        redshift of photon decoupling.
+    """
+    return 1090.0

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
@@ -327,6 +330,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."""

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)

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
@@ -182,7 +183,8 @@ def get_Cl(self, ells, nl, ks) -> dict:
         )
         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
@@ -200,6 +202,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]}
@@ -219,10 +226,23 @@ 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):
+            return {("CMBL", "POS", i, j): C[:, i - 1, j - 1]}
+
+        def cmbl_she_rule(C, i, j):
+            block = C[:, i - 1, j - 1]
+            return {("CMBL", "SHE", i, j): 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
@@ -237,11 +257,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

tests/test_cmb.py

@@ -0,0 +1,133 @@
+"""
+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.trapz(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, "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,
+    }
+
+    # 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)

tests/test_cosmology_background.py

@@ -52,6 +52,7 @@ def test_background_required_attributes():
         "N_eff",
         "interface_args",
         "rdrag",
+        "z_star",
     }
     assert attributes_required == attributes_found