CMB_ACT_Likelihood.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jeffmcm1977/CMBAnalysis_SummerSchool/blob/master/CMB_ACT_Likelihood.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CMB+ACT: Likelihoods and Sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is a combination of the work covered in the CMB Analysis Workbook Part 08, and the ACT DR4/5 Section 11 Likelihood notebook. We'll start from the basics of getting theory spectra from a set of cosmological parameters, through an example with the ACT likelihood, and MCMC sampling. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook Info | Value\n",
    "---|:--\n",
    "Dependencies | `numpy`, `scipy`, `healpy`, `matplotlib`, `pixell`, `pyactlike`, `camb`, `pandas`, `getdist`\n",
    "Data products | ACTPol_lcdm.paramnames.csv\n",
    "| ACTPol_lcdm.paramnames\n",
    "| ACTPol_lcdm_1.txt\n",
    "Memory usage | Low\n",
    "Contributors | Zack Li, Jo Dunkley, Maya Mallaby-Kay, Zach Atkins, Erminia Calabrese, Renee Hlozek\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1: Preliminaries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will start by pulling CAMB python so that we can get it running. Get pycamb from https://pypi.python.org/pypi/camb/0.1.1."
   ]
  },
  {
   "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",
    ")\n",
    "!python -c \"import pandas\" || python -m pip install pandas\n",
    "!python -c \"import camb\" || python -m pip install camb\n",
    "!python -c \"import getdist\" || python -m pip install getdist\n",
    "!python -c \"import pyactlike\" || python -m pip install pyactlike@git+https://github.com/ACTCollaboration/pyactlike"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import pyactlike\n",
    "import camb\n",
    "from camb import model, initialpower\n",
    "\n",
    "# Replace these paths with the paths to your data\n",
    "path_to_cmb_notebook_data = '.'\n",
    "path_to_act_notebook_data = '/home/zatkins/act/DR4/DR4_notebooks_live/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will start by pulling initialising the CAMB params structure that we will use later. \n",
    "#This is similar to how you would change the params.ini file \n",
    "\n",
    "#Set up a new set of parameters for CAMB\n",
    "pars = camb.CAMBparams()\n",
    "# The base cosmology model is set with these params, the others are all set to their default values\n",
    "pars.set_cosmology(H0=67.5, ombh2=0.022, omch2=0.122, mnu=0.06, omk=0, tau=0.06)\n",
    "# The initial power spectrum is set here, separately from the rest of cosmology\n",
    "pars.InitPower.set_params(ns=0.965, r=0)\n",
    "\n",
    "# Set how far in multipole we want the power spectra, and turn on defaults for the params.\n",
    "pars.set_for_lmax(4400, lens_potential_accuracy=0);\n",
    "\n",
    "#calculate results for these parameters \n",
    "# this is like \"running\" camb from the command line, and is the same as how it is done in cosmomc\n",
    "results = camb.get_results(pars)\n",
    "\n",
    "#get dictionary of CAMB power spectra\n",
    "powers =results.get_cmb_power_spectra(pars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see what the power spectra are:\n",
    "for name in powers: print(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now we want to plot the total lensed and unlensed CMB power spectra\n",
    "totCL=powers['total']\n",
    "CMBOutscale = 7.43e12\n",
    "unlensedCL=powers['unlensed_scalar']\n",
    "#Python CL arrays are all zero based (starting at L=0), Note L=0,1 entries will be zero by default.\n",
    "#The different CL are always in the order TT, EE, BB, TE (with BB=0 for unlensed scalar results).\n",
    "ls = np.arange(totCL.shape[0])\n",
    "fig, ax = plt.subplots(2,2, figsize = (15,15 ))\n",
    "ax[0,0].plot(ls,CMBOutscale*totCL[:,0], color='k', label='lensed')\n",
    "ax[0,0].plot(ls,CMBOutscale*unlensedCL[:,0], color='r', label='unlensed')\n",
    "ax[0,0].legend(loc=\"upper left\", bbox_to_anchor=[0, 1],\n",
    "           ncol=2, frameon=False)\n",
    "ax[0,0].set_xlabel(r'Multipole $\\ell$', fontsize=20)\n",
    "ax[0,0].set_ylabel(r'$\\ell(\\ell+1)C_\\ell/2\\pi$', fontsize=20)\n",
    "ax[0,0].set_title(r'$TT$')\n",
    "ax[0,1].axis('off')\n",
    "ax[1,0].plot(ls,CMBOutscale*totCL[:,1], color='k', label='lensed')\n",
    "ax[1,0].plot(ls,CMBOutscale*unlensedCL[:,1], color='m', label='unlensed')\n",
    "ax[1,0].legend(loc=\"upper left\", bbox_to_anchor=[0, 1],\n",
    "           ncol=2,frameon=False)\n",
    "ax[1,0].set_title(r'$EE$')\n",
    "ax[1,0].set_xlabel(r'Multipole $\\ell$', fontsize=20)\n",
    "ax[1,0].set_ylabel(r'$\\ell(\\ell+1)C_\\ell/2\\pi$', fontsize=20)\n",
    "ax[1,1].plot(ls,CMBOutscale*totCL[:,3], color='k', label='lensed')\n",
    "ax[1,1].plot(ls,CMBOutscale*unlensedCL[:,3], color='c', label='unlensed')\n",
    "ax[1,1].legend(loc=\"upper left\", bbox_to_anchor=[0, 1],\n",
    "           ncol=2, frameon=False)\n",
    "ax[1,1].set_title(r'$TE$');\n",
    "ax[1,1].set_xlabel(r'Multipole $\\ell$', fontsize=20)\n",
    "ax[1,1].set_ylabel(r'$\\ell(\\ell+1)C_\\ell/2\\pi$', fontsize=20)\n",
    "\n",
    "for ax in ax.reshape(-1): ax.set_xlim([2,4400])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now that we have run camb once and know how to do it, we can also compute the Cls over a range of parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#You can calculate spectra for different primordial power spectra without recalculating everything\n",
    "#for example, let's plot the BB spectra as a function of r\n",
    "pars.WantTensors = True\n",
    "results = camb.get_transfer_functions(pars)\n",
    "lmax=2000\n",
    "rs = np.linspace(0,0.2,6)\n",
    "for r in rs:\n",
    "    inflation_params = initialpower.InitialPowerLaw()\n",
    "    inflation_params.set_params(ns=0.96, r=r)\n",
    "    results.power_spectra_from_transfer(inflation_params)\n",
    "    cl = results.get_total_cls(lmax)\n",
    "    plt.loglog(np.arange(lmax+1),cl[:,2], label='r = %.2f'%r)\n",
    "plt.xlim([2,lmax])\n",
    "plt.xlabel(r'Multipole $\\ell$', fontsize=20)\n",
    "plt.ylabel(r'$\\ell(\\ell+1)C_\\ell/2\\pi$', fontsize=20)\n",
    "plt.legend( loc='lower right', frameon=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part Two: Simplified fake likelihood for some TT data. \n",
    "We'll start with a very simple diagonal likelihood that is easy to understand. We will replace this with the real ACT likelihood later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Fake likelihood based on realistic high-ell noise for CMB S4 data \n",
    "\n",
    "def s4_tt_likelihood(modell,loaddata=True,path=path_to_cmb_notebook_data):\n",
    "    if loaddata: \n",
    "        # if it is the first time, load the data\n",
    "        data = np.loadtxt('./binned_errors.dat', unpack=True)\n",
    "    modeltt = np.zeros(len(data[0]))\n",
    "    inds = (data[0]-0.5)\n",
    "    inds = inds.astype(int)\n",
    "    for i,ind in enumerate(inds):\n",
    "        modeltt[i] = modell[ind]  \n",
    "    loglike = (data[1]-modeltt)**2/(2.*data[2]**2)\n",
    "    loglike=-np.sum(loglike,axis=0)\n",
    "    return loglike\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing we notice is that we shouldn't just be taking the model spectrum at that bin, but we should be <font color='orange'> binning </font> the theory.\n",
    "\n",
    "### Binning the data\n",
    "We'll start by using a top hat binning function from some lower bound to an upper bound, with evenly spaced bins.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in the ell values for the data\n",
    "low_bin = 500\n",
    "high_bin = 2500\n",
    "nbin = 20\n",
    "\n",
    "# First we need to make sure that the theory power spectra are in Cls and not Dl = l^2Cl/2pi\n",
    "dl_tt = CMBOutscale*totCL[:,0]\n",
    "cl_tt = 2*np.pi*CMBOutscale*totCL[:,0]/ls**2\n",
    "\n",
    "# now making a vector of binned ell and a placeholder vector of cls\n",
    "ellbin = np.linspace(low_bin,high_bin,nbin)\n",
    "ellmids = ellbin[0:-1]+ (ellbin[1:]-ellbin[0:-1])/2\n",
    "clttbin = 0*ellmids\n",
    "\n",
    "for c,ell in enumerate(ellbin[:-1]):\n",
    "    inds = np.where((ls > ell)& (ls <= ellbin[c+1]))[0]\n",
    "    clttbin[c] = np.mean(cl_tt[inds])\n",
    "\n",
    "# Now transform the Cl values back to Dl\n",
    "dlttbin = ellmids**2*clttbin/(2*np.pi)\n",
    "plt.plot(ls,dl_tt, label='unbinned')\n",
    "plt.plot(ellmids,dlttbin, label='binned')\n",
    "plt.xlabel(r'Multipole $\\ell$', fontsize=20)\n",
    "plt.ylabel(r'$\\ell(\\ell+1)C_\\ell/2\\pi$', fontsize=20)\n",
    "plt.axis([500,2500,0,2000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'> EXERCISE: </font> Play around with the above binning by changing the range and number of bins. Can you think about how you would change the binning window function from this simple flat (top-hat) function?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing binning inside the likelihood\n",
    "Often this binning happens inside the likelihood itself. So let's modify the simple likelihood above to bin the theory\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Fake likelihood based on realistic high-ell noise for CMB S4 data \n",
    "\n",
    "def s4_tt_likelihood_binned(modell,loaddata=True,path=path_to_cmb_notebook_data):\n",
    "    if loaddata: \n",
    "        # if it is the first time, load the data\n",
    "        data = np.loadtxt('./binned_errors.dat', unpack=True)\n",
    "    \n",
    "    modeltt = np.zeros(len(data[0]))\n",
    "    model_ell = np.arange(len(modell))\n",
    "    \n",
    "    # Making our binning vector\n",
    "    midpt_ells = data[0]\n",
    "    delta_ells = midpt_ells[2]-midpt_ells[1]\n",
    "    \n",
    "    ell_bins = midpt_ells - delta_ells/2.\n",
    "    ell_bins[0] = 10 # In this case we don't want to sum ells below 10\n",
    "    \n",
    "    ell_bins=np.append(ell_bins, ell_bins[-1]+ delta_ells/2. )\n",
    "\n",
    "    # converting to Cl before averaging\n",
    "    cltt = 2*np.pi*modell/(model_ell**2)\n",
    "    \n",
    "    for i,ell in enumerate(ell_bins[0:-1]):\n",
    "        inds = np.where((model_ell > ell)& \n",
    "                        (model_ell <= ell_bins[i+1]))[0]\n",
    "        \n",
    "        modeltt[i] = np.mean(cltt[inds])\n",
    "    # converting back to Dl to do the chi^2 computation\n",
    "    modeltt = modeltt*data[0]**2/(2*np.pi)\n",
    "    loglike = (data[1]-modeltt)**2/(2.*data[2]**2)\n",
    "    loglike=-np.sum(loglike,axis=0)\n",
    "    return loglike\n",
    "   \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's compare\n",
    "loglike = s4_tt_likelihood(dl_tt)\n",
    "loglike_binned = s4_tt_likelihood_binned(dl_tt)\n",
    "print(loglike, loglike_binned)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the difference between the two results run on the same data vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part Three: Real data likelihoods\n",
    "\n",
    "Let's now look at the real ACT likelihood. Here are we reading in the ACT-only best-fit spectrum as theory, and we also need to define the `yp2` needed in the ACTlikelihood to 1.001, and for comparison's sake, we load a new spectrum -- that which maximizes the likelihood at the new value of `\n",
    "yp2`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "like = pyactlike.ACTPowerSpectrumData()\n",
    "# Read in the act only spectra\n",
    "filename = like.data_dir + \"/bf_ACTPol_lcdm.minimum.theory_cl\"\n",
    "\n",
    "tt_lmax = 6000\n",
    "ell, dell_tt, dell_te, dell_ee = np.genfromtxt(\n",
    "    filename,\n",
    "    delimiter=None,\n",
    "    unpack=True,\n",
    "    max_rows=tt_lmax - 1,\n",
    "    usecols=(0, 1, 2, 3),\n",
    ")\n",
    "\n",
    "# Set the new yp2 value\n",
    "yp2_act_only = 1.001\n",
    "\n",
    "# Now let's look at the chi2 using pyactlike\n",
    "like = pyactlike.ACTPowerSpectrumData()\n",
    "chi2 = -2 * like.loglike(dell_tt, dell_te, dell_ee, yp2_act_only)\n",
    "print(\"ACT only chi2 = \" + \"{0:.12f}\".format(chi2))\n",
    "print(\"Expected:       279.005057627002\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's take a moment to compare our theoretical predictions to the data. In order to do se we want to identify the TT, TE and EE bins, data and errors. We will look at the 'wide' patch data. The particular spectrum (TT, EE, etc.) and coadded patch information is set by the indices applied to the `like` instance variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wmin=130 #first element of TT data vector for wide patch\n",
    "wmax=130+40 #last element of TT for wide patch\n",
    "\n",
    "# Read in the ell values for the data\n",
    "TT_bval=like.bval[wmin:wmax]\n",
    "\n",
    "# Compute the Dl values using Dl = cl*ell*(ell+1)/(2*pi) \n",
    "TT_dat=like.X_data[wmin:wmax]*(like.bval[wmin:wmax]+1)*like.bval[wmin:wmax]/(2.*np.pi)\n",
    "\n",
    "# Read in the error which also needs to be converted to an error in Dl\n",
    "TT_err=like.X_sig[wmin:wmax]*(like.bval[wmin:wmax]+1)*like.bval[wmin:wmax]/(2.*np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we can plot this data along with the theory from above \n",
    "plt.figure(figsize=(15,6))\n",
    "\n",
    "plt.plot(ell, dell_tt, \"-\", label=\"ACT theory\")\n",
    "plt.errorbar(TT_bval,TT_dat,yerr = TT_err, fmt='bo',label=\"ACT data\")\n",
    "plt.xlabel(r\"Multipole, $\\ell$\")\n",
    "plt.ylabel(r\"$C_{\\ell} \\ell (\\ell+1) / 2 \\pi$ $[\\mu$K$^2]$\")\n",
    "plt.legend()\n",
    "plt.xlim(0,4300)\n",
    "plt.ylim(0,6000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By eye this looks pretty good, our data points seem to fit our theory extremely well.  We can test this more explicitly by looking at the residuals and printing out a simplified diagonal $\\chi^2$.  It's worth noting that in order to get the full $\\chi^2$ one would need to use the full likelihood (with `like.loglike`) instead of this simplified version.\n",
    "\n",
    "Now identify the bandpower window functions to bin the theory, plot residuals, and also print simplified diagonal $\\chi^2$ (again, for fully-accurate $\\chi^2$, use the full likelihood!). First, we obtain a binned theory spectrum, to compare with the binned data stored inside `like`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We begin by creating an array of ell values the range from 2 to 5000\n",
    "l_list = np.array(range(2, like.tt_lmax + 1))\n",
    "\n",
    "# Our data above was binned and so in order to compare to theory we \n",
    "# will need to similarly bin the theory curve from above\n",
    "\n",
    "# Start by reading in the Dl values annd converting to cls\n",
    "cltt = np.zeros(like.lmax_win)\n",
    "cltt[1 : like.tt_lmax] = (dell_tt[: like.tt_lmax - 1] / l_list / (l_list + 1.0) * 2.0 * np.pi)\n",
    "\n",
    "# now we bin these cl values just as we did for the data above\n",
    "bmax, lmax_win = like.bmax, like.lmax_win\n",
    "cth_tt = like.win_func_w[2 * bmax : 3 * bmax, 1:lmax_win] @ cltt[1:lmax_win] \n",
    "\n",
    "# We convert the binned cl values back to binned Dl by multiplying through by l(l+1)/2pi\n",
    "dth_tt = cth_tt[like.b0:like.b0+like.nbintt]*(like.bval[wmin:wmax]+1)*like.bval[wmin:wmax]/(2.*np.pi) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we have binned data and binned theory and so can plot the residuals\n",
    "plt.figure(figsize=(15,6))\n",
    "plt.errorbar(TT_bval,TT_dat-dth_tt,yerr = TT_err, fmt='bo',label=\"ACT-wide-ACT-theory)\")\n",
    "plt.plot(TT_bval, (TT_dat-dth_tt)*0.)\n",
    "plt.xlabel(r\"Multipole, $\\ell$\")\n",
    "plt.ylabel(r\"$(C_{\\ell}^{TT}-C_{\\ell}^{theory}) \\ell (\\ell+1) / 2 \\pi$ $[\\mu$K$^2]$\")\n",
    "plt.legend()\n",
    "plt.xlim(0,4300)\n",
    "plt.ylim(-100,70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, from the residuals and errors we loaded earlier, we can evaluate our \"diagonal\" $\\chi^2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can also find the chisq of this binned data and print it out\n",
    "chi_act_TT = (TT_dat-dth_tt)**2/TT_err**2\n",
    "print(\"ACT TT-wide diagonal chi2 = \" + \"{0:.2f}\".format(np.sum(chi_act_TT))+ \" for 40 data points\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Repeating for TE, EE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wmin_TE, wmax_TE= 130+40, 130+40+45   #set the bounds for TE\n",
    "wmin_EE, wmax_EE= 130+85, 130+85+45   #set the bounds for EE\n",
    "\n",
    "# As before we read in the ell, Dl, and error values\n",
    "# For TE\n",
    "TE_bval=like.bval[wmin_TE:wmax_TE]\n",
    "TE_dat=like.X_data[wmin_TE:wmax_TE]*like.bval[wmin_TE:wmax_TE]*(like.bval[wmin_TE:wmax_TE]+1)/(2.*np.pi)\n",
    "TE_err=like.X_sig[wmin_TE:wmax_TE]*like.bval[wmin_TE:wmax_TE]*(like.bval[wmin_TE:wmax_TE]+1)/(2.*np.pi)\n",
    "\n",
    "# For EE\n",
    "EE_bval=like.bval[wmin_EE:wmax_EE]\n",
    "EE_dat=like.X_data[wmin_EE:wmax_EE]*like.bval[wmin_EE:wmax_EE]*(like.bval[wmin_EE:wmax_EE]+1)/(2.*np.pi)\n",
    "EE_err=like.X_sig[wmin_EE:wmax_EE]*like.bval[wmin_EE:wmax_EE]*(like.bval[wmin_EE:wmax_EE]+1)/(2.*np.pi)\n",
    "\n",
    "# Start by binning the EE thoery\n",
    "clee = np.zeros(like.lmax_win)\n",
    "clee[1 : like.tt_lmax] = (dell_ee[: like.tt_lmax - 1] / l_list / (l_list + 1.0) * 2.0 * np.pi)\n",
    "bmax, lmax_win = like.bmax, like.lmax_win\n",
    "cth_ee = like.win_func_w[9 * bmax : 10 * bmax, 1:lmax_win] @ clee[1:lmax_win] \n",
    "dth_ee = cth_ee[:like.nbinee]*like.bval[wmin_EE:wmax_EE]*(like.bval[wmin_EE:wmax_EE]*+1)/(2.*np.pi) \n",
    "\n",
    "# Next let's bin the TE theory\n",
    "clte = np.zeros(like.lmax_win)\n",
    "clte[1 : like.tt_lmax] = (dell_te[: like.tt_lmax - 1] / l_list / (l_list + 1.0) * 2.0 * np.pi)\n",
    "bmax, lmax_win = like.bmax, like.lmax_win\n",
    "cth_te = like.win_func_w[6 * bmax : 7 * bmax, 1:lmax_win] @ clte[1:lmax_win] \n",
    "dth_te = cth_te[:like.nbinte]*like.bval[wmin_TE:wmax_TE]*(like.bval[wmin_TE:wmax_TE]+1)/(2.*np.pi) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the residuals\n",
    "plt.figure(figsize=(18,6))\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.errorbar(TE_bval,TE_dat-dth_te,yerr = TE_err, fmt='ro',label=\"ACT-wide-ACT-theory)\")\n",
    "plt.plot(TE_bval, (TE_dat-dth_te)*0.)\n",
    "plt.yticks(fontsize=15)\n",
    "plt.xlim(0,4300)\n",
    "plt.xlabel(r\"Multipole, $\\ell$\")\n",
    "plt.ylabel(r\"$(C_{\\ell}^{TE}-C_{\\ell}^{theory}) \\ell (\\ell+1) / 2 \\pi$ $[\\mu$K$^2]$\")\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.errorbar(EE_bval,EE_dat-dth_ee,yerr = EE_err, fmt='ro',label=\"ACT-wide-ACT-theory)\")\n",
    "plt.plot(EE_bval, (EE_dat-dth_ee)*0.)\n",
    "plt.yticks(fontsize=15)\n",
    "plt.xlim(0,4300)\n",
    "plt.xlabel(r\"Multipole, $\\ell$\")\n",
    "plt.ylabel(r\"$(C_{\\ell}^{EE}-C_{\\ell}^{theory}) \\ell (\\ell+1) / 2 \\pi$ $[\\mu$K$^2]$\")\n",
    "plt.legend()\n",
    "\n",
    "# Print out the chi^2 values\n",
    "chi_act_TE = (TE_dat-dth_te)**2/TE_err**2\n",
    "print(\"ACT TE-wide diagonal chi2 = \" + \"{0:.2f}\".format(np.sum(chi_act_TE))+ \" for 45 data points\")\n",
    "\n",
    "chi_act_EE = (EE_dat-dth_ee)**2/EE_err**2\n",
    "print(\"ACT EE-wide diagonal chi2 = \" + \"{0:.2f}\".format(np.sum(chi_act_EE))+ \" for 45 data points\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font> Repeat for deep patch (wherever you see data being accessed from the likelihood objects in the code above, change the number \"130\" to a \"0\"). Comment on any similarities/differences.\n",
    "\n",
    "Optional next step: use the tools later in the notebook to see which cosmological parameters are better-constrained by which patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stepping in parameter space\n",
    "We are now going to explore generating different model 'universes' to fit the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first generate an arbitrary set of spectra from CAMB\n",
    "cosmo_params = camb.model.CAMBparams()\n",
    "\n",
    "# set H0 to the ACT+WMAP best-fit result, set the universe curvature to be slightly positive\n",
    "# set lmax for this simulation to match the rest of the notebook \n",
    "cosmo_params.set_cosmology(H0 = 67.9, omk = 0.1)\n",
    "cosmo_params.set_for_lmax(6000)\n",
    "\n",
    "# retrieve the spectra\n",
    "res = camb.get_results(cosmo_params)\n",
    "spectra = res.get_cmb_power_spectra(params = cosmo_params, spectra = ('total',), CMB_unit = 'muK')['total']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot what our new model looks like (visually it is hard to tell the difference to our ACT-only maximum likelihood theory), but against the data the problem is clear:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can plot the new, theory spectra as before\n",
    "ell_new = np.arange(spectra.shape[0])\n",
    "dell_tt_new, dell_ee_new, dell_bb_new, dell_te_new = spectra.T\n",
    "\n",
    "plt.figure(figsize = (14, 5))\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.plot(ell_new, dell_tt_new, \"-\", label=\"TT - new theory\")\n",
    "plt.errorbar(TT_bval,TT_dat,yerr = TT_err, fmt='bo',label=\"TT - ACT data\")\n",
    "plt.ylabel(r\"$C_{\\ell} \\ell (\\ell+1) / 2 \\pi$ $[\\mu$K$^2]$\", fontsize = 14)\n",
    "plt.xlabel(r\"Multipole, $\\ell$\", fontsize = 14)\n",
    "plt.legend()\n",
    "plt.xlim(0,3000)\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.plot(ell_new, dell_ee_new, \"-\", label=\"EE - new theory\")\n",
    "plt.errorbar(EE_bval,EE_dat,yerr = EE_err, fmt='bo', label=\"EE - ACT data\")\n",
    "plt.xlabel(r\"Multipole, $\\ell$\", fontsize = 14)\n",
    "plt.legend()\n",
    "plt.xlim(0,3000)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramsvec=np.array([67.5,0.022, 0.122, 0, 0.06, 0.965])\n",
    "pars = camb.CAMBparams()\n",
    "pars.set_cosmology(H0=paramsvec[0], ombh2=paramsvec[1], omch2=paramsvec[2], mnu=0.06, omk=paramsvec[3], tau=paramsvec[4])\n",
    "pars.InitPower.set_params(ns=paramsvec[5], r=0)\n",
    "pars.set_for_lmax(4400, lens_potential_accuracy=0);\n",
    "\n",
    "#calculate results for these parameters\n",
    "results = camb.get_results(pars)\n",
    "\n",
    "#get dictionary of CAMB power spectra\n",
    "powers =results.get_cmb_power_spectra(pars)\n",
    "totCL=powers['total']\n",
    "model = totCL\n",
    "cltt = totCL[:,0]*CMBOutscale\n",
    "loglike = s4_tt_likelihood(cltt)\n",
    "print(loglike)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing the $\\chi^2$ for the ACT likelihood confirms this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at the chi2 using pyactlike\n",
    "like = pyactlike.ACTPowerSpectrumData()\n",
    "chi2 = -2 * like.loglike(dell_tt_new, dell_te_new, dell_ee_new, yp2_act_only)\n",
    "print(\"ACT only chi2 = \" + \"{0:.12f}\".format(chi2))\n",
    "print(\"Expected:       279.004901885481\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stepping in parameter space\n",
    "We are now going to call CAMB with a param vector in the same way as above, and compute the log likelihood for this. We want to take a step in this 6-D parameter space specified by the step vector. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramsvec=np.array([67.5,0.022, 0.122, 0, 0.06, 0.965])\n",
    "# Using this code above, we can take a gaussian step specified by the step vector below\n",
    "stepvec = np.array([0.1,0.0001, 0.0001, 0, 0.005, 0.001])\n",
    "nsteps = 2\n",
    "loglike = np.zeros(nsteps)\n",
    "for i in range(nsteps):\n",
    "    if i==0:\n",
    "        # First step\n",
    "        step = paramsvec\n",
    "    else:\n",
    "        # Take a Gaussian step from the previous position\n",
    "        step = step+np.random.randn(len(paramsvec))*stepvec\n",
    "    # Initialise the CMAB params    \n",
    "    pars = camb.CAMBparams()\n",
    "    pars.set_cosmology(H0=step[0], ombh2=step[1], omch2=step[2], mnu=0.06, omk=step[3], tau=step[4])\n",
    "    pars.InitPower.set_params(ns=step[5], r=0)\n",
    "    pars.set_for_lmax(4400, lens_potential_accuracy=0)\n",
    "    # Compute the spectra\n",
    "    powers =results.get_cmb_power_spectra(pars)\n",
    "    totCL=powers['total']\n",
    "    model = totCL\n",
    "    cltt = totCL[:,0]*CMBOutscale\n",
    "    loglike[i] = s4_tt_likelihood(cltt)\n",
    "print('loglike vector =', loglike)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part Four: Preliminaries of MCMC\n",
    "We are now ready to do the MCMC. We'll define the simplest/ugliest version of the Metropolis Hastings algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mcmc_mh(ratln):\n",
    "    accept=False\n",
    "    r1 = np.random.rand()\n",
    "    # If the step is definitely better, we want to accept it.\n",
    "    # If it isn't necessarily better, we want to throw a random number and step if we exceed it\n",
    "    if np.exp(ratln) > r1:\n",
    "        accept=True\n",
    "    return accept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using this code above, we can take a gaussian step specified by the step vector below\n",
    "stepvec = np.array([0.1,0.0001, 0.0001, 0, 0.005, 0.001])\n",
    "# stepvec = np.array([0.1,0, 0, 0, 0, 0])\n",
    "\n",
    "steps = 10\n",
    "loglike = np.zeros(steps)\n",
    "stepskeep = np.zeros((steps,len(paramsvec)+1))\n",
    "for i in range(steps):\n",
    "    \n",
    "    if i==0:\n",
    "        step = paramsvec\n",
    "        accept=True\n",
    "        pars = camb.CAMBparams()\n",
    "        pars.set_cosmology(H0=step[0], ombh2=step[1], omch2=step[2], mnu=0.06, omk=step[3], tau=step[4])\n",
    "        pars.InitPower.set_params(ns=step[5], r=0)\n",
    "        pars.set_for_lmax(4400, lens_potential_accuracy=0)\n",
    "        powers =results.get_cmb_power_spectra(pars)\n",
    "        totCL=powers['total']\n",
    "        model = totCL\n",
    "        cltt = totCL[:,0]*CMBOutscale\n",
    "        loglike[i] = s4_tt_likelihood(cltt)\n",
    "        #print loglike[i]\n",
    "        stepskeep[i,0:len(paramsvec)] = step\n",
    "        stepskeep[i,len(paramsvec)]= loglike[i]\n",
    "    else:\n",
    "        \n",
    "        step = stepskeep[i-1,0:len(paramsvec)]+np.random.randn(len(paramsvec))*stepvec \n",
    "        #print step\n",
    "        pars = camb.CAMBparams()\n",
    "        # Put the param vector into the camb structure\n",
    "        pars.set_cosmology(H0=step[0], ombh2=step[1], omch2=step[2], mnu=0.06, omk=step[3], tau=step[4])\n",
    "        pars.InitPower.set_params(ns=step[5], r=0)\n",
    "        pars.set_for_lmax(4400, lens_potential_accuracy=0)\n",
    "        # compute the power spectrum\n",
    "        powers =results.get_cmb_power_spectra(pars)\n",
    "        totCL=powers['total']\n",
    "        model = totCL\n",
    "        cltt = totCL[:,0]*CMBOutscale\n",
    "        # compute the likelihood\n",
    "        loglike[i] = s4_tt_likelihood(cltt)\n",
    "        rat = loglike[i]-loglike[i-1]\n",
    "        accept = mcmc_mh(rat)\n",
    "        \n",
    "        if accept:   \n",
    "            stepskeep[i,0:len(paramsvec)] = step\n",
    "            stepskeep[i,len(paramsvec)] = loglike[i]\n",
    "        else:\n",
    "            stepskeep[i,0:len(paramsvec)] = stepskeep[i-1,0:len(paramsvec)]\n",
    "            loglike[i] = loglike[i-1]\n",
    "            stepskeep[i,len(paramsvec)] = loglike[i]\n",
    "        \n",
    "    \n",
    "            \n",
    "np.savetxt('chain.txt', stepskeep, delimiter=' ', fmt='%.3e')\n",
    "print('we are done')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We don't actually want to read in the data every time. \n",
    "<font color='red'>EXERCISE: </font> Change the likelihood function to only read in the data the first time it is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your notes here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### This code is really ugly (and slow)!\n",
    "<font color='red'>EXERCISE: </font>  Write functions/modules to speed up the MCMC code above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part Five: Analyzing MCMC chains\n",
    "\n",
    "We will use MCMC chains generated using the ACT likelihood to look at parameters. The Likelihood we've introduced above can be used to solve for the best fit cosmological parameters that fit the ACT data.  This process and our results are discussed more in [Choi et al (2020)](https://phy-act1.princeton.edu/public/saiola/act_dr4_C20.pdf)  and  [Aiola et al (2020)](https://phyact1.princeton.edu/public/saiola/act_dr4_A20.pdf) but here we will show you how to use one of the MCMC chains to look at the parameters.\n",
    "We begin by reading in a chain and a file with the corresponding column names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Begin by reading in a chain and names of the columns\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "# Read in the column names\n",
    "names = pd.read_csv(path_to_act_notebook_data + \"ACTPol_lcdm.paramnames.csv\", names = [\"parameter\", \"latex parameter name\"])\n",
    "chain = np.genfromtxt(path_to_act_notebook_data + \"ACTPol_lcdm_1.txt\")\n",
    "names[\"parameter\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Making 1D parameter plots and 2D contour plots\n",
    "Now that we've read in the chain we can look at the data more closely.  It's often useful to look at the 1D and 2D parameter plots, e.g. looking at just H0 or looking at H0 with Omega matter.\n",
    "\n",
    "In this section we will demonstrate two ways that users can look at both of these plots.  The first is a straightforward example using matplotlib to make 1D and 2D plots for this chain.  The other option is to use [GetDist](https://getdist.readthedocs.io/en/latest/), a public software developed by Antony Lewis ([Lewis 2019](https://arxiv.org/abs/1910.13970)) specifically for rendering CMB MCMC chains.\n",
    "\n",
    "We will start with the straightforward matplotlib example.  Here we will use H0 and $\\Omega_m$ but you could use any of the parameters from the list above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We begin with the 1D plot of H0\n",
    "\n",
    "# It's worth noting that since the first two columns of the chains are iteration and likelihood we add +2\n",
    "# to the column number shown above when calling the columns in our scatter plot (i.e. we want H0 which is the \n",
    "# 7th parameter above so we plot chain[:, 9]) to get H0\n",
    "\n",
    "# we use a histogram to look at the data here\n",
    "plt.figure(figsize = (7.5,5))\n",
    "plt.hist(chain[:,9], bins = 100, density = True)\n",
    "plt.xlabel(\"$H_0$ [km/s/MPc]\", fontsize = 16)\n",
    "plt.ylabel(\"Probability Density\", fontsize = 16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's easy to generalize to multi-dimensional projections of the likelihood, e.g. by examining the joint distribution of $H_0$ and $\\Omega_m$ now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now let's try a 2D example adding in Omega Matter\n",
    "plt.figure(figsize = (7.5, 5))\n",
    "plt.scatter(chain[:,9], chain[:,11])\n",
    "plt.ylabel(\"$\\Omega_m$\", fontsize = 16)\n",
    "plt.xlabel(\"$H_0$ [km/s/MPc]\", fontsize = 16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now make the same plots using getdist.  We need to start by reading in the files again using the getdist approach.  From there we will walk through how to make the 2 plots above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from getdist import loadMCSamples, plots\n",
    "\n",
    "# Read in the samples\n",
    "samples = loadMCSamples(path_to_act_notebook_data + 'ACTPol_lcdm', path_to_act_notebook_data + \"ACTPol_lcdm_1.txt\")\n",
    "\n",
    "# Let's again start with the 1D version\n",
    "g = plots.get_single_plotter()\n",
    "g.plot_1d(samples, 'H0', marker_color='k')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see the above plot is similar to our histogram but is much prettier!  We can also use getdist for 2D plots by using the `plot_2d` function from getdist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can now create a 2D contour plot using the samples we loaded above\n",
    "g = plots.get_single_plotter()\n",
    "g.plot_2d(samples, ['H0', 'omegam'], filled=True, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Look at correlation between other parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Compare these 1D and 2D with the equivalent for the chain you ran before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part Six: (Exercise) Improving the MCMC alogoritm \n",
    "We now want to check the acceptance/rejection ratio of the chains. In general we want it to be between 0.2-0.4. To change this, you change the size of the steps in each parameter direction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Modify your code above to compute the acceptance/rejection ratio while the steps are being taken."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, it helps to be at a) a good part in parameter space when you start and b) to not step using a diagonal step matrix, but to step using the correlation between parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Modify your stepping function to take a covariance matrix (determined from a shorter run of the chain) and to step using this covariance matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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
}