gh-268: including CMB lensing

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."""

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."""

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):

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

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."""

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
@@ -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

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
 

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.