{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# CMB Tutorial \u2014 Companion Notebook\n", "\n", "Every code cell from the tutorial in order. Run top to bottom; each cell defines what later cells use.\n", "\n", "**Prerequisites:** `numpy`, `scipy`, `matplotlib`, optional `numba`. CAMB is used only in validation cells.\n", "\n", "```\npip install numpy scipy matplotlib jupyter\npip install camb # for validation only\n```\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup\n", "\n", "`np.trapz` was renamed to `np.trapezoid` in numpy 2.0. This shim keeps the rest of the notebook working either way.\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "import numpy as np\n", "if not hasattr(np, \"trapz\"):\n", " np.trapz = np.trapezoid\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Prerequisites\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell 1\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "from scipy import integrate, interpolate, optimize, special\n", "\n", "# Optional: numba speeds up the inner Boltzmann RHS by ~10x.\n", "# Falls back gracefully if not installed.\n", "try:\n", " from numba import njit\n", " _jit = njit(cache=False)\n", "except ImportError:\n", " _jit = lambda f: f\n", "\n", "# Default: Planck 2018 best-fit LCDM\n", "params = {\n", " 'omega_b_h2': 0.02237, # physical baryon density\n", " 'omega_c_h2': 0.1200, # physical cold dark matter density\n", " 'h': 0.6736, # H_0 / (100 km/s/Mpc)\n", " 'n_s': 0.9649, # scalar spectral index\n", " 'A_s': 2.1e-9, # scalar amplitude at k_pivot = 0.05 Mpc^-1\n", " 'tau_reion': 0.0544, # reionisation optical depth\n", " 'N_eff': 3.044, # effective number of relativistic species\n", " 'T_cmb': 2.7255, # CMB temperature today (K)\n", " 'Y_He': 0.245, # primordial helium mass fraction\n", " 'k_pivot': 0.05, # pivot scale (1/Mpc)\n", " 'ell_max': 2500, # maximum multipole we will compute\n", "}\n", "\n", "plt.rcParams.update({'figure.figsize': (7, 4.5), 'font.size': 11})\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 1 \u2014 Background\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: physical constants\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Physical constants used throughout the tutorial.\n", "c_km_s = 2.99792458e5 # speed of light (km/s)\n", "k_B = 1.380649e-23 # Boltzmann constant (J/K)\n", "h_P = 6.62607015e-34 # Planck constant (J s)\n", "m_e = 9.1093837015e-31 # electron mass (kg)\n", "m_H = 1.673575e-27 # hydrogen atom mass (kg)\n", "m_He4 = 6.646479073e-27 # He-4 atom mass (kg)\n", "not4 = m_He4 / m_H # He/H mass ratio (~3.97, not exactly 4)\n", "sigma_T = 6.6524587321e-29 # Thomson cross section (m^2)\n", "G = 6.67430e-11 # gravitational constant (m^3/kg/s^2)\n", "Mpc_in_m = 3.0856775814913673e22 # 1 Mpc in metres\n", "sigma_SB = 5.670374419e-8 # Stefan-Boltzmann constant (W/m^2/K^4)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `init_background`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def init_background(params):\n", " \"\"\"Convert cosmological parameters to internal `grho_i` densities,\n", " where grho_i = 8*pi*G*rho_{i,0} in Mpc^-2 (CAMB convention).\n", " \"\"\"\n", " h = params['h']\n", " H0 = 100 * h / c_km_s # H_0 in 1/Mpc (c = 1)\n", " grhocrit_h2 = 3 * (100.0 / c_km_s)**2 # 3*(H_100/c)^2 in 1/Mpc^2\n", "\n", " # Photon energy density from the blackbody formula\n", " T_cmb = params['T_cmb']\n", " rho_gamma = 4 * sigma_SB / (c_km_s * 1e3)**3 * T_cmb**4 # kg/m^3\n", " H100_SI = 100 * 1e3 / Mpc_in_m # 1/s\n", " rho_crit_100 = 3 * H100_SI**2 / (8 * np.pi * G) # kg/m^3\n", " omega_gamma = rho_gamma / rho_crit_100\n", "\n", " # Internal densities: 8*pi*G*rho_{i,0} in Mpc^-2\n", " grhog = grhocrit_h2 * omega_gamma # photons (a^-4)\n", " grhornomass = grhog * 7/8 * (4/11)**(4/3) * params['N_eff'] # massless nu (a^-4)\n", " grhoc = grhocrit_h2 * params['omega_c_h2'] # CDM (a^-3)\n", " grhob = grhocrit_h2 * params['omega_b_h2'] # baryons (a^-3)\n", "\n", " # Cosmological constant: Omega_L = 1 - Omega_m - Omega_r (flat universe)\n", " omega_m = (params['omega_c_h2'] + params['omega_b_h2']) / h**2\n", " omega_r = omega_gamma * (1 + 7/8 * (4/11)**(4/3) * params['N_eff']) / h**2\n", " grhov = grhocrit_h2 * (1 - omega_m - omega_r) * h**2 # Lambda (const)\n", "\n", " # Two thermodynamic quantities used much later (chapter 2): the\n", " # Thomson opacity prefactor, and the He/H number-fraction f_He.\n", " rho_b_SI = params['omega_b_h2'] * rho_crit_100\n", " n_H_Mpc = (1 - params['Y_He']) * rho_b_SI / m_H * Mpc_in_m**3\n", " akthom = (sigma_T / Mpc_in_m**2) * n_H_Mpc\n", " f_He = params['Y_He'] / (not4 * (1 - params['Y_He']))\n", "\n", " return {\n", " 'H0': H0, 'h': h,\n", " 'grhog': grhog, 'grhornomass': grhornomass,\n", " 'grhoc': grhoc, 'grhob': grhob, 'grhov': grhov,\n", " 'akthom': akthom, 'f_He': f_He,\n", " 'Y_He': params['Y_He'], 'T_cmb': T_cmb,\n", " }\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: Friedmann evaluators\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def total_density_a4(a, bg):\n", " \"\"\"Returns 8*pi*G*rho_tot*a^4 -- the polynomial in eq. (above) that\n", " sits inside the square root of the Friedmann equation.\"\"\"\n", " return (bg['grhog'] + bg['grhornomass']\n", " + (bg['grhoc'] + bg['grhob']) * a\n", " + bg['grhov'] * a**4)\n", "\n", "def dtau_da(a, bg):\n", " \"\"\"d(eta)/da = 1/(a^2 H) = sqrt(3 / total_density_a4), in Mpc.\"\"\"\n", " return np.sqrt(3.0 / total_density_a4(a, bg))\n", "\n", "def hubble(a, bg):\n", " \"\"\"Hubble parameter H(a) = sqrt(total_density_a4 / 3) / a^2, in 1/Mpc (c=1).\"\"\"\n", " return np.sqrt(total_density_a4(a, bg) / 3.0) / a**2\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: integrals over the background\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def conformal_time(a, bg):\n", " \"\"\"eta(a) = int_0^a da'/(a'^2 H), in Mpc.\"\"\"\n", " a = np.atleast_1d(a)\n", " out = np.array([\n", " integrate.quad(dtau_da, 0, ai, args=(bg,), limit=100, epsrel=1e-8)[0]\n", " for ai in a])\n", " return out.squeeze()\n", "\n", "def sound_horizon(a, bg):\n", " \"\"\"Comoving sound horizon r_s(a) = int_0^a c_s da'/(a'^2 H), in Mpc.\"\"\"\n", " def integrand(ap):\n", " R = 0.75 * bg['grhob'] * ap / bg['grhog']\n", " return 1.0 / np.sqrt(total_density_a4(ap, bg) * (1.0 + R))\n", " return integrate.quad(integrand, 0, a)[0]\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `build_background`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def build_background(params):\n", " \"\"\"Top-level: prepare the background for downstream use.\"\"\"\n", " bg = init_background(params)\n", " bg['tau0'] = float(conformal_time(1.0, bg))\n", " a_eq = (bg['grhog'] + bg['grhornomass']) / (bg['grhoc'] + bg['grhob'])\n", " bg['a_eq'] = a_eq\n", " bg['z_eq'] = 1.0 / a_eq - 1.0\n", " bg['tau_eq'] = float(conformal_time(a_eq, bg))\n", " return bg\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: run it\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "bg = build_background(params)\n", "print(f\"H_0 = {bg['H0']:.4e} 1/Mpc = {bg['H0']*c_km_s:.2f} km/s/Mpc\")\n", "print(f\"tau_0 = {bg['tau0']:.1f} Mpc\")\n", "print(f\"z_eq = {bg['z_eq']:.0f}\")\n", "print(f\"tau_eq = {bg['tau_eq']:.2f} Mpc\")\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot $H(z)/H_0$\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "a_grid = np.geomspace(1e-7, 1.0, 400)\n", "z_grid = 1.0 / a_grid - 1.0\n", "H_grid = np.array([hubble(ai, bg) for ai in a_grid])\n", "\n", "# Where does matter equal Lambda?\n", "rho_m = (bg['grhoc'] + bg['grhob']) / a_grid**3\n", "rho_L = bg['grhov'] * np.ones_like(a_grid)\n", "a_lambda = a_grid[np.argmin(np.abs(rho_m - rho_L))]\n", "z_lambda = 1.0 / a_lambda - 1.0\n", "\n", "fig, ax = plt.subplots()\n", "sel = (z_grid > 1e-3) & (z_grid < 1e5)\n", "ax.loglog(z_grid[sel], H_grid[sel] / bg['H0'], 'k-', lw=1.6)\n", "ax.axvline(bg['z_eq'], color='C3', ls=':', alpha=0.7)\n", "ax.text(bg['z_eq']*1.3, 2, f\"$z_\\\\mathrm{{eq}} = {bg['z_eq']:.0f}$\", color='C3')\n", "ax.axvline(z_lambda, color='C2', ls=':', alpha=0.7)\n", "ax.text(z_lambda*1.3, 2, f\"$z_\\\\Lambda = {z_lambda:.2f}$\", color='C2')\n", "ax.set_xlabel('Redshift $z$')\n", "ax.set_ylabel('$H(z) / H_0$')\n", "ax.set_title('Hubble expansion history')\n", "ax.grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot $Omega_i(a)$\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "rho_r = (bg['grhog'] + bg['grhornomass']) / a_grid**4\n", "rho_m = (bg['grhoc'] + bg['grhob']) / a_grid**3\n", "rho_L = bg['grhov'] * np.ones_like(a_grid)\n", "rho_tot = rho_r + rho_m + rho_L\n", "\n", "fig, ax = plt.subplots()\n", "ax.loglog(a_grid, rho_m/rho_tot, label='Matter', color='C0', lw=1.6)\n", "ax.loglog(a_grid, rho_r/rho_tot, label='Radiation', color='C3', lw=1.6)\n", "ax.loglog(a_grid, rho_L/rho_tot, label=r'$\\Lambda$', color='C2', lw=1.6)\n", "ax.axvline(bg['a_eq'], color='gray', ls=':', alpha=0.5)\n", "ax.axvline(a_lambda, color='gray', ls=':', alpha=0.5)\n", "ax.text(bg['a_eq']*1.3, 1.2e-5, r'$a_\\mathrm{eq}$', color='gray')\n", "ax.text(a_lambda*1.3, 1.2e-5, r'$a_\\Lambda$', color='gray')\n", "ax.set_xlabel('Scale factor $a$')\n", "ax.set_ylabel(r'$\\Omega_i(a)$')\n", "ax.set_title('Composition of the universe')\n", "ax.set_ylim(1e-6, 2); ax.set_xlim(1e-7, 1)\n", "ax.legend(loc='lower right'); ax.grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 2 \u2014 Recombination\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: atomic constants\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "c_SI = c_km_s * 1e3 # speed of light (m/s)\n", "\n", "# Atomic transition wavenumbers (1/m)\n", "L_H_ion = 1.096787737e7 # H ionization\n", "L_H_alpha = 8.225916453e6 # H Lyman-alpha\n", "L_He1_ion = 1.98310772e7 # HeI ionization\n", "L_He2_ion = 4.389088863e7 # HeII ionization\n", "L_He_2s = 1.66277434e7 # HeI 2^1 S_0\n", "L_He_2p = 1.71134891e7 # HeI 2^1 P_1\n", "\n", "# Decay/transition rates\n", "Lambda_2s1s = 8.2245809 # H 2s->1s two-photon (s^-1)\n", "Lambda_He = 51.3 # HeI 2s->1s two-photon (s^-1)\n", "A2P_s = 1.798287e9 # HeI singlet Einstein A (s^-1)\n", "A2P_t = 177.58 # HeI triplet Einstein A (s^-1)\n", "L_He_2Pt = 1.690871466e7 # HeI triplet 2^3 P (1/m)\n", "L_He_2St = 1.5985597526e7 # HeI triplet 2^3 S (1/m)\n", "L_He2St_ion = 3.8454693845e6 # HeI 2^3 S ionisation (1/m)\n", "sigma_He_2Ps = 1.436289e-22 # HeI singlet photoionisation sigma (m^2)\n", "sigma_He_2Pt = 1.484872e-22 # HeI triplet photoionisation sigma (m^2)\n", "\n", "# RECFAST fudge factors (calibrated against full level codes)\n", "RECFAST_fudge = 1.125\n", "AGauss1, AGauss2 = -0.14, 0.079\n", "zGauss1, zGauss2 = 7.28, 6.73\n", "wGauss1, wGauss2 = 0.18, 0.33\n", "\n", "# Derived constants\n", "CR = 2 * np.pi * m_e * k_B / h_P**2 # Saha coefficient (1/m^2/K)\n", "CB1 = h_P * c_SI * L_H_ion / k_B # H ionisation / k_B (K)\n", "CB1_He1 = h_P * c_SI * L_He1_ion / k_B\n", "CB1_He2 = h_P * c_SI * L_He2_ion / k_B\n", "CDB = h_P * c_SI * (L_H_ion - L_H_alpha) / k_B\n", "CDB_He = h_P * c_SI * (L_He1_ion - L_He_2s) / k_B\n", "CK = (1.0 / L_H_alpha)**3 / (8 * np.pi) # lambda_alpha^3 / (8 pi) (m^3)\n", "CK_He = (1.0 / L_He_2p)**3 / (8 * np.pi)\n", "CL = h_P * c_SI * L_H_alpha / k_B\n", "CL_He = h_P * c_SI * L_He_2s / k_B\n", "Bfact = h_P * c_SI * (L_He_2p - L_He_2s) / k_B\n", "CL_PSt = h_P * c_SI * (L_He_2Pt - L_He_2St) / k_B\n", "CB1_He2St = h_P * c_SI * L_He2St_ion / k_B\n", "CL_He_2St = h_P * c_SI * L_He_2St / k_B\n", "a_rad = 4 * sigma_SB / c_SI # radiation constant\n", "CT = (8/3) * (sigma_T / (m_e * c_SI)) * a_rad # Compton cooling\n", "barssc0 = k_B / (m_H * c_SI**2) # baryon sound speed prefactor (1/K)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `solve_recombination` -- part 1 -- Saha and RHS\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def solve_recombination(bg, params):\n", " \"\"\"Solve the ionisation history x_e(z) using the RECFAST equations.\n", "\n", " Three-variable ODE for hydrogen ionisation x_H, helium ionisation x_He,\n", " and matter temperature T_m. Saha equilibrium at high z, switching to\n", " ODEs as each species departs from equilibrium.\n", " \"\"\"\n", " T_cmb = bg['T_cmb']\n", " f_He = bg['f_He']\n", "\n", " # Present-day hydrogen number density (m^-3)\n", " H100_SI = 100 * 1e3 / Mpc_in_m\n", " rho_crit_100 = 3 * H100_SI**2 / (8 * np.pi * G)\n", " Nnow = (1 - bg['Y_He']) * (params['omega_b_h2'] * rho_crit_100) / m_H\n", "\n", " # Hubble in SI for use inside the ODE\n", " omega_m = (params['omega_b_h2'] + params['omega_c_h2']) / params['h']**2\n", " a_eq = (bg['grhog'] + bg['grhornomass']) / (bg['grhoc'] + bg['grhob'])\n", " z_eq = 1.0 / a_eq - 1.0\n", " H0_SI = bg['H0'] * c_SI / Mpc_in_m\n", " def Hz_SI(z):\n", " return hubble(1.0 / (1 + z), bg) * c_SI / Mpc_in_m\n", "\n", " # --- Saha equations for helium ---\n", " def saha_He2(z):\n", " \"\"\"He++ -> He+ Saha: returns total x_e per H atom.\"\"\"\n", " T = T_cmb * (1 + z)\n", " rhs = (CR * T_cmb / (1 + z))**1.5 * np.exp(-CB1_He2 / T) / Nnow\n", " return 0.5 * (np.sqrt((rhs - 1 - f_He)**2\n", " + 4 * (1 + 2 * f_He) * rhs) - (rhs - 1 - f_He))\n", "\n", " def saha_He1(z):\n", " \"\"\"He+ -> He0 Saha: returns x_He = n(He+) / n_He.\"\"\"\n", " T = T_cmb * (1 + z)\n", " rhs = 4.0 * (CR * T_cmb / (1 + z))**1.5 * np.exp(-CB1_He1 / T) / Nnow\n", " x0 = 0.5 * (np.sqrt((rhs - 1)**2 + 4 * (1 + f_He) * rhs) - (rhs - 1))\n", " return min((x0 - 1.0) / f_He, 1.0)\n", "\n", " # --- RECFAST 3-variable RHS: dy/dz for y = [x_H, x_He, T_mat] ---\n", " def recfast_rhs(z, y):\n", " x_H = max(y[0], 0.0)\n", " x_He = max(y[1], 0.0)\n", " T_mat = max(y[2], 0.5)\n", " x = x_H + f_He * x_He\n", " T_rad = T_cmb * (1 + z)\n", " n_H = Nnow * (1 + z)**3\n", " n_He = f_He * n_H\n", " Hz = Hz_SI(z)\n", "\n", " # ---- f1: Hydrogen Peebles equation ----\n", " t4 = T_mat / 1e4\n", " Rdown = 1e-19 * 4.309 * t4**(-0.6166) / (1 + 0.6703 * t4**0.5300)\n", " Rup = Rdown * (CR * T_mat)**1.5 * np.exp(-CDB / T_mat)\n", " K = CK / Hz * (1.0\n", " + AGauss1 * np.exp(-((np.log(1 + z) - zGauss1) / wGauss1)**2)\n", " + AGauss2 * np.exp(-((np.log(1 + z) - zGauss2) / wGauss2)**2))\n", " fu = RECFAST_fudge\n", " n_1s = n_H * max(1 - x_H, 1e-30)\n", " f1 = ((x * x_H * n_H * Rdown\n", " - Rup * (1 - x_H) * np.exp(-CL / T_mat))\n", " * (1 + K * Lambda_2s1s * n_1s)\n", " / (Hz * (1 + z) * (1.0/fu + K * Lambda_2s1s * n_1s / fu\n", " + K * Rup * n_1s)))\n", " # ---- f2: Helium singlet ODE + triplet correction ----\n", " if x_He < 1e-15:\n", " f2 = 0.0\n", " else:\n", " T_0 = 10.0**0.477121\n", " T_1 = 10.0**5.114\n", " sq_0 = np.sqrt(T_mat / T_0)\n", " sq_1 = np.sqrt(T_mat / T_1)\n", " Rdown_He = 10.0**(-16.744) / (sq_0 * (1 + sq_0)**0.289\n", " * (1 + sq_1)**1.711)\n", " Rup_He = 4.0 * Rdown_He * (CR * T_mat)**1.5 * np.exp(-CDB_He / T_mat)\n", " He_Boltz = np.exp(min(Bfact / T_mat, 500.0))\n", "\n", " n_He_ground = n_He * max(1 - x_He, 1e-30)\n", " tauHe_s = A2P_s * CK_He * 3 * n_He_ground / Hz\n", " pHe_s = ((1 - np.exp(-tauHe_s)) / tauHe_s\n", " if tauHe_s > 1e-7 else 1.0 - tauHe_s / 2.0)\n", "\n", " if x_H < 0.9999999:\n", " Doppler_s = c_SI * L_He_2p * np.sqrt(2 * k_B * T_mat\n", " / (m_H * not4 * c_SI**2))\n", " gamma_2Ps = (3 * A2P_s * f_He * (1 - x_He) * c_SI**2\n", " / (np.sqrt(np.pi) * sigma_He_2Ps * 8 * np.pi\n", " * Doppler_s * max(1 - x_H, 1e-30)\n", " * (c_SI * L_He_2p)**2))\n", " AHcon_s = A2P_s / (1 + 0.36 * gamma_2Ps**0.86)\n", " K_He = 1.0 / max((A2P_s * pHe_s + AHcon_s) * 3 * n_He_ground, 1e-300)\n", " else:\n", " K_He = 1.0 / max(A2P_s * pHe_s * 3 * n_He_ground, 1e-300)\n", "\n", " f2 = ((x * x_He * n_H * Rdown_He\n", " - Rup_He * (1 - x_He) * np.exp(-CL_He / T_mat))\n", " * (1 + K_He * Lambda_He * n_He_ground * He_Boltz)\n", " / (Hz * (1 + z)\n", " * (1 + K_He * (Lambda_He + Rup_He)\n", " * n_He_ground * He_Boltz)))\n", "\n", " # Triplet channel\n", " if x_He > 5e-9:\n", " a_trip = 10.0**(-16.306); b_trip = 0.761\n", " Rdown_trip = a_trip / (sq_0 * (1 + sq_0)**(1.0 - b_trip)\n", " * (1 + sq_1)**(1.0 + b_trip))\n", " Rup_trip = (4.0/3.0) * Rdown_trip * (CR * T_mat)**1.5 * np.exp(-CB1_He2St / T_mat)\n", " tauHe_t = A2P_t * n_He_ground * 3 / (8 * np.pi * Hz * L_He_2Pt**3)\n", " pHe_t = ((1 - np.exp(-tauHe_t)) / tauHe_t\n", " if tauHe_t > 1e-7 else 1.0 - tauHe_t / 2.0)\n", "\n", " if x_H < 0.99999:\n", " Doppler_t = c_SI * L_He_2Pt * np.sqrt(2 * k_B * T_mat\n", " / (m_H * not4 * c_SI**2))\n", " gamma_2Pt = (3 * A2P_t * f_He * (1 - x_He) * c_SI**2\n", " / (np.sqrt(np.pi) * sigma_He_2Pt * 8 * np.pi\n", " * Doppler_t * max(1 - x_H, 1e-30)\n", " * (c_SI * L_He_2Pt)**2))\n", " AHcon_t = A2P_t / (1 + 0.66 * gamma_2Pt**0.9) / 3.0\n", " CfHe_t = (A2P_t * pHe_t + AHcon_t) * np.exp(-CL_PSt / T_mat)\n", " else:\n", " CfHe_t = A2P_t * pHe_t * np.exp(-CL_PSt / T_mat)\n", " denom = Rup_trip + CfHe_t\n", " CfHe_t = CfHe_t / denom if denom > 1e-300 else 0.0\n", "\n", " f2 += ((x * x_He * n_H * Rdown_trip\n", " - (1 - x_He) * 3 * Rup_trip * np.exp(-CL_He_2St / T_mat))\n", " * CfHe_t / (Hz * (1 + z)))\n", "\n", " # ---- f3: Matter temperature ----\n", " x_safe = max(x, 1e-30)\n", " timeTh = (1.0 / (CT * T_rad**4)) * (1 + x + f_He) / x_safe\n", " timeH = 2.0 / (3.0 * H0_SI * (1 + z)**1.5)\n", " if timeTh < 1e-3 * timeH:\n", " # Tightly coupled: implicit form\n", " dHdz = (H0_SI**2 / (2 * Hz)) * omega_m * (\n", " 4 * (1 + z)**3 / (1 + z_eq) + 3 * (1 + z)**2)\n", " epsilon = Hz * (1 + x + f_He) / (CT * T_rad**3 * x_safe)\n", " f3 = (T_cmb\n", " + epsilon * (1 + f_He) / (1 + f_He + x)\n", " * (f1 + f_He * f2) / x_safe\n", " - epsilon * dHdz / Hz\n", " + 3 * epsilon / (1 + z))\n", " else:\n", " # Loose coupling: Compton cooling + adiabatic\n", " f3 = (CT * T_rad**4 * x_safe / (1 + x + f_He)\n", " * (T_mat - T_rad) / (Hz * (1 + z))\n", " + 2 * T_mat / (1 + z))\n", "\n", " return [f1, f2, f3]\n", " # --- Build z grid ---\n", " z_start, z_end, nz = 10000, 0, 20000\n", " z_arr = np.linspace(z_start, z_end, nz + 1)\n", " xH_arr = np.ones(nz + 1)\n", " xHe_arr = np.ones(nz + 1)\n", " xe_total = np.empty(nz + 1)\n", "\n", " # --- Phase 1: Saha equilibrium ---\n", " he_ode_idx = None\n", " for i, z in enumerate(z_arr):\n", " if z > 8000.0:\n", " xH_arr[i] = 1.0; xHe_arr[i] = 1.0\n", " xe_total[i] = 1.0 + 2.0 * f_He\n", " elif z > 5000.0:\n", " xH_arr[i] = 1.0; xHe_arr[i] = 1.0\n", " xe_total[i] = saha_He2(z)\n", " elif z > 3500.0:\n", " xH_arr[i] = 1.0; xHe_arr[i] = 1.0\n", " xe_total[i] = 1.0 + f_He\n", " elif z > 0:\n", " x_He = saha_He1(z)\n", " xHe_arr[i] = x_He\n", " xH_arr[i] = 1.0\n", " xe_total[i] = 1.0 + f_He * x_He\n", " if x_He < 0.99:\n", " he_ode_idx = i\n", " break\n", " else:\n", " break\n", " if he_ode_idx is None:\n", " he_ode_idx = len(z_arr) - 1\n", "\n", " # --- Phase 2: 3-variable ODE from He departure to z = 0 ---\n", " z_ode = z_arr[he_ode_idx:]\n", " z_ode = z_ode[z_ode >= 0]\n", " Tmat_arr = T_cmb * (1.0 + z_arr)\n", " if len(z_ode) > 1:\n", " y0 = [xH_arr[he_ode_idx], xHe_arr[he_ode_idx],\n", " T_cmb * (1 + z_ode[0])]\n", " sol = integrate.solve_ivp(\n", " recfast_rhs,\n", " [z_ode[0], z_ode[-1]], y0,\n", " t_eval=z_ode, method='Radau', rtol=1e-6, atol=1e-10,\n", " max_step=5.0,\n", " )\n", " n_sol = min(sol.y.shape[1], len(z_arr) - he_ode_idx)\n", " idx = he_ode_idx + n_sol\n", " xH_arr[he_ode_idx:idx] = sol.y[0, :n_sol]\n", " xHe_arr[he_ode_idx:idx] = sol.y[1, :n_sol]\n", " Tmat_arr[he_ode_idx:idx] = sol.y[2, :n_sol]\n", " xe_total[he_ode_idx:] = xH_arr[he_ode_idx:] + f_He * xHe_arr[he_ode_idx:]\n", " return z_arr, xe_total, Tmat_arr\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `build_thermodynamics`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def build_thermodynamics(bg, params):\n", " \"\"\"Build thermodynamic tables: opacity, optical depth, visibility function,\n", " plus derived recombination quantities (z_*, tau_*, r_s, k_D, ...).\"\"\"\n", " z_arr, xe_arr, Tmat_rec = solve_recombination(bg, params)\n", "\n", " # Reionisation layer: CAMB tanh model\n", " f_He = bg['f_He']\n", " delta_z = 0.5\n", " he_reion_z = 3.5\n", " he_reion_dz = 0.4\n", " he_reion_zstart = he_reion_z + 5 * he_reion_dz\n", " f_re = 1.0 + f_He\n", " z_rev, xe_rev = z_arr[::-1], xe_arr[::-1]\n", "\n", " def build_reion_xe(z_eval, z_re):\n", " z_reion_start = z_re + 8 * delta_z\n", " x_e_freeze = np.interp(z_reion_start, z_rev, xe_rev)\n", " window_var_mid = (1 + z_re)**1.5\n", " window_var_delta = 1.5 * (1 + z_re)**0.5 * delta_z\n", " xod = np.clip((window_var_mid - (1 + z_eval)**1.5) / window_var_delta, -100, 100)\n", " x_H_reion = (f_re - x_e_freeze) * (np.tanh(xod) + 1) / 2 + x_e_freeze\n", " xod_he = np.clip((he_reion_z - z_eval) / he_reion_dz, -100, 100)\n", " x_he_extra = np.where(z_eval < he_reion_zstart,\n", " f_He * (np.tanh(xod_he) + 1) / 2, 0.0)\n", " return x_H_reion + x_he_extra\n", "\n", " def compute_reion_optical_depth(z_re):\n", " z_reion_start = z_re + 8 * delta_z\n", " def integrand(z):\n", " a = 1.0 / (1 + z)\n", " return build_reion_xe(z, z_re) * bg['akthom'] * dtau_da(a, bg)\n", " return integrate.quad(integrand, 0, z_reion_start, limit=200)[0]\n", "\n", " # Solve for z_re matching the target reionisation tau\n", " def tau_residual(z_re):\n", " return compute_reion_optical_depth(z_re) - params['tau_reion']\n", " z_re = optimize.brentq(tau_residual, 2.0, 30.0, xtol=1e-8, rtol=1e-8)\n", "\n", " # Refine the z grid around z_re\n", " z_reion_lo = max(0.01, z_re - 8 * delta_z)\n", " z_reion_hi = z_re + 8 * delta_z\n", " z_dense = np.linspace(z_reion_hi, z_reion_lo, 200)\n", " mask = (z_arr > z_reion_hi) | (z_arr < z_reion_lo)\n", " z_new = np.sort(np.concatenate([z_arr[mask], z_dense]))[::-1]\n", " xe_new = np.interp(z_new[::-1], z_arr[::-1], xe_arr[::-1])[::-1]\n", " Tmat_new = np.interp(z_new[::-1], z_arr[::-1], Tmat_rec[::-1])[::-1]\n", " z_arr, xe_arr, Tmat_rec = z_new, xe_new, Tmat_new\n", "\n", " # Final x_e: max of recombination + reionisation\n", " xe_final = np.maximum(build_reion_xe(z_arr, z_re), xe_arr)\n", "\n", " # Conformal time grid (monotonically increasing in tau)\n", " a_arr = 1.0 / (1 + z_arr)\n", " tau_arr = conformal_time(a_arr, bg)\n", "\n", " # Opacity, optical depth, visibility\n", " opacity = xe_final * bg['akthom'] / a_arr**2\n", " tau_optical = -np.flip(integrate.cumulative_trapezoid(\n", " np.flip(opacity), np.flip(tau_arr), initial=0))\n", " exptau = np.exp(-tau_optical)\n", " visibility = opacity * exptau\n", "\n", " # Assemble the thermo dictionary\n", " thermo = {\n", " 'z_arr': z_arr, 'a_arr': a_arr, 'tau_arr': tau_arr,\n", " 'xe': xe_final, 'Tmat': Tmat_rec,\n", " 'opacity': opacity, 'tau_optical': tau_optical,\n", " 'exptau': exptau, 'visibility': visibility,\n", " 'z_reion': z_re,\n", " }\n", " thermo['opacity_interp'] = interpolate.CubicSpline(tau_arr, opacity)\n", " thermo['exptau_interp'] = interpolate.CubicSpline(tau_arr, exptau)\n", " thermo['visibility_interp'] = interpolate.CubicSpline(tau_arr, visibility)\n", "\n", " # Baryon sound speed\n", " dlnT_dln_a = np.gradient(np.log(np.maximum(Tmat_rec, 1e-30)), np.log(a_arr))\n", " barssc = barssc0 * (1.0 - 0.75 * bg['Y_He'] + (1.0 - bg['Y_He']) * xe_arr)\n", " cs2_b = np.maximum(barssc * Tmat_rec * (1.0 - dlnT_dln_a / 3.0), 0.0)\n", " thermo['cs2_b'] = cs2_b\n", "\n", " # Peak of g(tau): the surface of last scattering\n", " peak_idx = np.argmax(visibility)\n", " thermo['z_star'] = z_arr[peak_idx]\n", " thermo['tau_star'] = tau_arr[peak_idx]\n", " a_star = 1.0 / (1.0 + thermo['z_star'])\n", "\n", " # Sound horizon r_s and Silk damping scale k_D\n", " thermo['r_s'] = sound_horizon(a_star, bg)\n", " a_grid = np.linspace(a_arr[0], a_star, 5000)\n", " R = 0.75 * bg['grhob'] * a_grid / bg['grhog']\n", " dtauda_grid = dtau_da(a_grid, bg)\n", " kappa_dot = np.maximum(np.interp(a_grid, a_arr, xe_final) * bg['akthom'] / a_grid**2, 1e-30)\n", " integrand_D = (R**2 + 16.0*(1.0+R)/15.0) / (6.0*(1.0+R)**2 * kappa_dot) * dtauda_grid\n", " thermo['k_D'] = 1.0 / np.sqrt(np.trapz(integrand_D, a_grid))\n", "\n", " # FWHM widths of the visibility peaks (needed by make_k_grid / make_tau_grid)\n", " fwhm_to_sigma = 2.0 * np.sqrt(2.0 * np.log(2.0))\n", " half_max = visibility[peak_idx] / 2.0\n", " tau_left = np.interp(half_max, visibility[:peak_idx+1], tau_arr[:peak_idx+1])\n", " tau_right = np.interp(half_max, visibility[peak_idx:][::-1], tau_arr[peak_idx:][::-1])\n", " thermo['delta_tau_rec'] = (tau_right - tau_left) / fwhm_to_sigma\n", "\n", " z_rev, tau_rev = z_arr[::-1], tau_arr[::-1]\n", " thermo['tau_reion'] = np.interp(z_re, z_rev, tau_rev)\n", " thermo['delta_tau_reion'] = abs(\n", " np.interp(z_re + 6*delta_z, z_rev, tau_rev)\n", " - np.interp(max(0.01, z_re - 6*delta_z), z_rev, tau_rev)) / fwhm_to_sigma\n", "\n", " return thermo\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: build the thermodynamics\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "th = build_thermodynamics(bg, params)\n", "print(f\"z_* = {th['z_star']:.0f}\")\n", "print(f\"tau_* = {th['tau_star']:.1f} Mpc\")\n", "print(f\"r_s(tau_*) = {th['r_s']:.1f} Mpc\")\n", "print(f\"k_D = {th['k_D']:.3f} Mpc^-1\")\n", "print(f\"z_reion = {th['z_reion']:.2f}\")\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot $x_e(z)$\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, ax = plt.subplots(figsize=(7, 4.5))\n", "sel = (th['z_arr'] > 1) & (th['z_arr'] < 1e4)\n", "ax.loglog(th['z_arr'][sel], th['xe'][sel], 'k-', lw=1.6)\n", "ax.invert_xaxis()\n", "ax.set_xlim(1e4, 1)\n", "ax.set_xlabel(r'Redshift $z$')\n", "ax.set_ylabel(r'Free electron fraction $x_e$')\n", "ax.set_title('Recombination + reionisation history')\n", "ax.grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot $g(tau)$ and $dotkappa(tau)$\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, axes = plt.subplots(2, 1, sharex=True, figsize=(7, 6))\n", "sel = th['tau_arr'] > 0\n", "axes[0].semilogx(th['tau_arr'][sel], th['visibility'][sel], 'k-', lw=1.4)\n", "axes[0].set_ylabel(r'Visibility $g(\\tau)$')\n", "axes[0].set_title('Visibility function and Thomson opacity')\n", "axes[0].grid(True, which='both', alpha=0.3)\n", "axes[1].loglog(th['tau_arr'][sel], th['opacity'][sel], 'k-', lw=1.4)\n", "axes[1].set_ylabel(r'$\\dot\\kappa\\, [\\mathrm{Mpc}^{-1}]$')\n", "axes[1].set_xlabel(r'Conformal time $\\tau\\, [\\mathrm{Mpc}]$')\n", "axes[1].grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 3 \u2014 Perturbations\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: state-vector layout and hierarchy parameters\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Hierarchy truncation (increase for higher ell_max accuracy)\n", "LMAXG = 15 # photon temperature: Theta_0 ... Theta_LMAXG\n", "LMAXPOL = 15 # photon polarisation: E_2 ... E_LMAXPOL\n", "LMAXNR = 15 # massless neutrinos: N_0 ... N_LMAXNR\n", "\n", "# Indices into the flat state vector\n", "IX_ETAK = 0\n", "IX_CLXC = 1\n", "IX_CLXB = 2\n", "IX_VB = 3\n", "IX_G = 4 # Theta_0 at IX_G, Theta_1 at IX_G+1, ...\n", "IX_POL = IX_G + LMAXG + 1 # E_2 at IX_POL, E_3 at IX_POL+1, ...\n", "IX_R = IX_POL + LMAXPOL - 1 # N_0 at IX_R, N_1 at IX_R+1, ...\n", "NVAR = IX_R + LMAXNR + 1\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: cubic-spline evaluator\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "@_jit\n", "def _cubic_eval(x_knots, coeffs, t):\n", " \"\"\"Evaluate a scipy.interpolate.CubicSpline at point t.\"\"\"\n", " n = x_knots.shape[0] - 1\n", " lo, hi = 0, n - 1\n", " while lo < hi:\n", " mid = (lo + hi) >> 1\n", " if x_knots[mid + 1] < t:\n", " lo = mid + 1\n", " else:\n", " hi = mid\n", " dt = t - x_knots[lo]\n", " return ((coeffs[0, lo] * dt + coeffs[1, lo]) * dt\n", " + coeffs[2, lo]) * dt + coeffs[3, lo]\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `init_perturbation_grid`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def init_perturbation_grid(bg, thermo):\n", " \"\"\"Precompute a(tau) and background quantities on a fine conformal-time grid.\"\"\"\n", " # Build tau(a) by cumulative integration of dtau/da\n", " a_grid = np.logspace(-9, 0, 10000)\n", " dtauda_grid = np.array([dtau_da(a, bg) for a in a_grid])\n", " tau_grid = integrate.cumulative_trapezoid(dtauda_grid, a_grid, initial=0)\n", " a_of_tau = interpolate.CubicSpline(tau_grid, a_grid)\n", "\n", " # Radiation-era expansion rate\n", " grho_rad = bg['grhog'] + bg['grhornomass']\n", " adotrad = np.sqrt(grho_rad / 3.0)\n", "\n", " # Extend opacity, sound speed below the thermo grid (full ionisation)\n", " tau_thermo = thermo['tau_arr']\n", " tau_early = tau_grid[tau_grid < tau_thermo[0]]\n", " a_early = a_of_tau(tau_early)\n", " tau_ext = np.concatenate([tau_early, tau_thermo])\n", "\n", " opac_early = (1.0 + bg['f_He']) * bg['akthom'] / a_early**2\n", " opacity_interp = interpolate.CubicSpline(\n", " tau_ext, np.concatenate([opac_early, thermo['opacity']]))\n", "\n", " xe_early = 1.0 + 2.0 * bg['f_He']\n", " barssc_early = barssc0 * (1.0 - 0.75 * bg['Y_He']\n", " + (1.0 - bg['Y_He']) * xe_early)\n", " cs2_early = (4.0 / 3.0) * barssc_early * bg['T_cmb'] / a_early\n", " cs2_interp = interpolate.CubicSpline(\n", " tau_ext, np.concatenate([cs2_early, thermo['cs2_b']]))\n", "\n", " return {\n", " 'sp_a_x': a_of_tau.x, 'sp_a_c': a_of_tau.c,\n", " 'sp_op_x': opacity_interp.x, 'sp_op_c': opacity_interp.c,\n", " 'sp_cs_x': cs2_interp.x, 'sp_cs_c': cs2_interp.c,\n", " 'bg_vec': np.array([bg['grhog'], bg['grhornomass'], bg['grhoc'],\n", " bg['grhob'], bg['grhov']]),\n", " 'adotrad': adotrad, 'grho_rad': grho_rad, 'tau0': bg['tau0'],\n", " }\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `adiabatic_ics`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def adiabatic_ics(k, tau_start, bg, pgrid):\n", " \"\"\"Adiabatic initial conditions on super-horizon scales (k*tau << 1).\"\"\"\n", " x = k * tau_start\n", " x2 = x * x\n", "\n", " grho_rad = pgrid['grho_rad']\n", " Rv = bg['grhornomass'] / grho_rad # rho_nu/(rho_nu + rho_gamma)\n", " Rp15 = 4 * Rv + 15\n", "\n", " om = (bg['grhob'] + bg['grhoc']) / np.sqrt(3.0 * grho_rad)\n", " omtau = om * tau_start\n", "\n", " y0 = np.zeros(NVAR)\n", "\n", " # Metric variable etak = k * eta\n", " y0[IX_ETAK] = -k * (1.0 - x2 / 12.0 * (-10.0 / Rp15 + 1.0))\n", "\n", " # Photon monopole and dipole\n", " clxg_init = x2 / 3.0 * (1.0 - omtau / 5.0)\n", " qg_init = x2 * x / 27.0 * (1.0 - omtau / 5.0)\n", " y0[IX_G] = clxg_init\n", " y0[IX_G + 1] = qg_init\n", "\n", " # CDM and baryon density: delta = 3/4 delta_gamma for adiabatic mode\n", " y0[IX_CLXC] = 0.75 * clxg_init\n", " y0[IX_CLXB] = 0.75 * clxg_init\n", " y0[IX_VB] = 0.75 * qg_init\n", "\n", " # Massless neutrinos\n", " y0[IX_R] = clxg_init\n", " y0[IX_R + 1] = (4 * Rv + 23) / Rp15 * x2 * x / 27.0\n", " y0[IX_R + 2] = -4.0 / 3.0 * x2 / Rp15 * (1.0 + omtau / 4.0\n", " * (4*Rv - 5) / (2*Rv + 15))\n", " if LMAXNR >= 3:\n", " y0[IX_R + 3] = -4.0 / 21.0 / Rp15 * x2 * x\n", "\n", " # All higher multipoles and polarisation start at zero\n", " return y0\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: common background andEinstein terms (helper)\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "@_jit\n", "def _common_terms(tau, y, k, bg_vec, sp_a_x, sp_a_c):\n", " \"\"\"Background and Einstein-source terms used by both RHS and source builders.\"\"\"\n", " grhog, grhornomass, grhoc, grhob, grhov = bg_vec\n", " a = _cubic_eval(sp_a_x, sp_a_c, tau)\n", " a2 = a * a\n", " grhog_t = grhog / a2\n", " grhor_t = grhornomass / a2\n", " grhoc_t = grhoc / a\n", " grhob_t = grhob / a\n", " grho_a2 = grhog_t + grhor_t + grhoc_t + grhob_t + grhov * a2\n", " adotoa = np.sqrt(grho_a2 / 3.0)\n", "\n", " etak = y[IX_ETAK]\n", " clxc = y[IX_CLXC]; clxb = y[IX_CLXB]; vb = y[IX_VB]\n", " clxg = y[IX_G]; qg = y[IX_G + 1]; pig = y[IX_G + 2]\n", " clxr = y[IX_R]; qr = y[IX_R + 1]; pir = y[IX_R + 2]\n", "\n", " k2 = k * k\n", " dgrho = grhob_t*clxb + grhoc_t*clxc + grhog_t*clxg + grhor_t*clxr\n", " dgq = grhob_t*vb + grhog_t*qg + grhor_t*qr\n", " z = (0.5 * dgrho / k + etak) / adotoa\n", " sigma = z + 1.5 * dgq / k2\n", "\n", " return (a, adotoa, grhog_t, grhor_t, grhoc_t, grhob_t,\n", " dgrho, dgq, z, sigma,\n", " etak, clxc, clxb, vb, clxg, qg, pig, clxr, qr, pir)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `boltzmann_derivs` -- the Boltzmann RHS\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "@_jit\n", "def boltzmann_derivs(tau, y, k, bg_vec, sp_a_x, sp_a_c,\n", " sp_op_x, sp_op_c, sp_cs_x, sp_cs_c):\n", " \"\"\"dy/d(tau): the full Einstein-Boltzmann system in synchronous gauge.\"\"\"\n", " (a, adotoa, grhog_t, grhor_t, grhoc_t, grhob_t,\n", " dgrho, dgq, z, sigma,\n", " etak, clxc, clxb, vb, clxg, qg, pig, clxr, qr, pir) = \\\n", " _common_terms(tau, y, k, bg_vec, sp_a_x, sp_a_c)\n", " opacity = max(_cubic_eval(sp_op_x, sp_op_c, tau), 1e-30)\n", " cs2_b = max(_cubic_eval(sp_cs_x, sp_cs_c, tau), 0.0)\n", " photbar = grhog_t / grhob_t\n", " pb43 = 4.0 / 3.0 * photbar\n", " delta_p_b = cs2_b * clxb\n", "\n", " tight_coupling = (k / opacity < 0.01) and (1.0 / (opacity * tau) < 0.01)\n", " E2 = y[IX_POL] if LMAXPOL >= 2 else 0.0\n", " cothxor = 1.0 / tau\n", "\n", " dy = np.zeros(NVAR)\n", "\n", " # --- Metric: dot{eta} k = (1/2) dgq ---\n", " dy[IX_ETAK] = 0.5 * dgq\n", " # --- CDM (at rest in this gauge) ---\n", " dy[IX_CLXC] = -k * z\n", " # --- Baryons: continuity ---\n", " dy[IX_CLXB] = -k * (z + vb)\n", "\n", " if tight_coupling:\n", " # Tight coupling: photon-baryon fluid locked together\n", " pig_tc = 32.0 / 45.0 * k / opacity * (sigma + vb)\n", " polter = pig_tc / 4.0\n", "\n", " vbdot = (-adotoa * vb + k * delta_p_b\n", " + k / 4.0 * pb43 * (clxg - 2.0 * pig_tc)) / (1.0 + pb43)\n", " dy[IX_VB] = vbdot\n", "\n", " dy[IX_G] = -k * (4.0 / 3.0 * z + qg)\n", " qgdot = 4.0 / 3.0 * (-vbdot - adotoa * vb\n", " + k * delta_p_b) / pb43 + k / 3.0 * clxg \\\n", " - 2.0 * k / 3.0 * pig_tc\n", " dy[IX_G + 1] = qgdot\n", " dy[IX_G + 2] = opacity * (pig_tc - pig)\n", "\n", " if LMAXPOL >= 2:\n", " dy[IX_POL] = opacity * (pig_tc / 4.0 - E2)\n", "\n", " else:\n", " # Full Boltzmann hierarchy\n", " polter = pig / 10.0 + 9.0 / 15.0 * E2\n", " vbdot = -adotoa * vb + k * delta_p_b \\\n", " - photbar * opacity * (4.0 / 3.0 * vb - qg)\n", " dy[IX_VB] = vbdot\n", "\n", " dy[IX_G] = -k * (4.0 / 3.0 * z + qg)\n", " qgdot = 4.0 / 3.0 * (-vbdot - adotoa * vb + k * delta_p_b) / pb43 \\\n", " + k / 3.0 * clxg - 2.0 * k / 3.0 * pig\n", " dy[IX_G + 1] = qgdot\n", "\n", " Theta3 = y[IX_G + 3] if LMAXG >= 3 else 0.0\n", " dy[IX_G + 2] = (2.0 * k / 5.0 * qg - 3.0 * k / 5.0 * Theta3\n", " - opacity * (pig - polter) + 8.0 / 15.0 * k * sigma)\n", "\n", " for l in range(3, LMAXG):\n", " dy[IX_G + l] = (k * l / (2*l + 1) * y[IX_G + l - 1]\n", " - k * (l + 1) / (2*l + 1) * y[IX_G + l + 1]\n", " - opacity * y[IX_G + l])\n", " # Free-streaming closure\n", " dy[IX_G + LMAXG] = (k * y[IX_G + LMAXG - 1]\n", " - (LMAXG + 1) * cothxor * y[IX_G + LMAXG]\n", " - opacity * y[IX_G + LMAXG])\n", "\n", " # E-mode polarisation\n", " E3 = y[IX_POL + 1] if LMAXPOL >= 3 else 0.0\n", " dy[IX_POL] = -opacity * (E2 - polter) - k / 3.0 * E3\n", "\n", " for l in range(3, LMAXPOL):\n", " idx = IX_POL + l - 2\n", " polfac_l = (l + 3) * (l - 1) / (l + 1)\n", " dy[idx] = (-opacity * y[idx]\n", " + k * l / (2*l + 1) * y[idx - 1]\n", " - polfac_l * k / (2*l + 1) * y[idx + 1])\n", "\n", " idx_last = IX_POL + LMAXPOL - 2\n", " dy[idx_last] = (-opacity * y[idx_last]\n", " + k * LMAXPOL / (2*LMAXPOL + 1) * y[idx_last - 1]\n", " - (LMAXPOL + 3) * cothxor * y[idx_last])\n", "\n", " # --- Massless neutrinos ---\n", " dy[IX_R] = -k * (4.0 / 3.0 * z + qr)\n", " dy[IX_R + 1] = k / 3.0 * (clxr - 2.0 * pir)\n", " N3 = y[IX_R + 3] if LMAXNR >= 3 else 0.0\n", " dy[IX_R + 2] = 2.0 * k / 5.0 * qr - 3.0 * k / 5.0 * N3 + 8.0 / 15.0 * k * sigma\n", "\n", " for l in range(3, LMAXNR):\n", " dy[IX_R + l] = (k * l / (2*l + 1) * y[IX_R + l - 1]\n", " - k * (l + 1) / (2*l + 1) * y[IX_R + l + 1])\n", " dy[IX_R + LMAXNR] = (k * y[IX_R + LMAXNR - 1]\n", " - (LMAXNR + 1) * cothxor * y[IX_R + LMAXNR])\n", "\n", " return dy\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `evolve_k`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def evolve_k(k, bg, thermo, pgrid, tau_out):\n", " \"\"\"Evolve perturbations for a single wavenumber k from tau_start to tau_out[-1].\n", " Returns sol.y of shape (NVAR, len(tau_out)).\"\"\"\n", " tau_start = min(0.01 / k, tau_out[0] * 0.5)\n", " tau_start = max(tau_start, 0.1)\n", " y0 = adiabatic_ics(k, tau_start, bg, pgrid)\n", "\n", " bg_vec = pgrid['bg_vec']\n", " sp_a_x = pgrid['sp_a_x']; sp_a_c = pgrid['sp_a_c']\n", " sp_op_x = pgrid['sp_op_x']; sp_op_c = pgrid['sp_op_c']\n", " sp_cs_x = pgrid['sp_cs_x']; sp_cs_c = pgrid['sp_cs_c']\n", "\n", " sol = integrate.solve_ivp(\n", " lambda tau, y: boltzmann_derivs(\n", " tau, y, k, bg_vec, sp_a_x, sp_a_c, sp_op_x, sp_op_c, sp_cs_x, sp_cs_c),\n", " [tau_start, tau_out[-1]], y0, t_eval=tau_out,\n", " method='LSODA', rtol=1e-5, atol=1e-8, max_step=20.0)\n", " if not sol.success:\n", " raise RuntimeError(f\"ODE solver failed for k={k:.4e}: {sol.message}\")\n", " return sol.y\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: evolve one mode\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "pgrid = init_perturbation_grid(bg, th)\n", "k = 0.05\n", "tau_out = np.geomspace(0.5, bg['tau0'] * 0.999, 600)\n", "Y = evolve_k(k, bg, th, pgrid, tau_out)\n", "\n", "# Extract individual variables\n", "etak = Y[IX_ETAK]\n", "clxc = Y[IX_CLXC]; clxb = Y[IX_CLXB]; vb = Y[IX_VB]\n", "clxg = Y[IX_G]; qg = Y[IX_G + 1]; pig = Y[IX_G + 2]\n", "\n", "# Scale factor at each tau (for the x-axis)\n", "a_of_tau = np.array([_cubic_eval(pgrid['sp_a_x'], pgrid['sp_a_c'], t)\n", " for t in tau_out])\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot the perturbations\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, axes = plt.subplots(3, 1, sharex=True, figsize=(7, 8))\n", "m = (a_of_tau > 1e-6) & (a_of_tau < 1.0)\n", "\n", "axes[0].semilogx(a_of_tau[m], clxg[m], label=r'$\\delta_\\gamma$', color='C0')\n", "axes[0].semilogx(a_of_tau[m], clxc[m], label=r'$\\delta_c$', color='C1')\n", "axes[0].semilogx(a_of_tau[m], clxb[m], label=r'$\\delta_b$', color='C2')\n", "axes[0].set_ylabel('Density contrast')\n", "axes[0].set_title(rf'Perturbation evolution at $k = {k}$ Mpc$^{{-1}}$')\n", "axes[0].legend(loc='upper left'); axes[0].grid(True, alpha=0.3)\n", "axes[0].set_ylim(-3, 8)\n", "\n", "axes[1].semilogx(a_of_tau[m], qg[m] * 3 / 4, label=r'$v_\\gamma$', color='C0')\n", "axes[1].semilogx(a_of_tau[m], vb[m], label=r'$v_b$', color='C2')\n", "axes[1].set_ylabel('Velocities')\n", "axes[1].legend(); axes[1].grid(True, alpha=0.3)\n", "\n", "axes[2].semilogx(a_of_tau[m], pig[m], 'C0-', label=r'$\\pi_\\gamma$')\n", "ax2 = axes[2].twinx()\n", "ax2.semilogx(a_of_tau[m], etak[m] / k, 'C1-', label=r'$\\eta$')\n", "axes[2].set_ylabel(r'$\\pi_\\gamma$', color='C0')\n", "ax2.set_ylabel(r'$\\eta$ (metric)', color='C1')\n", "axes[2].set_xlabel('Scale factor $a$')\n", "axes[2].grid(True, alpha=0.3)\n", "plt.tight_layout(); plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 4 \u2014 Line of sight\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `build_sources`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def build_sources(tau, y, k, pgrid, thermo):\n", " \"\"\"Source function building blocks at a single (k, tau) point.\n", "\n", " Returns (ISW, monopole, sigma+vb, vis, polter, source_E) -- the\n", " raw pieces that get combined into the three temperature Bessel\n", " channels and the E-mode source.\n", " \"\"\"\n", " (a, adotoa, grhog_t, grhor_t, grhoc_t, grhob_t,\n", " dgrho, dgq, z, sigma,\n", " etak, clxc, clxb, vb, clxg, qg, pig, clxr, qr, pir) = _common_terms(\n", " tau, y, k, pgrid['bg_vec'], pgrid['sp_a_x'], pgrid['sp_a_c'])\n", " opacity = max(float(_cubic_eval(pgrid['sp_op_x'], pgrid['sp_op_c'], tau)),\n", " 1e-30)\n", " k2 = k * k\n", "\n", " E2 = y[IX_POL] if LMAXPOL >= 2 else 0.0\n", "\n", " # --- Weyl potential phi = (Phi + Psi)/2 in synchronous gauge ---\n", " dgpi = grhog_t * pig + grhor_t * pir\n", " phi = -((dgrho + 3.0 * dgq * adotoa / k) + dgpi) / (2.0 * k2)\n", "\n", " # --- Compute phi_dot directly from hierarchy equations ---\n", " polter = pig / 10.0 + 9.0 / 15.0 * E2\n", " Theta3 = y[IX_G + 3] if LMAXG >= 3 else 0.0\n", " N3 = y[IX_R + 3] if LMAXNR >= 3 else 0.0\n", " pigdot = (2*k/5*qg - 3*k/5*Theta3 - opacity*(pig - polter) + 8*k*sigma/15)\n", " pirdot = (2*k/5*qr - 3*k/5*N3 + 8*k*sigma/15)\n", " pidot_sum = grhog_t * pigdot + grhor_t * pirdot\n", " diff_rhopi = pidot_sum - 4.0 * adotoa * dgpi\n", " gpres_plus_grho = (4.0/3.0) * (grhog_t + grhor_t) + grhoc_t + grhob_t\n", " phidot = 0.5 * (adotoa * (-dgpi - 2.0 * k2 * phi) + dgq * k\n", " - diff_rhopi + k * sigma * gpres_plus_grho) / k2\n", "\n", " # --- Visibility and exp(-tau) at this time ---\n", " tau_min, tau_max = thermo['tau_arr'][0], thermo['tau_arr'][-1]\n", " if tau < tau_min or tau > tau_max:\n", " vis = 0.0\n", " exptau = np.exp(-thermo['tau_optical'][0]) if tau < tau_min else 1.0\n", " else:\n", " vis = float(thermo['visibility_interp'](tau))\n", " exptau = float(thermo['exptau_interp'](tau))\n", "\n", " # --- Three source components ---\n", " ISW = 2.0 * phidot * exptau # 2 phi_dot e^{-kappa}\n", " monopole = -etak / k + 2.0 * phi + clxg / 4.0 # Theta_0 + Psi_W + (3/8) delta_b...\n", " chi = pgrid['tau0'] - tau\n", " source_E = 15.0 / 8.0 * vis * polter / (chi**2 * k2) if chi > 0 else 0.0\n", "\n", " return ISW, monopole, sigma + vb, vis, polter, source_E\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `evolve_k` (source-function version)\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def evolve_k(k, bg, thermo, pgrid, tau_out):\n", " \"\"\"Evolve perturbations for wavenumber k, return source functions on tau_out.\n", " Returns (src_j0, src_j1, src_j2, src_E).\"\"\"\n", " tau_start = min(0.01 / k, tau_out[0] * 0.5)\n", " tau_start = max(tau_start, 0.1)\n", " y0 = adiabatic_ics(k, tau_start, bg, pgrid)\n", "\n", " bg_vec = pgrid['bg_vec']\n", " sp_a_x = pgrid['sp_a_x']; sp_a_c = pgrid['sp_a_c']\n", " sp_op_x = pgrid['sp_op_x']; sp_op_c = pgrid['sp_op_c']\n", " sp_cs_x = pgrid['sp_cs_x']; sp_cs_c = pgrid['sp_cs_c']\n", "\n", " sol = integrate.solve_ivp(\n", " lambda tau, y: boltzmann_derivs(\n", " tau, y, k, bg_vec, sp_a_x, sp_a_c, sp_op_x, sp_op_c, sp_cs_x, sp_cs_c),\n", " [tau_start, tau_out[-1]], y0, t_eval=tau_out,\n", " method='LSODA', rtol=1e-5, atol=1e-8, max_step=20.0)\n", " if not sol.success:\n", " raise RuntimeError(f\"ODE solver failed for k={k:.4e}: {sol.message}\")\n", "\n", " ntau = len(tau_out)\n", " ISW_arr = np.zeros(ntau)\n", " monopole_arr = np.zeros(ntau)\n", " sigma_vb_arr = np.zeros(ntau)\n", " vis_arr = np.zeros(ntau)\n", " polter_arr = np.zeros(ntau)\n", " src_E = np.zeros(ntau)\n", " for i, tau in enumerate(tau_out):\n", " (ISW_arr[i], monopole_arr[i], sigma_vb_arr[i],\n", " vis_arr[i], polter_arr[i], src_E[i]) = build_sources(\n", " tau, sol.y[:, i], k, pgrid, thermo)\n", "\n", " # Assemble the three temperature channels\n", " src_j0 = ISW_arr + vis_arr * (monopole_arr + 0.625 * polter_arr)\n", " src_j1 = vis_arr * sigma_vb_arr\n", " src_j2 = 1.875 * vis_arr * polter_arr\n", " return src_j0, src_j1, src_j2, src_E\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 5 \u2014 Power spectra\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `make_k_grid`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def make_k_grid(N, mode, bg, thermo, params,\n", " k_min=1e-5, k_max=0.5,\n", " ell_min=2, ell_max=2500, n_ell_samples=30,\n", " n_eval=5000):\n", " \"\"\"Optimal non-uniform k-grid.\n", " mode='cl' : optimised for the C_ell integral\n", " mode='ode' : optimised for source-function interpolation.\n", " \"\"\"\n", " x = np.linspace(np.log(k_min), np.log(k_max), n_eval)\n", " k = np.exp(x)\n", " primordial = k ** (params['n_s'] + 2)\n", " acoustic_curv = (1.0 / thermo['r_s']) ** 2\n", "\n", " if mode == \"cl\":\n", " damped = primordial * np.exp(-1.0 * (k / thermo['k_D']) ** 2)\n", " sigma_k = 1.0 / thermo['delta_tau_rec']\n", " chi_star = bg['tau0'] - thermo['tau_star']\n", " ells = np.unique(np.geomspace(ell_min, ell_max,\n", " n_ell_samples).astype(int))\n", " raw_weight = np.zeros_like(k)\n", " for ell in ells:\n", " envelope = np.exp(-0.5 * ((k - ell / chi_star) / (3.0 * sigma_k)) ** 2)\n", " curv = np.maximum(acoustic_curv, sigma_k ** 2) * envelope\n", " raw_weight += curv * damped\n", " floor = 1e-6 * np.max(raw_weight)\n", " else:\n", " damped = primordial * np.exp(-1.0 * (k / thermo['k_D']) ** 2)\n", " smooth_curv = (k * thermo['r_s']) ** 2 / bg['tau_eq'] ** 2\n", " raw_weight = np.maximum(acoustic_curv, smooth_curv) * damped\n", " floor = 0.005 * np.max(raw_weight)\n", "\n", " density = (raw_weight + floor) ** (1.0 / 3.0)\n", " dx = x[1] - x[0]\n", " cdf = np.cumsum(density) * dx\n", " cdf -= cdf[0]; cdf /= cdf[-1]\n", " grid = np.exp(np.interp(np.linspace(0, 1, N), cdf, x))\n", " grid[0] = k_min; grid[-1] = k_max\n", " return grid\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `make_tau_grid`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def make_tau_grid(N, k_max, bg, thermo,\n", " tau_min=1.0, tau_max=None, n_eval=10000):\n", " \"\"\"Optimal non-uniform tau-grid for the line-of-sight integral.\"\"\"\n", " if tau_max is None:\n", " tau_max = bg['tau0']\n", " tau = np.linspace(tau_min, tau_max, n_eval)\n", " tau_star = thermo['tau_star']\n", " delta_tau_rec = thermo['delta_tau_rec']\n", "\n", " g_rec = np.exp(-0.5 * ((tau - tau_star) / delta_tau_rec) ** 2)\n", " g_broad = np.exp(-0.5 * ((tau - tau_star) / thermo['r_s']) ** 2)\n", " weight = g_rec / delta_tau_rec**2 + g_broad * (k_max / np.sqrt(3.0))**2\n", " g_reion = np.exp(-0.5 * ((tau - thermo['tau_reion'])\n", " / thermo['delta_tau_reion']) ** 2)\n", " weight += 0.3 * g_reion / thermo['delta_tau_reion'] ** 2\n", "\n", " density = (weight + 0.005 * np.max(weight)) ** (1.0 / 3.0)\n", " dtau = tau[1] - tau[0]\n", " cdf = np.cumsum(density) * dtau\n", " cdf -= cdf[0]; cdf /= cdf[-1]\n", " grid = np.interp(np.linspace(0, 1, N), cdf, tau)\n", " grid[0] = tau_min; grid[-1] = tau_max\n", " return grid\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: Bessel tables\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def _build_bessel_tables(ells_compute, x_max, dx):\n", " \"\"\"Precompute j_l and j_{l+1} on a uniform x-grid.\"\"\"\n", " n_x = int(np.ceil(x_max / dx)) + 1\n", " x_lo = 0.0\n", " x_arr = x_lo + dx * np.arange(n_x)\n", " inv_dx = 1.0 / dx\n", " nell = len(ells_compute)\n", " jl_tab = np.zeros((nell, n_x))\n", " jl1_tab = np.zeros((nell, n_x))\n", " for il, ell in enumerate(ells_compute):\n", " jl_tab[il] = special.spherical_jn(ell, x_arr)\n", " jl1_tab[il] = special.spherical_jn(ell + 1, x_arr)\n", " return x_lo, inv_dx, n_x, jl_tab, jl1_tab\n", "\n", "@_jit\n", "def _interp_uniform_table(x, x0, inv_dx, n_x, vals):\n", " \"\"\"Linear interpolation on a uniform grid.\"\"\"\n", " out = np.zeros_like(x)\n", " flat = x.ravel()\n", " flat_out = out.ravel()\n", " for i in range(flat.size):\n", " u = (flat[i] - x0) * inv_dx\n", " j = int(u)\n", " if 0 <= j < n_x - 1:\n", " t = u - j\n", " flat_out[i] = (1 - t) * vals[j] + t * vals[j + 1]\n", " return out\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `compute_spectra`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def compute_spectra(bg, thermo, params, N_k_ode=200, N_k_fine=4000, N_tau=1000,\n", " k_arr=None, k_fine=None, tau_out=None):\n", " \"\"\"Main CMB pipeline: evolve all k modes, do LOS integration, assemble C_ell.\"\"\"\n", " pgrid = init_perturbation_grid(bg, thermo)\n", " tau0 = bg['tau0']\n", " tau_star = thermo['tau_star']\n", "\n", " if k_arr is None:\n", " k_arr = make_k_grid(N_k_ode, \"ode\", bg, thermo, params)\n", " nk = len(k_arr)\n", " if k_fine is None:\n", " k_fine = make_k_grid(N_k_fine, \"cl\", bg, thermo, params,\n", " k_min=k_arr[0], k_max=k_arr[-1])\n", " if tau_out is None:\n", " tau_out = make_tau_grid(N_tau, k_arr[-1], bg, thermo,\n", " tau_min=1.0, tau_max=tau0 - 1)\n", " ntau = len(tau_out)\n", "\n", " # Warm up numba JIT (no-op without numba), then evolve all modes sequentially.\n", " # A parallel multiprocessing version is in nanocmb.py if you need it.\n", " boltzmann_derivs(tau_out[0], np.zeros(NVAR), k_arr[0],\n", " pgrid['bg_vec'], pgrid['sp_a_x'], pgrid['sp_a_c'],\n", " pgrid['sp_op_x'], pgrid['sp_op_c'],\n", " pgrid['sp_cs_x'], pgrid['sp_cs_c'])\n", " results = [evolve_k(k, bg, thermo, pgrid, tau_out) for k in k_arr]\n", "\n", " sources_j0 = np.array([r[0] for r in results])\n", " sources_j1 = np.array([r[1] for r in results])\n", " sources_j2 = np.array([r[2] for r in results])\n", " sources_E = np.array([r[3] for r in results])\n", "\n", " # Akima interpolation onto fine k-grid\n", " nk_fine = len(k_fine)\n", " lnk_ode = np.log(k_arr)\n", " lnk_fine = np.log(k_fine)\n", " src_fine_j0 = np.zeros((nk_fine, ntau))\n", " src_fine_j1 = np.zeros((nk_fine, ntau))\n", " src_fine_j2 = np.zeros((nk_fine, ntau))\n", " src_fine_E = np.zeros((nk_fine, ntau))\n", " for it in range(ntau):\n", " for src_in, src_out in [(sources_j0, src_fine_j0), (sources_j1, src_fine_j1),\n", " (sources_j2, src_fine_j2), (sources_E, src_fine_E)]:\n", " src_out[:, it] = interpolate.Akima1DInterpolator(lnk_ode, src_in[:, it])(lnk_fine)\n", "\n", " # ell-grid: dense at low ell, sparser at high ell\n", " ell_max = params['ell_max']\n", " ells_compute = np.unique(np.concatenate([\n", " np.arange(2, 40, 2),\n", " np.arange(40, 200, 5),\n", " np.arange(200, ell_max + 1, 50),\n", " ]))\n", " ells_compute = ells_compute[ells_compute <= ell_max]\n", " nell = len(ells_compute)\n", "\n", " chi_arr = tau0 - tau_out\n", " chi_star = tau0 - tau_star\n", " chi_max = chi_arr.max()\n", "\n", " Delta_T = np.zeros((nell, nk_fine))\n", " Delta_E = np.zeros((nell, nk_fine))\n", " x_2d_full = k_fine[:, None] * chi_arr[None, :]\n", "\n", " x0_tab, inv_dx_tab, n_x_tab, jl_tab, jl1_tab = _build_bessel_tables(\n", " ells_compute, float(np.max(x_2d_full)) + 2.0, 0.03)\n", "\n", " def _compute_ell_transfer(il, ell):\n", " # Restrict the k range to where j_ell has support\n", " x_lo = max(0.0, ell - 4.0 * ell**(1.0/3.0))\n", " k_lo = x_lo / chi_max if chi_max > 0 else 0\n", " k_hi = (ell + 2500) / chi_star if chi_star > 0 else k_fine[-1]\n", " ik_lo = max(0, np.searchsorted(k_fine, k_lo) - 1)\n", " ik_hi = min(nk_fine, np.searchsorted(k_fine, k_hi) + 1)\n", "\n", " x_2d = x_2d_full[ik_lo:ik_hi, :]\n", " jl = _interp_uniform_table(x_2d, x0_tab, inv_dx_tab, n_x_tab, jl_tab[il])\n", " jl_next = _interp_uniform_table(x_2d, x0_tab, inv_dx_tab, n_x_tab, jl1_tab[il])\n", "\n", " with np.errstate(divide='ignore', invalid='ignore'):\n", " inv_x = np.where(x_2d > 1e-30, 1.0 / x_2d, 0.0)\n", " jl_d = np.where(x_2d > 1e-30, ell * inv_x * jl - jl_next, 0.0)\n", " jl_dd = np.where(x_2d > 1e-30,\n", " -2.0 * inv_x * jl_d\n", " + (ell * (ell + 1) * inv_x * inv_x - 1.0) * jl,\n", " 0.0)\n", " integrand_T = (src_fine_j0[ik_lo:ik_hi, :] * jl\n", " + src_fine_j1[ik_lo:ik_hi, :] * jl_d\n", " + src_fine_j2[ik_lo:ik_hi, :] * jl_dd)\n", " integrand_E = src_fine_E[ik_lo:ik_hi, :] * jl\n", " Delta_T[il, ik_lo:ik_hi] = np.trapz(integrand_T, tau_out, axis=1)\n", " Delta_E[il, ik_lo:ik_hi] = np.trapz(integrand_E, tau_out, axis=1)\n", "\n", " for il, ell in enumerate(ells_compute):\n", " _compute_ell_transfer(il, int(ell))\n", "\n", " # Primordial power and assembly: C_l = 4 pi int dlnk P(k) Delta_T^2\n", " Pk = params['A_s'] * (k_fine / params['k_pivot'])**(params['n_s'] - 1.0)\n", " Cl_TT, Cl_EE, Cl_TE = [np.trapz(Pk * d, lnk_fine, axis=1)\n", " for d in (Delta_T**2, Delta_E**2, Delta_T * Delta_E)]\n", "\n", " # Normalise to D_l = l(l+1) C_l / (2 pi) (E mode has extra spin-2 factor)\n", " ells_f = ells_compute.astype(float)\n", " norm = 4.0 * np.pi * ells_f * (ells_f + 1) / (2.0 * np.pi)\n", " ctnorm = (ells_f**2 - 1.0) * (ells_f + 2) * ells_f\n", " Cl_TT *= norm\n", " Cl_EE *= norm * ctnorm\n", " Cl_TE *= norm * np.sqrt(ctnorm)\n", "\n", " # Convert dimensionless (dT/T)^2 to mu K^2\n", " T0_muK2 = (params['T_cmb'] * 1e6) ** 2\n", " Cl_TT *= T0_muK2; Cl_EE *= T0_muK2; Cl_TE *= T0_muK2\n", "\n", " # Interpolate to all integer ell\n", " ells_all = np.arange(2, ell_max + 1)\n", " Dl_TT, Dl_EE, Dl_TE = [interpolate.CubicSpline(ells_compute, cl)(ells_all)\n", " for cl in (Cl_TT, Cl_EE, Cl_TE)]\n", "\n", " return {\n", " 'ells': ells_all, 'Dl_TT': Dl_TT, 'Dl_EE': Dl_EE, 'Dl_TE': Dl_TE,\n", " 'k_fine': k_fine, 'ells_compute': ells_compute,\n", " 'Delta_T': Delta_T, 'Delta_E': Delta_E,\n", " }\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: compute the scalar spectra\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "result = compute_spectra(bg, th, params)\n", "ells = result['ells']\n", "Dl_TT = result['Dl_TT']\n", "Dl_EE = result['Dl_EE']\n", "Dl_TE = result['Dl_TE']\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot the scalar spectra\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, ax = plt.subplots(figsize=(8, 5))\n", "ax.plot(ells, Dl_TT, color='C0', label='TT')\n", "ax.plot(ells, Dl_EE, color='C1', label='EE')\n", "ax.plot(ells, np.abs(Dl_TE), color='C2', label='|TE|')\n", "ax.set_xlabel(r'Multipole $\\ell$')\n", "ax.set_ylabel(r'$\\mathcal{D}_\\ell\\,[\\mu K^2]$')\n", "ax.set_title('Scalar CMB power spectra')\n", "ax.set_xscale('log'); ax.set_yscale('log')\n", "ax.set_xlim(2, ells[-1]); ax.legend(); ax.grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: compare with CAMB\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "import camb\n", "pars = camb.set_params(H0=100 * params['h'],\n", " ombh2=params['omega_b_h2'],\n", " omch2=params['omega_c_h2'],\n", " As=params['A_s'], ns=params['n_s'],\n", " tau=params['tau_reion'])\n", "pars.set_for_lmax(params['ell_max'], lens_potential_accuracy=2)\n", "camb_powers = camb.get_results(pars).get_cmb_power_spectra(\n", " pars, CMB_unit='muK', raw_cl=False)\n", "unlensed = camb_powers['unlensed_scalar'] # shape (lmax+1, 4); cols = TT,EE,BB,TE\n", "\n", "ells_camb = np.arange(unlensed.shape[0])\n", "Dl_TT_camb = unlensed[:, 0]\n", "Dl_EE_camb = unlensed[:, 1]\n", "Dl_TE_camb = unlensed[:, 3]\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot tutorial vs CAMB with residuals\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, axes = plt.subplots(2, 1, sharex=True, figsize=(8, 6),\n", " gridspec_kw={'height_ratios': [3, 1]})\n", "\n", "# Top: spectra\n", "axes[0].plot(ells, Dl_TT, 'k-', lw=1.4, label='tutorial')\n", "axes[0].plot(ells_camb, Dl_TT_camb, 'C3--', lw=1.0, label='CAMB')\n", "axes[0].plot(ells, Dl_EE, 'k-', lw=1.4)\n", "axes[0].plot(ells_camb, Dl_EE_camb, 'C3--', lw=1.0)\n", "axes[0].set_yscale('log'); axes[0].set_xscale('log')\n", "axes[0].set_ylabel(r'$\\mathcal{D}_\\ell$ [$\\mu K^2$]')\n", "axes[0].set_xlim(2, params['ell_max'])\n", "axes[0].legend(); axes[0].grid(True, which='both', alpha=0.3)\n", "axes[0].set_title('Validation: tutorial vs CAMB')\n", "\n", "# Residuals\n", "common = np.intersect1d(ells, ells_camb)\n", "ix_t = np.isin(ells, common)\n", "ix_c = np.isin(ells_camb, common)\n", "resid_TT = 100.0 * (Dl_TT[ix_t] - Dl_TT_camb[ix_c]) / Dl_TT_camb[ix_c]\n", "resid_EE = 100.0 * (Dl_EE[ix_t] - Dl_EE_camb[ix_c]) / Dl_EE_camb[ix_c]\n", "axes[1].plot(common, resid_TT, 'C0-', lw=1.0, label='TT')\n", "axes[1].plot(common, resid_EE, 'C1-', lw=1.0, label='EE')\n", "axes[1].axhline(0, color='gray', lw=0.5)\n", "axes[1].set_xscale('log')\n", "axes[1].set_xlabel(r'Multipole $\\ell$')\n", "axes[1].set_ylabel('residual [%]')\n", "axes[1].set_ylim(-5, 5)\n", "axes[1].legend(); axes[1].grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 6 \u2014 Lensing\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `evolve_phi`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def evolve_phi(k, bg, thermo, pgrid, tau_out):\n", " \"\"\"Evolve mode k, return the Weyl potential Psi_W(tau) at each tau_out point.\"\"\"\n", " tau_start = max(1e-4, min(1e-3 / k, tau_out[0] * 0.5))\n", " y0 = adiabatic_ics(k, tau_start, bg, pgrid)\n", "\n", " bg_vec = pgrid['bg_vec']\n", " sp_a_x = pgrid['sp_a_x']; sp_a_c = pgrid['sp_a_c']\n", " sp_op_x = pgrid['sp_op_x']; sp_op_c = pgrid['sp_op_c']\n", " sp_cs_x = pgrid['sp_cs_x']; sp_cs_c = pgrid['sp_cs_c']\n", "\n", " sol = integrate.solve_ivp(\n", " boltzmann_derivs, (tau_start, tau_out[-1]), y0,\n", " args=(k, bg_vec, sp_a_x, sp_a_c, sp_op_x, sp_op_c, sp_cs_x, sp_cs_c),\n", " t_eval=tau_out, method='LSODA', rtol=1e-6, atol=1e-10)\n", " if not sol.success:\n", " return np.zeros(len(tau_out))\n", "\n", " psi_arr = np.zeros(len(tau_out))\n", " for i, tau in enumerate(tau_out):\n", " y = sol.y[:, i]\n", " (a, adotoa, grhog_t, grhor_t, _, _,\n", " _, _, _, sigma,\n", " etak, _, _, _, _, _, pig, _, _, pir) = _common_terms(\n", " tau, y, k, bg_vec, sp_a_x, sp_a_c)\n", " dgpi = grhog_t * pig + grhor_t * pir\n", " eta = etak / k\n", " # Psi_W = eta - (a'/a) sigma/k - 3 dgpi / (4 k^2)\n", " psi_arr[i] = eta - adotoa * sigma / k - 0.75 * dgpi / (k * k)\n", " return psi_arr\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `lensing_potential`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def lensing_potential(params, bg, thermo, pgrid,\n", " ells=None, N_k=40, N_chi=120, ell_los_max=5):\n", " \"\"\"Compute C_L^{phi phi} via hybrid LOS (low L) + Limber (high L).\"\"\"\n", " tau_star = float(thermo['tau_star'])\n", " tau0 = float(bg['tau0'])\n", " chi_star = tau0 - tau_star\n", "\n", " if ells is None:\n", " ells = np.unique(np.concatenate([\n", " np.arange(2, 30),\n", " np.geomspace(30, 2500, 80).astype(int)]))\n", " ells = np.asarray(ells, dtype=int)\n", "\n", " ell_max = int(ells.max())\n", " chi_min_eff = 0.01 * chi_star\n", " k_max_needed = (ell_max + 0.5) / chi_min_eff\n", " k_max = min(max(k_max_needed * 1.2, 1.0), 5.0)\n", " k_arr = np.geomspace(1e-5, k_max, N_k)\n", "\n", " chi_arr = np.linspace(0.001 * chi_star, 0.999 * chi_star, N_chi)\n", " tau_sorted = np.sort(tau0 - chi_arr)\n", " chi_sorted = tau0 - tau_sorted\n", "\n", " phi_grid_sorted = np.zeros((N_k, N_chi))\n", " for i_k, k in enumerate(k_arr):\n", " phi_grid_sorted[i_k, :] = evolve_phi(k, bg, thermo, pgrid, tau_sorted)\n", " phi_grid_asc = phi_grid_sorted[:, ::-1]\n", "\n", " A_s = params['A_s']; n_s = params['n_s']; k_pivot = params['k_pivot']\n", "\n", " # --- LOS branch (ell <= ell_los_max) ---\n", " ells_los = ells[ells <= ell_los_max]\n", " Cl_los = np.zeros(len(ells_los))\n", " W_lens = (chi_star - chi_sorted) / (chi_star * chi_sorted)\n", " for i_ell, ell in enumerate(ells_los):\n", " Delta = np.zeros(N_k)\n", " for i_k, k in enumerate(k_arr):\n", " jl = special.spherical_jn(int(ell), k * chi_sorted)\n", " integrand = -2.0 * phi_grid_sorted[i_k] * W_lens * jl\n", " Delta[i_k] = -np.trapz(integrand, chi_sorted)\n", " P_R = A_s * (k_arr / k_pivot)**(n_s - 1.0)\n", " Cl_los[i_ell] = 4.0 * np.pi * np.trapz(P_R * Delta**2, np.log(k_arr))\n", "\n", " # --- Limber branch (ell > ell_los_max) ---\n", " chi_asc = chi_sorted[::-1]\n", " phi_interp = interpolate.RegularGridInterpolator(\n", " (np.log(k_arr), chi_asc), phi_grid_asc,\n", " bounds_error=False, fill_value=0.0)\n", " ells_limber = ells[ells > ell_los_max]\n", " Cl_limber = np.zeros(len(ells_limber))\n", " for i_ell, ell in enumerate(ells_limber):\n", " nu = ell + 0.5\n", " k_lim = nu / chi_asc\n", " valid = (k_lim >= k_arr[0]) & (k_lim <= k_arr[-1])\n", " if not np.any(valid):\n", " continue\n", " chi_v = chi_asc[valid]\n", " k_v = k_lim[valid]\n", " T_psi = phi_interp((np.log(k_v), chi_v))\n", " Delta2_R = A_s * (k_v / k_pivot)**(n_s - 1.0)\n", " chiW2 = (chi_star - chi_v)**2 / (chi_star**2 * chi_v)\n", " Cl_limber[i_ell] = (8.0 * np.pi**2 / nu**3) * np.trapz(\n", " T_psi**2 * Delta2_R * chiW2, chi_v)\n", "\n", " Cl_phi = np.zeros(len(ells))\n", " Cl_phi[ells <= ell_los_max] = Cl_los\n", " Cl_phi[ells > ell_los_max] = Cl_limber\n", " Dl_phi = (ells * (ells + 1.0))**2 * Cl_phi / (2.0 * np.pi)\n", "\n", " return {'ell': ells, 'Cl_phi_phi': Cl_phi, 'Dl_phi_phi': Dl_phi,\n", " 'k_arr': k_arr, 'chi_star': chi_star}\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: plot $C_L^{phiphi}$\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "phi = lensing_potential(params, bg, th, init_perturbation_grid(bg, th))\n", "\n", "fig, ax = plt.subplots(figsize=(7, 4.5))\n", "ax.loglog(phi['ell'], phi['Dl_phi_phi'], 'k-', lw=1.5)\n", "ax.set_xlabel(r'$L$')\n", "ax.set_ylabel(r'$[L(L+1)]^2 C_L^{\\phi\\phi}/(2\\pi)$')\n", "ax.set_title('CMB lensing potential power spectrum')\n", "ax.set_xlim(2, 2500); ax.grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: factorials and the normalising prefactor\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "from math import lgamma\n", "from scipy.special import eval_legendre, eval_jacobi\n", "\n", "def _log_fact(n):\n", " return lgamma(n + 1.0)\n", "\n", "def _norm_factor(l, m, mp):\n", " \"\"\"sqrt[(l-m)!(l+m)! / ((l-mp)!(l+mp)!)] computed via log-gamma.\"\"\"\n", " log = 0.5 * (_log_fact(l - m) + _log_fact(l + m)\n", " - _log_fact(l - mp) - _log_fact(l + mp))\n", " return np.exp(log)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: single-$ell$ Wigner evaluator (canonical s and s prime)\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def _d_general(l, m, mp, beta):\n", " \"\"\"d^l_{m, mp}(beta) for arrays of beta. Assumes m >= |mp| and m + mp >= 0.\"\"\"\n", " sin_h = np.sin(beta / 2.0)\n", " cos_h = np.cos(beta / 2.0)\n", " cos_b = np.cos(beta)\n", " norm = _norm_factor(l, m, mp)\n", " s = sin_h ** (m - mp) if m != mp else np.ones_like(beta)\n", " c = cos_h ** (m + mp) if (m + mp) != 0 else np.ones_like(beta)\n", " n = l - m\n", " if n < 0:\n", " return np.zeros_like(beta)\n", " jacobi = eval_jacobi(n, m - mp, m + mp, cos_b)\n", " return norm * s * c * jacobi\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `wigner_d_array`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def wigner_d_array(l_max, m, mp, beta):\n", " \"\"\"Compute d^l_{m, mp}(beta) for l = 0..l_max at all beta values.\n", " Returns array of shape (l_max + 1, len(beta)).\"\"\"\n", " beta = np.atleast_1d(beta).astype(float)\n", " d = np.zeros((l_max + 1, len(beta)))\n", "\n", " # (0, 0) reduces to Legendre P_l(cos beta)\n", " if m == 0 and mp == 0:\n", " cos_b = np.cos(beta)\n", " for l in range(l_max + 1):\n", " d[l] = eval_legendre(l, cos_b)\n", " return d\n", "\n", " # Map (m, mp) to canonical form using symmetries.\n", " sign = 1.0\n", " m_eff, mp_eff = m, mp\n", " if abs(mp_eff) > abs(m_eff):\n", " sign *= (-1.0) ** (m_eff - mp_eff)\n", " m_eff, mp_eff = mp_eff, m_eff\n", " if m_eff < 0 or (m_eff + mp_eff) < 0:\n", " sign *= (-1.0) ** (m_eff - mp_eff)\n", " m_eff, mp_eff = -m_eff, -mp_eff\n", " if abs(mp_eff) > abs(m_eff):\n", " sign *= (-1.0) ** (m_eff - mp_eff)\n", " m_eff, mp_eff = mp_eff, m_eff\n", "\n", " l_min = max(abs(m_eff), abs(mp_eff))\n", " for l in range(l_min, l_max + 1):\n", " d[l] = sign * _d_general(l, m_eff, mp_eff, beta)\n", " return d\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `lensed_cls_TT_EE_TE` (CL05)\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "from scipy.special import roots_legendre\n", "\n", "def lensed_cls_TT_EE_TE(ell_arr, Cl_TT, Cl_EE, Cl_TE, Cl_phi_phi, n_theta=None):\n", " \"\"\"Curved-sky lensed TT, EE, TE via the CL05 small-deflection expansion.\n", " Recommended: n_theta >= 2 * max(ell_arr) for sub-percent TT/EE/TE.\"\"\"\n", " ell_arr = np.asarray(ell_arr, dtype=int)\n", " lmax = int(ell_arr.max())\n", " if n_theta is None:\n", " n_theta = max(2 * lmax + 200, 500)\n", " ell_full = np.arange(lmax + 1)\n", " def _pad(C):\n", " return np.interp(ell_full, ell_arr, np.asarray(C), left=0.0, right=0.0)\n", " CTT, CEE, CTE, Cpp = _pad(Cl_TT), _pad(Cl_EE), _pad(Cl_TE), _pad(Cl_phi_phi)\n", "\n", " # Gauss-Legendre nodes and weights for the beta integral\n", " x, w = roots_legendre(n_theta)\n", " beta = np.arccos(x)\n", "\n", " # Wigner-d tables on the beta grid\n", " d_00 = wigner_d_array(lmax, 0, 0, beta)\n", " d_22 = wigner_d_array(lmax, 2, 2, beta)\n", " d_2m2 = wigner_d_array(lmax, 2, -2, beta)\n", " d_02 = wigner_d_array(lmax, 0, 2, beta)\n", "\n", " ell_w = (2 * ell_full + 1) / (4.0 * np.pi)\n", " Lpp = ell_full * (ell_full + 1.0)\n", "\n", " # Total deflection variance and beta-dependent variance\n", " sigma2_alpha = np.sum((2 * ell_full + 1) * Lpp * Cpp / (4.0 * np.pi))\n", " C_gl_beta = np.einsum('l,lb->b', (2*ell_full + 1) * Lpp * Cpp / (4.0*np.pi), d_00)\n", " sigma2_beta = np.maximum(sigma2_alpha - C_gl_beta, 0.0)\n", "\n", " # Per-l Gaussian lensing modulation: exp(- L(L+1) sigma^2(beta) / 2)\n", " smoothing_per_l = np.exp(-0.5 * Lpp[:, None] * sigma2_beta[None, :])\n", " weight = ell_w[:, None] * smoothing_per_l\n", "\n", " # Lensed correlation functions\n", " xi_TT_lens = np.einsum('lb,lb->b', CTT[:, None] * weight, d_00)\n", " xi_p_lens = np.einsum('lb,lb->b', CEE[:, None] * weight, d_22)\n", " xi_m_lens = np.einsum('lb,lb->b', CEE[:, None] * weight, d_2m2)\n", " xi_TE_lens = -np.einsum('lb,lb->b', CTE[:, None] * weight, d_02)\n", "\n", " # Inverse transforms back to multipoles\n", " Cl_TT_lensed = 2.0 * np.pi * np.einsum('b,lb,b->l', xi_TT_lens, d_00, w)\n", " Cl_EE_lensed = (np.pi * np.einsum('b,lb,b->l', xi_p_lens, d_22, w)\n", " + np.pi * np.einsum('b,lb,b->l', xi_m_lens, d_2m2, w))\n", " Cl_TE_lensed = -2.0 * np.pi * np.einsum('b,lb,b->l', xi_TE_lens, d_02, w)\n", "\n", " return {\n", " 'ell': ell_arr,\n", " 'Cl_TT_lensed': np.interp(ell_arr, ell_full, Cl_TT_lensed),\n", " 'Cl_EE_lensed': np.interp(ell_arr, ell_full, Cl_EE_lensed),\n", " 'Cl_TE_lensed': np.interp(ell_arr, ell_full, Cl_TE_lensed),\n", " 'sigma2_alpha': sigma2_alpha,\n", " }\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `lensed_BB_from_EE_flat_sky`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def lensed_BB_from_EE_flat_sky(ell_arr, Cl_EE_unl, Cl_phi_phi, n_l=200, n_phi=80):\n", " \"\"\"Flat-sky BB from lensing of E, Hu-Okamoto / Lewis-Challinor formula.\"\"\"\n", " ell_arr = np.asarray(ell_arr, dtype=int)\n", " lmax = int(ell_arr.max())\n", " ell_full = np.arange(lmax + 1)\n", " CEE_full = np.interp(ell_full, ell_arr, np.asarray(Cl_EE_unl), left=0.0, right=0.0)\n", " Cpp_full = np.interp(ell_full, ell_arr, np.asarray(Cl_phi_phi), left=0.0, right=0.0)\n", "\n", " # Log-spaced l' grid; uniform alpha grid for the angular integral\n", " l_p = np.geomspace(2.0, float(lmax), n_l)\n", " Cpp_p = np.interp(l_p, ell_full, Cpp_full)\n", " dlnlp = np.gradient(np.log(l_p))\n", "\n", " alpha = np.linspace(0.0, 2.0 * np.pi, n_phi, endpoint=False)\n", " da = 2.0 * np.pi / n_phi\n", " cos_a, sin_a = np.cos(alpha), np.sin(alpha)\n", " sin2_a = sin_a ** 2\n", "\n", " Cl_BB_lens = np.zeros(len(ell_full))\n", " for L in ell_full:\n", " if L < 2:\n", " continue\n", " Lm_sq = L*L + l_p[None, :]**2 - 2.0 * L * l_p[None, :] * cos_a[:, None]\n", " Lm_sq = np.maximum(Lm_sq, 0.25)\n", " Lm = np.sqrt(Lm_sq)\n", " dot = L * l_p[None, :] * cos_a[:, None] - l_p[None, :]**2\n", " Lmm = L - l_p[None, :] * cos_a[:, None]\n", " sin2_2phi = 4.0 * l_p[None, :]**2 * sin2_a[:, None] * Lmm**2 / Lm_sq**2\n", " CEE_at = np.interp(Lm.ravel(), ell_full, CEE_full).reshape(Lm.shape)\n", " integrand = dot**2 * sin2_2phi * Cpp_p[None, :] * CEE_at\n", " Cl_BB_lens[L] = (1.0 / (2.0 * np.pi)**2) * np.sum(\n", " integrand * l_p[None, :]**2 * dlnlp[None, :] * da)\n", "\n", " return np.interp(ell_arr, ell_full, Cl_BB_lens)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `lens_spectra`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def lens_spectra(ell_arr, Cl_TT, Cl_EE, Cl_TE, Cl_phi_phi,\n", " Cl_BB_unl=None, n_theta=None, n_l_BB=200, n_phi_BB=80):\n", " \"\"\"Full lensed CMB spectra: TT, EE, BB, TE.\n", " Cl_BB_unl: optional unlensed BB (e.g. tensors from chapter 7).\"\"\"\n", " res = lensed_cls_TT_EE_TE(ell_arr, Cl_TT, Cl_EE, Cl_TE, Cl_phi_phi,\n", " n_theta=n_theta)\n", " Cl_BB_from_EE = lensed_BB_from_EE_flat_sky(\n", " ell_arr, Cl_EE, Cl_phi_phi, n_l=n_l_BB, n_phi=n_phi_BB)\n", " if Cl_BB_unl is not None:\n", " # Smooth the input BB (e.g. tensor) by the same Gaussian lensing envelope\n", " sigma2 = res['sigma2_alpha']\n", " Lpp = ell_arr * (ell_arr + 1.0)\n", " Cl_BB_in_smoothed = np.asarray(Cl_BB_unl) * np.exp(-0.5 * Lpp * sigma2)\n", " else:\n", " Cl_BB_in_smoothed = np.zeros_like(ell_arr, dtype=float)\n", " res['Cl_BB_lensed'] = Cl_BB_from_EE + Cl_BB_in_smoothed\n", " res['Cl_BB_from_EE_lensing'] = Cl_BB_from_EE\n", " return res\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: produce lensed scalar spectra\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "ells_unl = result['ells']\n", "Cl_TT_unl = result['Dl_TT'] * 2*np.pi / (ells_unl*(ells_unl+1))\n", "Cl_EE_unl = result['Dl_EE'] * 2*np.pi / (ells_unl*(ells_unl+1))\n", "Cl_TE_unl = result['Dl_TE'] * 2*np.pi / (ells_unl*(ells_unl+1))\n", "Cl_phi_interp = np.interp(ells_unl, phi['ell'], phi['Cl_phi_phi'])\n", "lensed = lens_spectra(ells_unl, Cl_TT_unl, Cl_EE_unl, Cl_TE_unl, Cl_phi_interp)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: lensed vs unlensed\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, axes = plt.subplots(2, 3, figsize=(15, 8),\n", " gridspec_kw={'height_ratios': [3, 1]})\n", "for col, label, unl, lens in [\n", " (0, 'TT', Cl_TT_unl, lensed['Cl_TT_lensed']),\n", " (1, 'EE', Cl_EE_unl, lensed['Cl_EE_lensed']),\n", " (2, 'BB', np.zeros_like(Cl_TT_unl), lensed['Cl_BB_lensed'])]:\n", " top, bot = axes[0, col], axes[1, col]\n", " Dl_unl = ells_unl*(ells_unl+1)/(2*np.pi) * unl\n", " Dl_lens = ells_unl*(ells_unl+1)/(2*np.pi) * lens\n", " if label == 'BB':\n", " top.loglog(ells_unl, Dl_lens, 'C3-', lw=1.4, label='lensed')\n", " top.set_ylim(1e-5, 1)\n", " bot.axis('off')\n", " else:\n", " top.loglog(ells_unl, Dl_unl, 'C0--', lw=1.2, label='unlensed')\n", " top.loglog(ells_unl, Dl_lens, 'C3-', lw=1.4, label='lensed')\n", " resid = 100*(Dl_lens - Dl_unl)/np.maximum(Dl_unl, 1e-30)\n", " bot.semilogx(ells_unl, resid, 'k-', lw=1.0)\n", " bot.axhline(0, color='gray', lw=0.5)\n", " bot.set_xlabel(r'$\\ell$'); bot.set_ylabel('residual [%]')\n", " bot.set_xlim(2, 2500); bot.grid(True, which='both', alpha=0.3)\n", " top.set_title(label); top.set_xlim(2, 2500); top.legend()\n", " top.set_ylabel(rf'$\\mathcal{{D}}_\\ell^{{{label}}}\\,[\\mu K^2]$')\n", " top.grid(True, which='both', alpha=0.3)\n", "plt.tight_layout(); plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 7 \u2014 Tensors\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: tensor truncation and indices\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "LMAX_T_T = 10 # tensor temperature truncation\n", "LMAX_E_T = 10\n", "LMAX_B_T = 10\n", "\n", "IX_TENS_H = 0 # h_T\n", "IX_TENS_HDOT = 1 # h_T dot\n", "IX_TENS_T = 2 # Theta^T_2, _3, ...\n", "IX_TENS_E = IX_TENS_T + (LMAX_T_T - 1) # E^T_2, _3, ...\n", "IX_TENS_B = IX_TENS_E + (LMAX_E_T - 1) # B^T_2, _3, ...\n", "NVAR_TENS = IX_TENS_B + (LMAX_B_T - 1)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `_tensor_kappa`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def _tensor_kappa(ell):\n", " \"\"\"Spin-2 angular-momentum coupling coefficient (Hu & White 1997).\"\"\"\n", " if ell < 2:\n", " return 0.0\n", " return np.sqrt((ell**2 - 4) * (ell**2 - 1)) / ell\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: tensor Boltzmann RHS\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def _tensor_rhs(tau, y, k, bg_vec, sp_a_x, sp_a_c, sp_op_x, sp_op_c):\n", " \"\"\"Tensor Boltzmann hierarchy in Hu-White conventions.\"\"\"\n", " grhog, grhornomass, grhoc, grhob, grhov = bg_vec\n", " a = _cubic_eval(sp_a_x, sp_a_c, tau)\n", " a2 = a * a\n", " grho_a2 = grhog/a2 + grhornomass/a2 + grhoc/a + grhob/a + grhov*a2\n", " adotoa = np.sqrt(grho_a2 / 3.0)\n", " opacity = max(_cubic_eval(sp_op_x, sp_op_c, tau), 1e-30)\n", "\n", " h_T = y[IX_TENS_H]\n", " h_Tdot = y[IX_TENS_HDOT]\n", "\n", " Theta = np.zeros(LMAX_T_T + 2)\n", " E = np.zeros(LMAX_E_T + 2)\n", " B = np.zeros(LMAX_B_T + 2)\n", " for ell in range(2, LMAX_T_T + 1):\n", " Theta[ell] = y[IX_TENS_T + ell - 2]\n", " for ell in range(2, LMAX_E_T + 1):\n", " E[ell] = y[IX_TENS_E + ell - 2]\n", " for ell in range(2, LMAX_B_T + 1):\n", " B[ell] = y[IX_TENS_B + ell - 2]\n", "\n", " Pi_T = 0.1 * (Theta[2] - np.sqrt(6.0) * E[2])\n", "\n", " dy = np.zeros(NVAR_TENS)\n", " # GW wave equation\n", " dy[IX_TENS_H] = h_Tdot\n", " dy[IX_TENS_HDOT] = -2.0 * adotoa * h_Tdot - k * k * h_T\n", "\n", " # Temperature hierarchy with metric source delta_{l,2} h_Tdot\n", " for ell in range(2, LMAX_T_T + 1):\n", " kappa_l = _tensor_kappa(ell)\n", " kappa_lp1 = _tensor_kappa(ell + 1)\n", " T_lm1 = Theta[ell - 1] if ell > 2 else 0.0\n", " T_lp1 = Theta[ell + 1] if ell + 1 <= LMAX_T_T else 0.0\n", " if ell == 2:\n", " src = -h_Tdot - opacity * (Theta[ell] - Pi_T)\n", " elif ell == LMAX_T_T:\n", " src = -(ell + 1) / tau * Theta[ell] - opacity * Theta[ell]\n", " dy[IX_TENS_T + ell - 2] = k * T_lm1 + src\n", " continue\n", " else:\n", " src = -opacity * Theta[ell]\n", " dy[IX_TENS_T + ell - 2] = (k / (2.0 * ell + 1.0)) * (\n", " kappa_l * T_lm1 - kappa_lp1 * T_lp1) + src\n", "\n", " # E-B coupled hierarchies\n", " for ell in range(2, LMAX_E_T + 1):\n", " kappa_l = _tensor_kappa(ell); kappa_lp1 = _tensor_kappa(ell + 1)\n", " E_lm1 = E[ell - 1] if ell > 2 else 0.0\n", " E_lp1 = E[ell + 1] if ell + 1 <= LMAX_E_T else 0.0\n", " B_l = B[ell] if ell <= LMAX_B_T else 0.0\n", " if ell == 2:\n", " src_E = opacity * (-(E[ell] - np.sqrt(6.0) * Pi_T / 2.0))\n", " elif ell == LMAX_E_T:\n", " src_E = -(ell + 1) / tau * E[ell] - opacity * E[ell]\n", " dy[IX_TENS_E + ell - 2] = k * E_lm1 + src_E\n", " continue\n", " else:\n", " src_E = -opacity * E[ell]\n", " dy[IX_TENS_E + ell - 2] = (k / (2.0 * ell + 1.0)) * (\n", " kappa_l * E_lm1 - kappa_lp1 * E_lp1\n", " ) - 2.0 * k * B_l / (ell * (ell + 1.0)) + src_E\n", "\n", " for ell in range(2, LMAX_B_T + 1):\n", " kappa_l = _tensor_kappa(ell); kappa_lp1 = _tensor_kappa(ell + 1)\n", " B_lm1 = B[ell - 1] if ell > 2 else 0.0\n", " B_lp1 = B[ell + 1] if ell + 1 <= LMAX_B_T else 0.0\n", " E_l = E[ell] if ell <= LMAX_E_T else 0.0\n", " if ell == LMAX_B_T:\n", " src_B = -(ell + 1) / tau * B[ell] - opacity * B[ell]\n", " dy[IX_TENS_B + ell - 2] = k * B_lm1 + src_B\n", " continue\n", " dy[IX_TENS_B + ell - 2] = (k / (2.0 * ell + 1.0)) * (\n", " kappa_l * B_lm1 - kappa_lp1 * B_lp1\n", " ) + 2.0 * k * E_l / (ell * (ell + 1.0)) - opacity * B[ell]\n", "\n", " return dy\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: `tensor_spectra`\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "def tensor_spectra(params, bg, thermo, pgrid,\n", " ells=None, r=0.05, n_t=None, N_k=200, N_tau=300):\n", " \"\"\"Tensor-mode CMB angular power spectra.\n", " r : tensor-to-scalar ratio at k_pivot\n", " n_t : tensor spectral index; None -> consistency relation n_t = -r/8\n", " \"\"\"\n", " if n_t is None:\n", " n_t = -r / 8.0\n", " A_s = params['A_s']; A_t = r * A_s; k_pivot = params['k_pivot']\n", "\n", " k_arr = np.geomspace(1e-5, 0.5, N_k)\n", " tau_min = max(1e-3, float(thermo['tau_arr'][0]))\n", " tau_max = float(bg['tau0']) * 0.9999\n", " tau_grid = np.geomspace(tau_min, tau_max, N_tau)\n", " vis = thermo['visibility_interp'](tau_grid)\n", " exptau = thermo['exptau_interp'](tau_grid)\n", "\n", " if ells is None:\n", " ells = np.unique(np.concatenate([\n", " np.arange(2, 30), np.arange(30, 200, 5),\n", " np.arange(200, 1001, 25)]))\n", " ells = np.asarray(ells, dtype=int)\n", "\n", " histories = []\n", " for k in k_arr:\n", " tau_start = max(1e-4, min(1e-3 / k, tau_grid[0] * 0.5))\n", " y0 = np.zeros(NVAR_TENS); y0[IX_TENS_H] = 1.0\n", " sol = integrate.solve_ivp(\n", " _tensor_rhs, (tau_start, tau_grid[-1]), y0,\n", " args=(k, pgrid['bg_vec'], pgrid['sp_a_x'], pgrid['sp_a_c'],\n", " pgrid['sp_op_x'], pgrid['sp_op_c']),\n", " t_eval=tau_grid, method='LSODA', rtol=1e-6, atol=1e-10)\n", " histories.append(sol.y if sol.success else None)\n", "\n", " chi = float(bg['tau0']) - tau_grid\n", "\n", " Cl_TT = np.zeros(len(ells)); Cl_EE = np.zeros(len(ells))\n", " Cl_BB = np.zeros(len(ells)); Cl_TE = np.zeros(len(ells))\n", " P_t = A_t * (k_arr / k_pivot)**n_t\n", "\n", " for i_k, k in enumerate(k_arr):\n", " sol = histories[i_k]\n", " if sol is None:\n", " continue\n", " h_Tdot = sol[IX_TENS_HDOT]\n", " Theta_2 = sol[IX_TENS_T]; E_2 = sol[IX_TENS_E]\n", " Pi_T = 0.1 * (Theta_2 - np.sqrt(6.0) * E_2)\n", " for i_ell, ell in enumerate(ells):\n", " ell_f = float(ell)\n", " ell_factor = np.sqrt((ell_f - 1) * ell_f * (ell_f + 1) * (ell_f + 2))\n", " x = np.maximum(k * chi, 1e-30)\n", " jl = special.spherical_jn(int(ell), x)\n", " F_T = ell_factor * jl / x**2\n", "\n", " # Temperature source: ISW-like -h_Tdot e^{-kappa} + visibility * Pi_T\n", " S_T = -h_Tdot * exptau + vis * Pi_T\n", " S_E = vis * Pi_T * np.sqrt(6.0) / 4.0\n", " S_B = vis * Pi_T * np.sqrt(6.0) / 4.0\n", " Delta_T = np.trapz(S_T * F_T, tau_grid)\n", " Delta_E = np.trapz(S_E * F_T, tau_grid)\n", " Delta_B = np.trapz(S_B * F_T, tau_grid)\n", "\n", " weight = 4.0 * np.pi * P_t[i_k]\n", " dlnk = np.log(k_arr[1] / k_arr[0]) if i_k > 0 else 0.0\n", " Cl_TT[i_ell] += weight * Delta_T**2 * dlnk\n", " Cl_EE[i_ell] += weight * Delta_E**2 * dlnk\n", " Cl_BB[i_ell] += weight * Delta_B**2 * dlnk\n", " Cl_TE[i_ell] += weight * Delta_T * Delta_E * dlnk\n", "\n", " T0_muK2 = (params['T_cmb'] * 1e6)**2\n", " Cl_TT *= T0_muK2; Cl_EE *= T0_muK2; Cl_BB *= T0_muK2; Cl_TE *= T0_muK2\n", " Dl = lambda C: ells * (ells + 1) / (2 * np.pi) * C\n", " return {'ell': ells,\n", " 'Cl_TT_tensor': Cl_TT, 'Cl_EE_tensor': Cl_EE,\n", " 'Cl_BB_tensor': Cl_BB, 'Cl_TE_tensor': Cl_TE,\n", " 'Dl_TT_tensor': Dl(Cl_TT), 'Dl_EE_tensor': Dl(Cl_EE),\n", " 'Dl_BB_tensor': Dl(Cl_BB), 'Dl_TE_tensor': Dl(Cl_TE)}\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: tensor BB at several values of $r$\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "pgrid = init_perturbation_grid(bg, th)\n", "results_t = {r: tensor_spectra(params, bg, th, pgrid, r=r)\n", " for r in (0.05, 0.01, 0.001)}\n", "\n", "fig, ax = plt.subplots(figsize=(7.5, 5))\n", "for r, color in zip([0.05, 0.01, 0.001], ['C3', 'C4', 'C0']):\n", " Dl_BB_t = results_t[r]['Dl_BB_tensor']\n", " ell_t = results_t[r]['ell']\n", " ax.loglog(ell_t, Dl_BB_t, color=color, lw=1.5, label=f'tensor BB, $r={r}$')\n", "# Overlay lensing BB\n", "ax.loglog(ells_unl, lensed['Cl_BB_lensed'] * ells_unl * (ells_unl + 1) / (2 * np.pi),\n", " 'k--', lw=1.4, label='lensing BB ($r=0$)')\n", "ax.set_xlabel(r'$\\ell$'); ax.set_ylabel(r'$\\mathcal{D}_\\ell^{BB}\\,[\\mu K^2]$')\n", "ax.set_title('Tensor BB vs lensing BB')\n", "ax.set_xlim(2, 2500); ax.set_ylim(1e-7, 1e-1)\n", "ax.legend(); ax.grid(True, which='both', alpha=0.3)\n", "plt.show()\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: total spectra and CAMB validation\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "import camb\n", "pars = camb.set_params(H0=100*params['h'],\n", " ombh2=params['omega_b_h2'],\n", " omch2=params['omega_c_h2'],\n", " As=params['A_s'], ns=params['n_s'],\n", " tau=params['tau_reion'], r=0.05)\n", "pars.set_for_lmax(2500, lens_potential_accuracy=2)\n", "pars.WantTensors = True\n", "camb_powers = camb.get_results(pars).get_cmb_power_spectra(\n", " pars, CMB_unit='muK', raw_cl=False)\n", "Dl_camb_lens = camb_powers['lensed_scalar']\n", "Dl_camb_tens = camb_powers['tensor']\n", "\n", "# Tutorial total: lensed scalar + tensor\n", "r_target = 0.05\n", "tens_r = results_t[r_target]\n", "Dl_BB_t_interp = np.interp(ells_unl, tens_r['ell'], tens_r['Dl_BB_tensor'])\n", "Dl_TT_t_interp = np.interp(ells_unl, tens_r['ell'], tens_r['Dl_TT_tensor'])\n", "Dl_EE_t_interp = np.interp(ells_unl, tens_r['ell'], tens_r['Dl_EE_tensor'])\n", "\n", "Dl_TT_tot = (ells_unl*(ells_unl+1)/(2*np.pi)) * lensed['Cl_TT_lensed'] + Dl_TT_t_interp\n", "Dl_EE_tot = (ells_unl*(ells_unl+1)/(2*np.pi)) * lensed['Cl_EE_lensed'] + Dl_EE_t_interp\n", "Dl_BB_tot = (ells_unl*(ells_unl+1)/(2*np.pi)) * lensed['Cl_BB_lensed'] + Dl_BB_t_interp\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Cell: final 4-panel validation plot\n" ] }, { "cell_type": "code", "metadata": {}, "source": [ "fig, axes = plt.subplots(2, 2, figsize=(12, 8))\n", "ells_camb = np.arange(Dl_camb_lens.shape[0])\n", "\n", "for ax, label, idx, col in [\n", " (axes[0,0], 'TT', 0, Dl_TT_tot),\n", " (axes[0,1], 'EE', 1, Dl_EE_tot),\n", " (axes[1,0], 'BB lensed (scalar)', 2, (ells_unl*(ells_unl+1)/(2*np.pi))*lensed['Cl_BB_lensed']),\n", " (axes[1,1], 'BB tensor (r=0.05)', 2, Dl_BB_t_interp),\n", " ]:\n", " if 'tensor' in label:\n", " camb_y = Dl_camb_tens[:, 2]\n", " elif 'lensed' in label:\n", " # lensing BB from CAMB = lensed_scalar BB\n", " camb_y = Dl_camb_lens[:, 2]\n", " else:\n", " camb_y = Dl_camb_lens[:, idx]\n", " ax.plot(ells_unl, col, 'k-', lw=1.4, label='tutorial')\n", " ax.plot(ells_camb, camb_y, 'C3--', lw=1.0, label='CAMB')\n", " ax.set_xscale('log'); ax.set_yscale('log')\n", " ax.set_xlim(2, 2500)\n", " ax.set_xlabel(r'$\\ell$')\n", " ax.set_ylabel(rf'$\\mathcal{{D}}_\\ell\\,[\\mu K^2]$')\n", " ax.set_title(label); ax.legend(); ax.grid(True, which='both', alpha=0.3)\n", "plt.tight_layout(); plt.show()\n" ], "outputs": [], "execution_count": null } ], "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3", "language": "python" }, "language_info": { "name": "python", "version": "3.10" } }, "nbformat": 4, "nbformat_minor": 5 }