CMB_School_Part_12.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jeffmcm1977/CMBAnalysis_SummerSchool/blob/master/CMB_School_Part_12.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lensing reconstruction tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import cmb_modules\" || ( \\\n",
    "    wget https://github.com/jeffmcm1977/CMBAnalysis_SummerSchool/raw/master/cmb_school.tar.gz && \\\n",
    "    tar xzvf cmb_school.tar.gz \\\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# We use some stuff we learned before\n",
    "import cmb_modules\n",
    "import lens_modules as lm\n",
    "np.random.seed(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will learn how to obtain the underlying lensing convergence -- the projected (dark) matter density -- from an observed map of the CMB. \n",
    "\n",
    "We will first simulate a lensed CMB map as we learnt in an earlier tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "patch_deg_width = 40. # patch width in degrees\n",
    "pix_size = 1.5 # pixel size in arcminutes\n",
    "ells,ucltt,lcltt,clkk = lm.get_theory()\n",
    "N,lensed,kappa,ly,lx,modlmap = lm.get_lensed(patch_deg_width,pix_size,ells,ucltt,clkk)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now convolve this map with a beam and add noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Beam\n",
    "beam_arcmin = 1.4\n",
    "def gauss_beam(ell,fwhm):\n",
    "    # A gaussian beam transfer function (map-space, i.e. one power)\n",
    "    tht_fwhm = np.deg2rad(fwhm / 60.)\n",
    "    return np.exp(-(tht_fwhm**2.)*(ell**2.) / (16.*np.log(2.)))\n",
    "# Evaluate the beam on an isotropic Fourier grid\n",
    "kbeam2d = gauss_beam(modlmap,beam_arcmin)\n",
    "# Filter the map to convolve it with a beam\n",
    "beamed = lm.filter_map(lensed,kbeam2d)\n",
    "\n",
    "# Noise\n",
    "noise_uk_arcmin = 1.5 # CMB-S4 level white noise\n",
    "# White noise\n",
    "Clnoise = (noise_uk_arcmin*np.pi/180./60.)**2.\n",
    "Dlnoise = Clnoise*ells*(ells+1.)/2./np.pi\n",
    "# Make a GRF noise map\n",
    "noise_map = cmb_modules.make_CMB_T_map(N,pix_size,ells,Dlnoise)\n",
    "\n",
    "# The observed map\n",
    "observed = beamed + noise_map\n",
    "\n",
    "plt.imshow(observed)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lensing reconstruction\n",
    "\n",
    "Let's expand the lens equation\n",
    "\n",
    "$\\newcommand{\\al}{\\boldsymbol{\\alpha}}$\n",
    "$\\newcommand{\\x}{\\boldsymbol{x}}$\n",
    "$\\newcommand{\\nab}{\\boldsymbol{\\nabla}}$\n",
    "$\\newcommand{\\uT}{{\\tilde{T}}}$\n",
    "$\\newcommand{\\uC}{{\\tilde{C}}}$\n",
    "$\\newcommand{\\dphi}{{\\nab\\phi}}$\n",
    "$\\newcommand{\\O}{{\\mathcal{O}}}$\n",
    "$\\newcommand{\\l}{{\\boldsymbol{\\ell}}}$\n",
    "$\\newcommand{\\L}{{\\boldsymbol{L}}}$\n",
    "$$T(\\x) = \\uT(\\x+\\al)$$\n",
    "\n",
    "in a Taylor series (and substitute $\\al=\\dphi$) in the weak-lensing limit of small $\\phi$,\n",
    "\n",
    "$$\n",
    "T(\\x) \\approx \\uT(\\x) + \\nabla \\uT \\cdot \\dphi + \\O(\\dphi^2)\n",
    "$$ \n",
    "\n",
    "HOMEWORK EXERCISE: Using 2d Fourier transforms,\n",
    "\n",
    "$$\n",
    "X(\\l) = \\int d^2\\x ~e^{-i\\l\\cdot\\x}X(\\x)\n",
    "$$\n",
    "$$\n",
    "X(\\x) = \\int \\frac{d^2\\l}{(2\\pi)^2} ~e^{i\\l\\cdot\\x}X(\\l)\n",
    "$$\n",
    "$$\n",
    "\\int d^2\\x e^{i\\l\\cdot\\x} = (2\\pi)^2 \\delta(\\l)\n",
    "$$\n",
    "\n",
    "show that to leading order in the lensing potential, lensing induces coupling between modes of the CMB map:\n",
    "\n",
    "\n",
    "$$\n",
    "\\langle T(\\l_1)T(\\l_2) \\rangle_{\\rm CMB} = \\phi(\\L)\\left[\\uC^{TT}_{\\ell_1}(\\L \\cdot \\l_1) + \\uC^{TT}_{\\ell_2}(\\L \\cdot \\l_2)\\right]\n",
    "$$\n",
    "\n",
    "where the averaging is over realizations of the unlensed CMB field $\\uT$. The above result motivates extracting modes of the lensing potential $\\phi(\\L)$ by taking a weighted average of products of CMB modes. The weights can be derived such that the noise in the reconstruction is minimized. The resulting estimator is:\n",
    "\n",
    "$$\n",
    "\\hat{\\kappa}(\\L) = -A(\\L) \\int d^2\\l_1 \\left[\\uC^{TT}_{\\ell_1}(\\L \\cdot \\l_1) + \\uC^{TT}_{\\ell_2}(\\L \\cdot \\l_2)\\right]\\frac{T(\\l_1)T(\\l_2)}{C^{TT}_{\\ell_1}C^{TT}_{\\ell_2}}\n",
    "$$\n",
    "\n",
    "HOMEWORK EXERCISE: Show that the above estimator can be written as a real-space multiplication\n",
    "\n",
    "$$\n",
    "\\hat{\\kappa} \\propto -\\nabla \\cdot \\left[\\left(\\nabla T_W\\right) T_H \\right]\n",
    "$$\n",
    "\n",
    "where $T_W$ is the temperature map Wiener filtered by $\\frac{\\uC^{TT}_{\\ell}}{C^{TT}_{\\ell}}$ and $T_H$ is the temperature map inverse variance filtered by $\\frac{1}{C^{TT}_{\\ell}}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we attempt to reconstruct a lensing map from the observed map\n",
    "def qe_reconstruct(tmap,unlensed_cmb_power_2d,total_cmb_power_2d,ellmin,ellmax,modlmap,ly,lx):\n",
    "    \"\"\"\n",
    "    The simplest form of the quadratic estimator is\n",
    "\n",
    "    kappa_recon = - div ( grad(T) T )\n",
    "    where grad(T) is Wiener filtered with (unlensed_cmb_power/total_cmb_power)\n",
    "    and T is filtered with (1/total_cmb_power)\n",
    "    where the T map is beam deconvolved and\n",
    "    where total_cmb_power = lensed_cmb_power + noise_power / beam^2\n",
    "\n",
    "    The reconstruction is limited to ellmin < ells < ellmax of the CMB temperature map.\n",
    "\n",
    "    The grad and div operations require knowledge of:\n",
    "    modlmap: absolute wavenumbers of Fourier pixels\n",
    "    ly: map of ly Fourier coordinate of each pixel\n",
    "    lx: map of lx Fourier coordinate of each pixel\n",
    "    (note modlmap = sqrt(ly**2 + lx**2)\n",
    "    \"\"\"\n",
    "    inv_noise_filter = lm.kmask((1./total_cmb_power_2d),modlmap,ellmin,ellmax)\n",
    "    grad_filter = lm.kmask((unlensed_cmb_power_2d/total_cmb_power_2d),modlmap,ellmin,ellmax)\n",
    "\n",
    "    gradTy,gradTx = lm.gradient(tmap,ly,lx)\n",
    "\n",
    "    # The Wiener filtered gradient Grad(T_G)\n",
    "    filtered_gradTy = lm.filter_map(gradTy,grad_filter)\n",
    "    filtered_gradTx = lm.filter_map(gradTx,grad_filter)\n",
    "    # The inverse variance filtered T_H\n",
    "    filtered_T = lm.filter_map(tmap,inv_noise_filter)\n",
    "    \n",
    "    ############## COMPLETE THIS!!!\n",
    "    # ADD A LINE BELOW TO GET ukappa = divergence(Grad(T_G) * T_H)\n",
    "    ############## COMPLETE THIS!!!\n",
    "    \n",
    "    # We filter with 1/L^2 at the end to make the normalization white\n",
    "    return -lm.filter_map(ukappa,lm.kmask(1/modlmap**2,modlmap,ellmin=2))\n",
    "\n",
    "# For this we also need to know how to calculate a divergence\n",
    "def div(imapy,imapx,ly,lx):\n",
    "    # Return divergence grad(Y)_y + grad(X)_x\n",
    "    ############## COMPLETE THIS!!!\n",
    "    # COMPLETE THIS FUNCTION THAT RETURNS THE DIVERGENCE\n",
    "    ############## COMPLETE THIS!!!\n",
    "    return divergence\n",
    "\n",
    "\n",
    "def interp(x,y,bounds_error=False,fill_value=0.,**kwargs):\n",
    "    # Just a simple interpolator that fills with zeros by default\n",
    "    from scipy.interpolate import interp1d\n",
    "    return interp1d(x,y,bounds_error=bounds_error,fill_value=fill_value,**kwargs)\n",
    "\n",
    "\n",
    "# To use this we need unlensed and total spectra interpolated on to the 2D Fourier grid\n",
    "### Prepare 2D Fourier space interpolations of theory spectra\n",
    "unlensed_cmb_power_2d = interp(ells,ucltt)(modlmap)\n",
    "total_cmb_power_2d = interp(ells,lcltt)(modlmap) + Clnoise/kbeam2d**2.\n",
    "\n",
    "\n",
    "### The noise was specified for a beam deconvolved map so we deconvolve the beam from our map\n",
    "tmap = lm.filter_map(observed,1/kbeam2d)\n",
    "\n",
    "### Next, we get the unnormalized reconstruction after choosing what multipoles to involve\n",
    "ellmin = 10\n",
    "ellmax = 3000\n",
    "ukappa = qe_reconstruct(tmap,unlensed_cmb_power_2d,total_cmb_power_2d,ellmin,ellmax,modlmap,ly,lx)\n",
    "\n",
    "# Calculating the normalization is a bit involved, so for now we focus on comparing\n",
    "# the largest scales of the reconstruction with that of the input, where the normalization is\n",
    "# expected to be somewhat constant\n",
    "\n",
    "### Filter the reconstruction and the input kappa to the same ellrange (where the normalization\n",
    "### is expected to be relatively white)\n",
    "kellmin = 10\n",
    "kellmax = 100\n",
    "ukappa_f = lm.filter_map(ukappa,lm.kmask(modlmap*0.+1.,modlmap,kellmin,kellmax))\n",
    "ikappa_f = lm.filter_map(kappa,lm.kmask(modlmap*0.+1.,modlmap,kellmin,kellmax))\n",
    "\n",
    "# Plot the filtered maps for comparison\n",
    "plt.imshow(ukappa_f,cmap='gray')\n",
    "plt.show()\n",
    "plt.clf()\n",
    "plt.imshow(ikappa_f,cmap='gray')\n",
    "plt.show()\n",
    "plt.clf()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should be able to see a correspondence of peaks and valleys by eye in the above reconstruction if you have implemented it correctly. To be absolute sure, we will next cross-correlate the reconstruction with the input.\n",
    "\n",
    "## Cross power with input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "EXERCISE: This looks good by eye! But are we sure it's not just a coincidence? The systematic way of checking that your lensing pipeline is working is to cross-correlate your reconstruction with the input kappa you provided. Calculate the binned 1d cross power and confirm it is non-zero. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Calculate binned 1d cross-power between reconstruction and input\n",
    "ell_max = 1000\n",
    "delta_ell = 10\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}