CMB_School_Part_07.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jeffmcm1977/CMBAnalysis_SummerSchool/blob/master/CMB_School_Part_07.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Map Analysis III\n",
    "\n",
    "### Jeff McMahon and Renée Hložek \n",
    "\n",
    "# CMB Polarization\n",
    "\n",
    "In this section we generalize our code to handle polarization.  This means creating T, E, and B maps following the same general procedure we used for generating the T maps (but with a corelation between T and E).  We then transform these maps into Q and U maps.  These calculations are carried out in the flat sky approximation, but the treatment in a curved sky is analogous. \n",
    "\n",
    "After simulating these maps we present the problems associated with a cut sky and the associated E-B mode mixing.  We then present Kendrick Smith's no-leak pseudo-$C_\\ell$ estimator as one of the possible solutions to this mixing problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Incorporating previous bits of code"
   ]
  },
  {
   "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 matplotlib.mlab as mlab\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "import astropy.io.fits as fits\n",
    "import numpy as np\n",
    "\n",
    "import constants as cs # the constants module\n",
    "import cmb_modules # the module of functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also read in the constants for map noise, beam shape, and point sources so we can use these in the excercises"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = cs.N\n",
    "N_iterations = cs.N_iterations\n",
    "c_min = cs.c_min\n",
    "c_max = cs.c_max\n",
    "X_width =cs.X_width\n",
    "Y_width = cs.Y_width\n",
    "beam_size_fwhp = cs.beam_size_fwhp\n",
    "\n",
    "pix_size = cs.pix_size\n",
    "\n",
    "Number_of_Sources  = cs.Number_of_Sources\n",
    "Amplitude_of_Sources = cs.Amplitude_of_Sources\n",
    "Number_of_Sources_EX = cs.Number_of_Sources_EX\n",
    "Amplitude_of_Sources_EX = cs.Amplitude_of_Sources_EX\n",
    "\n",
    "Number_of_SZ_Clusters  = cs.Number_of_SZ_Clusters\n",
    "Mean_Amplitude_of_SZ_Clusters = cs.Mean_Amplitude_of_SZ_Clusters\n",
    "SZ_beta = cs.SZ_beta\n",
    "SZ_Theta_core = cs.SZ_Theta_core\n",
    "\n",
    "white_noise_level = cs.white_noise_level\n",
    "atmospheric_noise_level = 0 #cs.atmospheric_noise_level\n",
    "one_over_f_noise_level = 0 #cs.one_over_f_noise_level\n",
    "\n",
    "## we will need a window function below, so we create it here\n",
    "window = (cmb_modules.cosine_window(N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in and plot the TT, EE, BB, and TE power spectra\n",
    "\n",
    "With polarization there are four cosmologial power spectra to work with.   As you can see here, the main challenge is handling four times the book keeping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbig = 10000\n",
    "# read in the input CMB spectra\n",
    "ell, DlTT,DlEE,DlBB, DlTE= np.loadtxt(\"CMB_fiducial_totalCls.dat\", usecols=(0, 1,2,3,4), unpack=True) \n",
    "\n",
    "##\n",
    "ell_big = np.arange(nbig)\n",
    "DlTT_big = np.zeros(nbig)\n",
    "DlTT_big[ell.astype(int)] = DlTT\n",
    "DlEE_big = np.zeros(nbig)\n",
    "DlEE_big[ell.astype(int)] = DlEE\n",
    "DlBB_big = np.zeros(nbig)\n",
    "DlBB_big[ell.astype(int)] = DlBB\n",
    "DlTE_big = np.zeros(nbig)\n",
    "DlTE_big[ell.astype(int)] = DlTE\n",
    "\n",
    "ell = ell_big\n",
    "DlTT = DlTT_big + 1e-3   ### the 1e-3 factor maps plotting easy\n",
    "DlEE = DlEE_big + 1e-3\n",
    "DlBB = DlBB_big + 1e-3\n",
    "DlTE = DlTE_big\n",
    "\n",
    "Temp_point_source_spectrum = DlTT[3000]*(ell/3000.)**2.\n",
    "Pol_point_source_spectrum = DlEE[4500]*(ell/4500.)**2.\n",
    "\n",
    "DlTT_PS = DlTT + Temp_point_source_spectrum   ### these are used for computing the transer functions\n",
    "DlEE_PS = DlEE + Pol_point_source_spectrum\n",
    "DlBB_PS = DlBB + Pol_point_source_spectrum\n",
    "\n",
    "\n",
    "plt.semilogy(ell,DlTT,'r')\n",
    "plt.semilogy(ell,DlEE,'g')\n",
    "plt.semilogy(ell,DlBB,'b')\n",
    "plt.title('TT (red), EE (green), BB (blue) spectra')\n",
    "plt.ylabel('$D^{XX}_{\\ell}$ [$\\mu$K$^2$]')\n",
    "plt.xlabel('$\\ell$')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(ell,DlTE,'y')\n",
    "plt.ylabel('$D^{TE}_{\\ell}$ [$\\mu$K$^2$]')\n",
    "plt.xlabel('$\\ell$')\n",
    "plt.title('TE spectrum')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These spectra probably look familiar to you.   If you wish to explore the cosmological dependence of these spectra, you can simulate these spectra using CAMB on the NASA Lambda website just as we did in Part two of this course.  A key apsect for simmulation and analysis is that the information in the TT, EE, and BB are entirely contained in the T, E, and B maps, respectively;  while the TE information is contained within the T and E maps through a paritial correlation that comes about from how the T and E fluctuations are sourced in the early universe.  We will have to build in this correlation when we simulate the triplet of T, E, and B maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate T, E, and B maps\n",
    "\n",
    "Here we introduce a function to create T, Q, and U maps.   To do this we follow the same procedure for creating T, E and B maps that we introduced in Part two of this course with the exception of introducing a T-E correlation.  We then transform into Q and U maps.\n",
    "\n",
    "The correlation is introduced by (1) generating Gaussian random fields for the T, B, and uncorrelated part of the E maps; (2) separating the E-mode spectrum into a piece that is correlated with T: \n",
    "$$C_\\ell^{EE, correlated} = r^2_\\ell C_\\ell^{EE} $$ \n",
    "where $$r_\\ell = \\frac {(C_\\ell^{TE})} {\\sqrt{C_\\ell^{EE}C_\\ell^{TT}}},$$ and a part that has represents the amplitude of the uncorrelated fluctuations (chosen such that the correlated and uncorelated E- fluctuations reproduce the EE spectrum: $$C_\\ell^{uncorrelated} = (1-r^2_\\ell) C_\\ell^{EE};$$ and then generating the E- map as a sum of realizations of these two pieces.\n",
    "\n",
    "\n",
    "\n",
    "The relation between E-B and Q-U maps is given by a position dependent rotation in fourier space.  We introduce an angle $\\psi= \\arctan \\frac {k_Y} {k_X}$, where $k_X$ and $k_Y$ are the wave numbers in fourier space.   With this angle in had the relation between E-B and Q-U is:\n",
    "$$\\tilde Q  = \\tilde E \\cos 2 \\psi - \\tilde B \\sin 2 \\psi$$\n",
    "$$\\tilde U  = \\tilde E \\sin 2 \\psi + \\tilde B \\cos 2 \\psi$$\n",
    "\n",
    "The ~ over the Q, U, E, and B are to remind you that these quantities are in Fourier space.  \n",
    "\n",
    "As this is a simple rotation, the inverse transfom is:\n",
    "$$\\tilde E  = \\tilde Q \\cos 2 \\psi + \\tilde U \\sin 2 \\psi$$\n",
    "$$\\tilde B  = -\\tilde Q \\sin 2 \\psi + \\tilde U \\cos 2 \\psi$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_CMB_maps(N,pix_size,ell,DlTT,DlEE,DlTE,DlBB):\n",
    "    \"makes a realization of a simulated CMB sky map\"\n",
    "\n",
    "    # convert Dl to Cl, we use np.divide to avoid dividing by zero.\n",
    "    dell = ell * (ell + 1) / 2 / np.pi\n",
    "    ClTT = np.divide(DlTT, dell, where=ell>1)\n",
    "    ClEE = np.divide(DlEE, dell, where=ell>1)\n",
    "    ClTE = np.divide(DlTE, dell, where=ell>1)\n",
    "    ClBB = np.divide(DlBB, dell, where=ell>1)\n",
    "    \n",
    "    # set the \\ell = 0 and \\ell =1 modes to zero as these are unmeasurmable and blow up with the above transform\n",
    "    ClTT[0:2] = 0.\n",
    "    ClEE[0:2] = 0.\n",
    "    ClTE[0:2] = 0.\n",
    "    ClBB[0:2] = 0.\n",
    "\n",
    "    # separate the correlated and uncorrelated part of the EE spectrum\n",
    "    correlated_part_of_E = np.divide(ClTE, np.sqrt(ClTT), where=ell>1)\n",
    "    uncorrelated_part_of_EE = ClEE - np.divide(ClTE**2., ClTT, where=ell>1)\n",
    "    \n",
    "    correlated_part_of_E[0:2] = 0.\n",
    "    uncorrelated_part_of_EE[0:2] = 0.\n",
    "    \n",
    "    # make a 2d coordinate system\n",
    "    ones = np.ones(N)\n",
    "    inds  = (np.arange(N) - N/2.) /(N-1.)\n",
    "    X = np.outer(ones,inds)\n",
    "    Y = np.transpose(X)\n",
    "    R = np.sqrt(X**2. + Y**2.)\n",
    "    ang = np.arctan2(Y,X)   ## we now need this angle to handle the EB <--> QU rotation\n",
    "    \n",
    "    # now make a set of 2d CMB masks for the T, E, and B maps\n",
    "    ell_scale_factor = 2. * np.pi / (pix_size/60. * np.pi/180.)\n",
    "    ell2d = R * ell_scale_factor\n",
    "    ClTT_expanded = np.zeros(int(ell2d.max())+1)\n",
    "    ClTT_expanded[0:(ClTT.size)] = ClTT\n",
    "    ClEE_uncor_expanded = np.zeros(int(ell2d.max())+1)\n",
    "    ClEE_uncor_expanded[0:(uncorrelated_part_of_EE.size)] = uncorrelated_part_of_EE\n",
    "    ClE_corr_expanded = np.zeros(int(ell2d.max())+1)\n",
    "    ClE_corr_expanded[0:(correlated_part_of_E.size)] = correlated_part_of_E\n",
    "    ClBB_expanded = np.zeros(int(ell2d.max())+1)\n",
    "    ClBB_expanded[0:(ClBB.size)] = ClBB\n",
    "    CLTT2d = ClTT_expanded[ell2d.astype(int)]\n",
    "    ClEE_uncor_2d = ClEE_uncor_expanded[ell2d.astype(int)]\n",
    "    ClE_corr2d = ClE_corr_expanded[ell2d.astype(int)]\n",
    "    CLBB2d = ClBB_expanded[ell2d.astype(int)]\n",
    "    \n",
    "    # now make a set of gaussian random fields that will be turned into the CMB maps\n",
    "    randomn_array_for_T = np.fft.fft2(np.random.normal(0,1,(N,N)))\n",
    "    randomn_array_for_E = np.fft.fft2(np.random.normal(0,1,(N,N))) \n",
    "    randomn_array_for_B = np.fft.fft2(np.random.normal(0,1,(N,N))) \n",
    "    \n",
    "    ## make the T, E, and B maps by multiplying the masks against the random fields\n",
    "    FT_2d = np.sqrt(CLTT2d) * randomn_array_for_T\n",
    "    FE_2d = np.sqrt(ClEE_uncor_2d) * randomn_array_for_E + ClE_corr2d* randomn_array_for_T\n",
    "    FB_2d = np.sqrt(CLBB2d) * randomn_array_for_B\n",
    "    \n",
    "    ## now conver E abd B to Q and U\n",
    "    FQ_2d = FE_2d* np.cos(2.*ang) - FB_2d * np.sin(2. *ang)\n",
    "    FU_2d = FE_2d* np.sin(2.*ang) + FB_2d * np.cos(2. *ang)\n",
    "    \n",
    "    ## convert from fourier space to real space\n",
    "    CMB_T = np.fft.ifft2(np.fft.fftshift(FT_2d)) /(pix_size /60.* np.pi/180.)\n",
    "    CMB_T = np.real(CMB_T)\n",
    "    CMB_Q = np.fft.ifft2(np.fft.fftshift(FQ_2d)) /(pix_size /60.* np.pi/180.)\n",
    "    CMB_Q = np.real(CMB_Q)\n",
    "    CMB_U = np.fft.ifft2(np.fft.fftshift(FU_2d)) /(pix_size /60.* np.pi/180.)\n",
    "    CMB_U = np.real(CMB_U)\n",
    "\n",
    "    ## optional code for spitting out E and B maps \n",
    "    CMB_E = np.fft.ifft2(np.fft.fftshift(FE_2d)) /(pix_size /60.* np.pi/180.)\n",
    "    CMB_E = np.real(CMB_E)\n",
    "    CMB_B = np.fft.ifft2(np.fft.fftshift(FB_2d)) /(pix_size /60.* np.pi/180.)\n",
    "    CMB_B = np.real(CMB_B)\n",
    "    \n",
    "    ## return the maps\n",
    "    return(CMB_T,CMB_Q,CMB_U,CMB_E,CMB_B)\n",
    "  ###############################\n",
    "\n",
    "\n",
    "    \n",
    "## make a CMB T map\n",
    "CMB_T,CMB_Q,CMB_U,CMB_E,CMB_B = make_CMB_maps(N,pix_size,ell,DlTT,DlEE,DlTE,DlBB)\n",
    "\n",
    "\n",
    "p = cmb_modules.Plot_CMB_Map(CMB_T,c_min,c_max,X_width,Y_width)\n",
    "p = cmb_modules.Plot_CMB_Map(CMB_Q,c_min/20.,c_max/20.,X_width,Y_width)\n",
    "p = cmb_modules.Plot_CMB_Map(CMB_U,c_min/20.,c_max/20.,X_width,Y_width)\n",
    "p = cmb_modules.Plot_CMB_Map(CMB_E,c_min/20.,c_max/20.,X_width,Y_width)\n",
    "p = cmb_modules.Plot_CMB_Map(CMB_B,c_min/200.,c_max/200.,X_width,Y_width)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the RMS of these maps: the T is ~100 $\\mu$K, the Q and U are ~ 4.5 $\\mu$K, the E is ~ 6.5 $\\mu$K, and the B map rms is $<1 \\mu K$.  Knowing these values is a good useful starting point for debuging simulations, calibration of real maps, and tracking the impact of filtering on maps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Add point sources and the SZ to these maps, next convolve with a beam, and add instrumental noise.   Assume the point sources have a fixed $3\\%$ polarization fraction with random polarization angles and the SZ is fully unpolarized.  The noise in the polarization maps is $\\sqrt 2$ larger than the temperature noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Directly visualize the polarization field\n",
    "\n",
    "In the above we have visualized the polarization field using two different methods: a decomposition into Stokes Q and U maps, and a decomposition into E- and B-modes. Sometimes it is useful to instead visualize the polarization field as a vector field of (headless) vectors that directly show the orientation and amplitude of the (linearly) polarized light.\n",
    "\n",
    "Below we generate a plot like this using vectors of length $p = \\sqrt{Q^2 + U^2}$ and orientation $\\psi_p$ given by $2 \\psi_p = \\arctan\\left( \\frac{U}{Q}\\right)$. We overlay this plot over the E mode map to demonstrate that the polarization field (which is dominated by E-modes) closely traces this map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_quiver(Q, U, X_width, Y_width, background=None):\n",
    "    '''Visualize Stokes Q, U as headless vectors'''\n",
    "    from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "    \n",
    "    # Smooth maps for nicer images\n",
    "    fwhm_pix = 30\n",
    "    Q = cmb_modules.convolve_map_with_gaussian_beam(N,pix_size,fwhm_pix,Q)\n",
    "    U = cmb_modules.convolve_map_with_gaussian_beam(N,pix_size,fwhm_pix,U)\n",
    "    if background is not None:\n",
    "        # If provided, we overplot the vectors on top of the smoothed background.\n",
    "        background = cmb_modules.convolve_map_with_gaussian_beam(N,pix_size,fwhm_pix,background)\n",
    "    \n",
    "    Q = Q[::int(fwhm_pix),::int(fwhm_pix)]\n",
    "    U = U[::int(fwhm_pix),::int(fwhm_pix)]\n",
    "    \n",
    "    p_amp = np.sqrt(Q**2 + U**2)\n",
    "    ang = np.arctan2(U, Q) / 2.\n",
    "    \n",
    "    u = p_amp * np.cos(ang)\n",
    "    v = p_amp * np.sin(ang)\n",
    "\n",
    "    x = np.linspace(0,X_width,u.shape[1])\n",
    "    y = np.linspace(0,X_width,u.shape[0])    \n",
    "        \n",
    "    fig, ax = plt.subplots(figsize=(10,10))\n",
    "    if background is not None:\n",
    "        im = ax.imshow(background, interpolation='bilinear', origin='lower',cmap=cm.RdBu_r,\n",
    "                       extent=([0,X_width,0,Y_width]))\n",
    "        divider = make_axes_locatable(ax)\n",
    "        cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "        cbar = plt.colorbar(im, cax=cax)\n",
    "        cbar.set_label('temperature [uK]', rotation=270)\n",
    "        \n",
    "    ax.quiver(x, y, u, v, headwidth=1, headlength=0, pivot='mid', units='xy',\n",
    "              scale_units='xy', scale=2 * p_amp.max(), linewidth=1)\n",
    "        \n",
    "    ax.set_ylabel('angle $[^\\circ]$')\n",
    "    ax.set_xlabel('angle $[^\\circ]$')\n",
    "    \n",
    "    plt.show(fig)\n",
    "\n",
    "plot_quiver(CMB_Q,CMB_U,X_width,Y_width,background=CMB_E)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Generate a polarization field dominated by B-mode polarization and plot the polarization vectors on top of the B-mode map. What is the difference between the E-mode field and the B-mode field?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Calculate the BB power spectrum using the code and Monte Carlo techniques we explored in Part Four.  Compare your results to the input spectra. Make sure to include a window (i.e. pixel mask).  \n",
    "\n",
    "\n",
    "<font color='red'>hint: </font> You can make use of the `QU2EB()` function below to convert Q and U maps to E and B maps. Because real data would be provided as Q and U maps, a pixel mask should be applied to the Q and U maps before converting to E and B."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def QU2EB(N,pix_size,Qmap,Umap,):\n",
    "    '''Calcalute E, B maps given input Stokes Q, U maps'''\n",
    "    \n",
    "    # Create 2d Fourier coordinate system.\n",
    "    ones = np.ones(N)\n",
    "    inds  = (np.arange(N) - N/2.) /(N-1.)\n",
    "    kX = np.outer(ones,inds) / (pix_size/60. * np.pi/180.)\n",
    "    kY = np.transpose(kX)\n",
    "    ang = np.arctan2(kY,kX)\n",
    "  \n",
    "    # Convert to Fourier domain.\n",
    "    fQ = np.fft.fftshift(np.fft.fft2(Qmap))\n",
    "    fU = np.fft.fftshift(np.fft.fft2(Umap))\n",
    "    \n",
    "    # Convert Q, U to E, B in Fourier domain.\n",
    "    fE = fQ * np.cos(2.*ang) + fU * np.sin(2. *ang)\n",
    "    fB = - fQ * np.sin(2.*ang) + fU * np.cos(2. *ang)\n",
    "    \n",
    "    # Convert E, B from Fourier to real space.\n",
    "    Emap = np.real(np.fft.ifft2(np.fft.fftshift(fE)))\n",
    "    Bmap = np.real(np.fft.ifft2(np.fft.fftshift(fB)))\n",
    "\n",
    "    return Emap, Bmap\n",
    "\n",
    "# Convert masked Q, U to E and B.\n",
    "Emap, Bmap = QU2EB(N,pix_size,CMB_Q * window,CMB_U * window)\n",
    "\n",
    "cmb_modules.Plot_CMB_Map(window,0,1,X_width,Y_width)\n",
    "cmb_modules.Plot_CMB_Map(Emap,-20,20,X_width,Y_width)\n",
    "cmb_modules.Plot_CMB_Map(Bmap,-2,2,X_width,Y_width)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your notes go here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Create a simulated map with $C_\\ell^{BB} = 0$ Calculate the BB power spectrum of this map using the estimates for the multiplicitive and addative bias you calculated in the previous excercise with non-zero $C_\\ell^{BB}$.   Why isn't your result consistent with $C_\\ell^{BB} = 0$   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "your notes go here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Calculate the TE power spectrum by adapting the techniques from Part Four."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "your notes go here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Kendrick Smith's no-leak B-mode estimator\n",
    "\n",
    "Kendrick Smith wrote a nice <a href=https://arxiv.org/abs/astro-ph/0511629>paper</a> describing his solution to the E-B mixing problem on cut skies.  He showed that the leakage is related to the window function used to apodize the maps before moving to fourier space and identified corrections (products of derivitives of the window funciton and the Q and U maps in fourier space) that can be added to create a clean estimate of the B-modes.\n",
    "\n",
    "We implement his method here so that you can apply it to your simulations and see that it cleans up the B-mode estimate you produced above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kendric_method_precompute_window_derivitives(win,pix_size):\n",
    "    delta = pix_size * np.pi /180. /60.\n",
    "    dwin_dx =    ((-1.) * np.roll(win,-2,axis =1)      +8. * np.roll(win,-1,axis =1)     - 8. *np.roll(win,1,axis =1)      +np.roll(win,2,axis =1) ) / (12. *delta)\n",
    "    dwin_dy =    ((-1.) * np.roll(win,-2,axis =0)      +8. * np.roll(win,-1,axis =0)     - 8. *np.roll(win,1,axis =0)      +np.roll(win,2,axis =0) ) / (12. *delta)\n",
    "    d2win_dx2 =  ((-1.) * np.roll(dwin_dx,-2,axis =1)  +8. * np.roll(dwin_dx,-1,axis =1) - 8. *np.roll(dwin_dx,1,axis =1)  +np.roll(dwin_dx,2,axis =1) ) / (12. *delta)\n",
    "    d2win_dy2 =  ((-1.) * np.roll(dwin_dy,-2,axis =0)  +8. * np.roll(dwin_dy,-1,axis =0) - 8. *np.roll(dwin_dy,1,axis =0)  +np.roll(dwin_dy,2,axis =0) ) / (12. *delta)\n",
    "    d2win_dxdy = ((-1.) * np.roll(dwin_dy,-2,axis =1)  +8. * np.roll(dwin_dy,-1,axis =1) - 8. *np.roll(dwin_dy,1,axis =1)  +np.roll(dwin_dy,2,axis =1) ) / (12. *delta)\n",
    "    return(dwin_dx,dwin_dy,d2win_dx2,d2win_dy2,d2win_dxdy)\n",
    "\n",
    "def kendrick_method_TQU_to_fourier_TEB(N,pix_size,Tmap,Qmap,Umap,window,dwin_dx,dwin_dy,d2win_dx2,d2win_dy2,d2win_dxdy):\n",
    "    ### the obvious FFTs\n",
    "    fft_TxW = np.fft.fftshift(np.fft.fft2(Tmap * window))\n",
    "    fft_QxW = np.fft.fftshift(np.fft.fft2(Qmap * window))\n",
    "    fft_UxW = np.fft.fftshift(np.fft.fft2(Umap * window))\n",
    "\n",
    "    ### the less obvious FFTs that go into the no-leak estiamte\n",
    "    fft_QxdW_dx = np.fft.fftshift(np.fft.fft2(Qmap * dwin_dx))\n",
    "    fft_QxdW_dy = np.fft.fftshift(np.fft.fft2(Qmap * dwin_dy))\n",
    "    fft_UxdW_dx = np.fft.fftshift(np.fft.fft2(Umap * dwin_dx))\n",
    "    fft_UxdW_dy = np.fft.fftshift(np.fft.fft2(Umap * dwin_dy))\n",
    "    fft_QU_HOT  = np.fft.fftshift(np.fft.fft2( (2. * Qmap * d2win_dxdy) + Umap * (d2win_dy2 - d2win_dx2) ))\n",
    "    \n",
    "    ### generate the polar coordinates needed to cary out the EB-QU conversion\n",
    "    ones = np.ones(N)\n",
    "    inds  = (np.arange(N) - N/2.) /(N-1.)\n",
    "    X = np.outer(ones,inds)\n",
    "    Y = np.transpose(X)\n",
    "    R = np.sqrt(X**2. + Y**2. + 1e-9)  ## the small offset regularizes the 1/ell factors below\n",
    "    ang =  np.arctan2(Y,X)\n",
    "    ell_scale_factor = 2. * np.pi / (pix_size/60. * np.pi/180.)\n",
    "    ell2d = R * ell_scale_factor\n",
    "        \n",
    "    #p=Plot_CMB_Map(np.real( ang),-np.pi,np.pi,N,N)\n",
    "    \n",
    "    \n",
    "    ### now compute the estimator\n",
    "    fTmap = fft_TxW\n",
    "    fEmap = fft_QxW * np.cos(2. * ang) + fft_UxW * np.sin(2. * ang)\n",
    "    fBmap = (fft_QxW * (-1. *np.sin(2. * ang)) + fft_UxW * np.cos(2. * ang))  ## this line is the nominal B estimator\n",
    "    fBmap = fBmap - complex(0,2.) / ell2d * (fft_QxdW_dx * np.sin(ang) + fft_QxdW_dy * np.cos(ang))\n",
    "    fBmap = fBmap - complex(0,2.) / ell2d * (fft_UxdW_dy * np.sin(ang) - fft_UxdW_dx * np.cos(ang))\n",
    "    fBmap = fBmap +  ell2d**(-2.) * fft_QU_HOT\n",
    "\n",
    "    ### return the complex fourier maps in 2d\n",
    "    return(fTmap,fEmap,fBmap)\n",
    "\n",
    "    \n",
    "    \n",
    "dwin_dx,dwin_dy,d2win_dx2,d2win_dy2,d2win_dxdy = kendric_method_precompute_window_derivitives(window,pix_size)\n",
    "\n",
    "### plots of the window function derivatives\n",
    "#p=Plot_CMB_Map(dwin_dx,-10,10,X_width,Y_width)\n",
    "#p=Plot_CMB_Map(dwin_dy,-10,10,X_width,Y_width)\n",
    "#p=Plot_CMB_Map(d2win_dx2,-10,10,X_width,Y_width)\n",
    "#p=Plot_CMB_Map(d2win_dy2,-10,10,X_width,Y_width)\n",
    "#p=Plot_CMB_Map(d2win_dxdy,-10,10,X_width,Y_width)\n",
    "\n",
    "fTmap,fEmap,fBmap = kendrick_method_TQU_to_fourier_TEB(N,pix_size,CMB_T,CMB_Q,CMB_U,window,dwin_dx,dwin_dy,d2win_dx2,d2win_dy2,d2win_dxdy)\n",
    "\n",
    "p=cmb_modules.Plot_CMB_Map(np.real( (np.fft.ifft2(np.fft.fftshift(fTmap)))),c_min,c_max,X_width,Y_width)\n",
    "p=cmb_modules.Plot_CMB_Map(np.real( (np.fft.ifft2(np.fft.fftshift(fEmap)))),c_min/20.,c_max/20.,X_width,Y_width)\n",
    "p=cmb_modules.Plot_CMB_Map(np.real( (np.fft.ifft2(np.fft.fftshift(fBmap))-np.mean(np.fft.ifft2(np.fft.fftshift(fBmap))))),c_min/400.,c_max/400.,X_width,Y_width)\n",
    "               \n",
    "               \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Calculate the TT, EE, BB, TE, TB, and EB power spectra using the code and techniques we explored in Part Four but with the Kendrick-method for B-modes.  Compare your results to the input spectra and to the spectra you produced with the naive E-B separation.  Pay special attention to the angular scale dependence of the differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "your comments go here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Note on E-B leakage correction\n",
    "\n",
    "The no-leak B-mode estimator in the above cells removes spurious B-modes at the level of the map. This is useful when the B-mode contribution has to be measured accurately from a small patch of the sky. In cases where some E-to-B leakage can be tolerated, the leakage can also be corrected for at the level of the power spectrum. If you want to learn more about this method you can work through [notebook 8](https://nbviewer.jupyter.org/github/ACTCollaboration/DR4_DR5_Notebooks/blob/bbc7767607c0e51a0a909ffdec617ce6c64d8388/Notebooks/Section_8_power_spectra_part_2.ipynb) of the DR4_DR5_Notebooks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>ADVANCED EXERCISE: </font> Take the `lens_map` function from the lensing notebook (part 10) and lens the Q and U maps of a polarization field that has no B-mode contribution. Plot the resulting B-mode power spectrum and compare to the CAMB B-mode power spectrum."
   ]
  },
  {
   "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
}