CMB_School_Part_08.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jeffmcm1977/CMBAnalysis_SummerSchool/blob/master/CMB_School_Part_08.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cosmological Analysis\n",
    "\n",
    "###  Renée Hložek and Jeff McMahon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to do a very rough example of an MCMC, using the <a href=\" https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm\"> Metroplis Hastings algorithm. </a>\n"
   ]
  },
  {
   "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 camb\" || python -m pip install camb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import CAMB\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import camb\n",
    "from camb import model, initialpower"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will start by 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, CMB_unit='muK')"
   ]
  },
  {
   "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",
    "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,totCL[:,0], color='k', label='lensed')\n",
    "ax[0,0].plot(ls,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_title(r'$TT$')\n",
    "ax[1,0].plot(ls,totCL[:,1], color='k', label='lensed')\n",
    "ax[1,0].plot(ls,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,1].plot(ls,totCL[:,3], color='k', label='lensed')\n",
    "ax[1,1].plot(ls,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",
    "\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.legend( loc='lower right', frameon=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now let's make a fake likelihood for CMB S4 TT data. NOTE that this is super simplified, in general we have a much more complicated likelihood, and any other spectra will be correlated."
   ]
  },
  {
   "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):\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",
    "    # plt.figure()\n",
    "    # plt.plot(data[0],modeltt)\n",
    "    # plt.errorbar(data[0], data[1], data[2])\n",
    "    \n",
    "    loglike = (data[1]-modeltt)**2/(2.*data[2]**2)\n",
    "    loglike = -np.sum(loglike,axis=0)\n",
    "    return loglike    "
   ]
  },
  {
   "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 binning the theory.\n",
    "<font color='red'>EXERCISE: </font>  Write a function to bin the theory over the same ell range as the binned data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's call the module for the spectrum we have above.\n",
    "model = totCL\n",
    "cltt = totCL[:,0]\n",
    "loglike = s4_tt_likelihood(cltt)\n",
    "print(loglike)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We are now going to call CAMB with a param vector in the same way as above, and compute the log likelihood for this."
   ]
  },
  {
   "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, CMB_unit='muK')\n",
    "totCL = powers['total']\n",
    "model = totCL\n",
    "cltt = totCL[:,0]\n",
    "loglike = s4_tt_likelihood(cltt)\n",
    "print(loglike)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 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": [
    "# 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, CMB_unit='muK')\n",
    "    totCL = powers['total']\n",
    "    model = totCL\n",
    "    cltt = totCL[:,0]\n",
    "    loglike[i] = s4_tt_likelihood(cltt)\n",
    "print('loglike vector =', loglike)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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",
    "\n",
    "steps = 100\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, CMB_unit='muK')\n",
    "        totCL = powers['total']\n",
    "        model = totCL\n",
    "        cltt = totCL[:,0]\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",
    "        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, CMB_unit='muK')\n",
    "        totCL = powers['total']\n",
    "        model = totCL\n",
    "        cltt = totCL[:,0]\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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your notes here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in the chain\n",
    "chain = np.loadtxt('chain.txt', unpack=True)\n",
    "print(np.shape(chain))\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.plot(chain[0,:], chain[1,:], marker='*', linestyle='None')\n",
    "plt.xlabel(r'$H_0$', fontsize=20)\n",
    "plt.ylabel(r'$\\Omega_bh^2$', fontsize=20)\n",
    "plt.show()"
   ]
  },
  {
   "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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "## Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your discussion here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "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
}