{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read and process Planck CMB power spectra with `healpy`\n",
    "- categories: [healpy,cmb,cosmology]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will load CMB power spectra generated by the cosmological parameters measured by Planck, we will create a realization of the Spherical Harmonics coefficients $a_{\\ell m}$, save it to disk, and use them to create maps at different resolutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import healpy as hp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter out all warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(plt.style.available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('seaborn-talk')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we go to the Planck data release page at IPAC, I prefer plain web pages instead of the funky javascript interface of the Planck Legacy Archive.\n",
    "The spectra are relegated to the **Ancillary data** section:\n",
    "\n",
    "<https://irsa.ipac.caltech.edu/data/Planck/release_3/ancillary-data/>\n",
    "\n",
    "We can choose one of the available spectra, I would like the spectrum from theory, `BaseSpectra`, and download it locally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://irsa.ipac.caltech.edu/data/Planck/release_3/ancillary-data/cosmoparams/COM_PowerSpect_CMB-base-plikHM-TTTEEE-lowl-lowE-lensing-minimum-theory_R3.01.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, I couldn't find any documentation about the format of those spectra, there are no units and we don't know if those spectra are just the $C_\\ell$ or $\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell$,\n",
    "it is probably a standard output of the `CAMB` software, anyway I prefer to cross check with some public plot.\n",
    "\n",
    "`pandas` has trouble with the `#` at the beginning of the title line, the easiest is to edit that out with a file editor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -3 COM_PowerSpect_CMB-base-plikHM-TTTEEE-lowl-lowE-lensing-minimum-theory_R3.01.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should not have `#` at the beginning of the first line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cl = pd.read_csv(\"COM_PowerSpect_CMB-base-plikHM-TTTEEE-lowl-lowE-lensing-minimum-theory_R3.01.txt\",\n",
    "                       delim_whitespace=True, index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cl.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cl.plot(logx=True, logy=True, grid=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare this to one of the plots from the Planck wiki:\n",
    "\n",
    "![](https://wiki.cosmos.esa.int/planckpla/images/a/a5/Mission_spectrum.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cl.TT.plot(logx=False, logy=False, grid=True)\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "\n",
    "plt.xlim([50, 2500]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Very good, they agree, therefore the data are in $\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell [\\mu K^2] $,\n",
    "let's transform to $C_\\ell [K^2]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cl.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(input_cl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmax = input_cl.index[-1]\n",
    "print(lmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl = input_cl.divide(input_cl.index * (input_cl.index+1) / (np.pi*2), axis=\"index\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl /= 1e12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\ell$ starts at 2, but all the `healpy` tools expect an array starting from zero, so let's extend that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl = cl.reindex(np.arange(0, lmax+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl = cl.fillna(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate a CMB map realization\n",
    "\n",
    "The power spectrum gives us the distribution of power with the angular scale, if we want to create a map we could use `hp.synfast`, however the realization will be different each time that we run `synfast`, we could always use the same seed to make it reproducible.\n",
    "However, in case we want to have the same realization of a CMB map at different resolutions ($N_{side}$), it is better to generate a realization of the spectrum in Spherical Harmonics space, creating a set of $a_{\\ell m}$, and then when needed transform them to a map with `hp.alm2map`.\n",
    "\n",
    "We also want to make sure that the $a_{\\ell m}$ have the same $\\ell_{max}$ of the input $C_\\ell$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 583\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm = hp.synalm((cl.TT, cl.EE, cl.BB, cl.TE), lmax=lmax, new=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "high_nside = 1024\n",
    "cmb_map = hp.alm2map(alm, nside=high_nside, lmax=lmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can double check that we got the order of magnitude correct, for example we can compare with one of the [published Planck maps](https://www.cosmos.esa.int/documents/387566/425793/2015_SMICA_CMB/c8c4c802-4b76-49da-b80a-0fb8d02c62b7?t=1423083319437), the scale is $\\pm 300 \\mu K$: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(cmb_map[0], min=-300*1e-6, max=300*1e-6, unit=\"K\", title=\"CMB Temperature\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Polarization is generally an order of magnitude lower:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in cmb_map[1:]:\n",
    "    hp.mollview(m, min=-30*1e-6, max=30*1e-6, unit=\"K\", title=\"CMB Polarization\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Double check the spectra of the maps\n",
    "\n",
    "As a final check, we can compare the spectra of the output maps to the input spectra."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_from_map = hp.anafast(cmb_map, lmax=lmax, use_pixel_weights=True) * 1e12 # in muK^2\n",
    "ell = np.arange(cl_from_map.shape[1])\n",
    "cl_from_map *= ell * (ell+1) / (2*np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.median(cl_from_map, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl_from_map[0], label=\"map TT\")\n",
    "input_cl.TT.plot(logx=False, logy=False, grid=True)\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cl.plot(logx=True, logy=True, grid=True)\n",
    "plt.plot(cl_from_map[0], label=\"map TT\")\n",
    "plt.plot(cl_from_map[1], label=\"map EE\")\n",
    "plt.plot(cl_from_map[2], label=\"map BB\")\n",
    "plt.plot(cl_from_map[3], label=\"map TE\")\n",
    "#plt.plot(cl_from_map[4], label=\"map EB\")\n",
    "#plt.plot(cl_from_map[5], label=\"map TB\")\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The scale is fine, we notice below that going through the discretization step of mapping the Spherical Harmonics into a map in pixel space and then trying to estimate the power spectrum again from that introduces some noise, on top of that, there is the issue of \"Cosmic variance\", i.e. the $C_\\ell$ estimates the variance of the $a_{\\ell m}$ coefficients, but large angular scales, low $\\ell$, there are only a small number of such coefficients, so the estimation itself is noisy [see this work by Benjamin Wandelt](https://www.slac.stanford.edu/econf/C030908/papers/THAT004.pdf):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl_from_map[0][:-2], label=\"map TT\")\n",
    "input_cl.TT.plot(logx=False, logy=False, grid=True)\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl_from_map[1][:-2], label=\"map EE\")\n",
    "input_cl.EE.plot(logx=False, logy=False, grid=True)\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl_from_map[2], label=\"map BB\")\n",
    "input_cl.BB.plot(logx=False, logy=False, grid=True)\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the $a_{\\ell m}$ to disk\n",
    "\n",
    "We can save in a standard FITS format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.write_alm(f\"Planck_bestfit_alm_seed_{seed}.fits\", alm, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls *alm*fits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a lower resolution map from the same $a_{\\ell m}$\n",
    "\n",
    "The $a_{\\ell m}$ can now be used to create the same realization of the CMB power spectrum at different resolution,\n",
    "it is not as straightforward as it might seem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_nside = 32\n",
    "low_nside_lmax = 3*low_nside - 1\n",
    "low_nside_cmb_map = hp.alm2map(alm, nside=low_nside)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_low_nside = hp.anafast(low_nside_cmb_map, use_pixel_weights=True, lmax=low_nside_lmax) * 1e12\n",
    "ell_low_nside = np.arange(cl_low_nside.shape[1])\n",
    "cl_low_nside *= ell_low_nside * (ell_low_nside+1) / (2*np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl_from_map[0], label=\"map TT\")\n",
    "plt.plot(cl_low_nside[0], label=\"low nside map TT\")\n",
    "input_cl.TT.plot(logx=False, logy=False, grid=True)\n",
    "plt.axhline(low_nside_cmb_map.std()**2)\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We notice that the spectra is dominated by noise. The issue arises from the fact that Spherical Harmonics transforms are well defined only for band-limited signals, that means that at high $\\ell$ the signal should be very low otherwise we are going to create artifacts, because we used $a_{\\ell m}$ with a $\\ell_{max}$ of ~2500 to create a map which could only represent a signal with a $\\ell_{max} < 3 N_{side} - 1$ (or $\\ell_{max} < 2 N_{side}$ to be conservative).\n",
    "\n",
    "Unfortunately the `lmax` argument of `hp.map2alm` doesn't help because it is designed to specify the $\\ell_{max}$ of the input $a_{\\ell m}$.\n",
    "\n",
    "We can fix the issue by zeroing the $a_{\\ell m}$ coefficients above a specific $\\ell$ using `hp.almxfl`, it automatically assumes that if a input $\\ell$ is not provided, it is zero, therefore we can just provide an array of ones with the length of our target $\\ell_{max}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip = np.ones(low_nside_lmax+1)\n",
    "alm_clipped = np.array([hp.almxfl(each, clip) for each in alm])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_nside_cmb_map_clipped = hp.alm2map(alm_clipped, nside=low_nside)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_low_nside_clipped = hp.anafast(low_nside_cmb_map_clipped, use_pixel_weights=True, lmax=low_nside_lmax) * 1e12\n",
    "cl_low_nside_clipped *= ell_low_nside * (ell_low_nside+1) / (2*np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl_from_map[0], label=\"map TT\", alpha=.5)\n",
    "plt.plot(cl_low_nside_clipped[0], label=\"low nside map TT clipped\", alpha=.8)\n",
    "input_cl.TT.plot(logx=False, logy=False, grid=True)\n",
    "plt.axvline(2*low_nside, color=\"black\", linestyle=\"--\", label=\"ell=2*N_side\")\n",
    "plt.xlim([0, 800])\n",
    "plt.ylabel(\"$\\dfrac{\\ell(\\ell+1)}{2\\pi} C_\\ell~[\\mu K^2]$\")\n",
    "plt.xlabel(\"$\\ell$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thanks to this the spectra is accurate at least up to $2 N_{side}$ and it looks has a small bias towards higher values at higher $\\ell$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can check by eye that the large scale features of the map (focus on a hot or cold spot) are the same in the 2 maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(cmb_map[0], title=\"CMB T map at $N_{side}$= \" + str(high_nside))\n",
    "hp.mollview(low_nside_cmb_map_clipped[0], title=\"CMB T map at $N_{side}$= \" + str(low_nside))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.write_map(f\"map_nside_{low_nside}_from_Planck_bestfit_alm_seed_{seed}_K_CMB.fits\", low_nside_cmb_map_clipped, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls *.fits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clip $a_{\\ell m}$ to a lower $\\ell_{max}$\n",
    "\n",
    "It was easy to use `hp.almxfl` to zero the coefficient above a threshold, however, the array remains very large, especially inconvenient if you need to save it to disk.\n",
    "\n",
    "We want instead to actually clip the coefficients using the `hp.Alm.getidx` function which gives us the indices of a specific $\\ell,m$ pair (they can be array):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lclip = low_nside_lmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "clipping_indices = []\n",
    "for m in range(lclip+1):\n",
    "    clipping_indices.append(hp.Alm.getidx(lmax, np.arange(m, lclip+1), m))\n",
    "print(clipping_indices[:2])\n",
    "clipping_indices = np.concatenate(clipping_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm_really_clipped = [each[clipping_indices] for each in alm]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_nside_cmb_map_really_clipped = hp.alm2map(alm_clipped, nside=low_nside)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(low_nside_cmb_map_clipped[0], title=\"zeroed - CMB T map at $N_{side}$= \" + str(low_nside))\n",
    "hp.mollview(low_nside_cmb_map_really_clipped[0], title=\"clipped - CMB T map at $N_{side}$= \" + str(low_nside))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(low_nside_cmb_map_clipped[0] - low_nside_cmb_map_really_clipped[0]).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.write_alm(f\"Planck_bestfit_alm_seed_{seed}_lmax_{lclip}_K_CMB.fits\", alm_really_clipped, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls *.fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SO",
   "language": "python",
   "name": "so"
  },
  "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}