2_x-ray-analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div  >\n",
    "<img src=\"https://raw.githubusercontent.com/threeML/astromodels/master/docs/media/transp_logo.png\" alt=\"drawing\" width=\"300\" align=\"right\"/>\n",
    " \n",
    "\n",
    "\n",
    "<div  >\n",
    "<img src=\"https://raw.githubusercontent.com/threeML/threeML/master/logo/logo_sq.png\" alt=\"drawing\" width=\"300\" align=\"right\"/>\n",
    "\n",
    "\n",
    "\n",
    "# X-ray Analysis with 3ML\n",
    "    \n",
    "While 3ML can handle a lot of different data/likelihood types a lot of attention was spent on making sure that users familiar with past community standards are able to easily adapt to the 3ML workflow. There are some guides for these users in the [documentation](https://threeml.readthedocs.io/en/stable/xspec_users.html).\n",
    "    \n",
    "X-ray analysis in 3ML is centered around the `OGIPLike` plugin which reads OGIP style PHAI/II, RMF, and ARF files. the OGIPLike plugin is a specialized version of the `DispersionSpectrumLike` plugin which deals with count data that are produced by convolving the model spectrum with the resonse of an instrument that suffers from energy dispersion. Thus, if you have an instrument you are designing and you don't like fits files... inherit from DispersionSpectrumLike and create your own unique plugin for ROOT, HDF5, txt, etc. files. The cool thing is that you can still fit your data along with normal OGIP type data... or any of the other plugins in the 3ML family. 3ML is a toolbox to bring instruments (and people) together. \n",
    "    \n",
    "    \n",
    "Let's explore the OGIPLike plugin\n",
    "\n",
    " \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## The OGIPLike plugin\n",
    "\n",
    "The OGIP plugin reads in standard OGIP files. **It will complain a lot if files are in the correct format!**. For PHAII files with multiple spectra, you can use the familiar `<filename>{<spectrum_number>}` format to specify file names or you can pass a spectrum number as an argument. \n",
    "\n",
    "<img src=\"https://cdn.pixabay.com/photo/2012/11/28/11/16/star-67705_960_720.jpg\" alt=\"drawing\" width=\"400\" align=\"center\"/>\n",
    "\n",
    "\n",
    "In the tutorial, there are some simulated Chandra data. Let's say that these data come from the observation of a a white drawf atmosphere. Let's see what we can do with these data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.064403Z",
     "start_time": "2021-08-24T15:06:45.830016Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from threeML import *\n",
    "# get xspec models\n",
    "#from astromodels.xspec import *\n",
    "\n",
    "# this is the basic interface to \n",
    "# logging. see the documentation for more\n",
    "update_logging_level(\"INFO\")\n",
    "silence_warnings()\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.units as u\n",
    "\n",
    "%matplotlib notebook\n",
    "from jupyterthemes import jtplot\n",
    "set_threeML_style()\n",
    "jtplot.style(context='notebook', fscale=1, ticks=True, grid=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.335727Z",
     "start_time": "2021-08-24T15:06:50.066685Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra = OGIPLike(name=\"chandra\", \n",
    "                   observation=\"c_data/obs.pha\", \n",
    "                   background=\"c_data/obs_bak.pha\",\n",
    "                   response=\"c_data/acis.rmf\",\n",
    "                   arf_file=\"c_data/acis.arf\",\n",
    "                   spectrum_number=1 )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Note that 3ML probed the type of data that were read in. As long as the data files have been appropriately labelled, the plugin will **choose the correct likelihood for you**. While freedom is a great thing, math is not a democracy and thus we follow the rules so that your fits are of the highest scientific rigour. \n",
    "\n",
    "In this case, the total observation and the background observation are Poisson distributed. Thus, the proper likelihood is a Poisson for the total observation conditional on the Poisson likelihood of the background. For now, we will not model the background. Therefore a profile likelihood will be choosen.\n",
    "\n",
    "Let's examine the properties of the plugin.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.344503Z",
     "start_time": "2021-08-24T15:06:50.339040Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.significance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.351389Z",
     "start_time": "2021-08-24T15:06:50.346807Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.significance_per_channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.543251Z",
     "start_time": "2021-08-24T15:06:50.353414Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.display_rsp()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.549440Z",
     "start_time": "2021-08-24T15:06:50.545354Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.exposure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.685851Z",
     "start_time": "2021-08-24T15:06:50.561473Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.view_count_spectrum(scale_background=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now, not all channels are great to use in an analysis. Thus, we can set our selections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.742051Z",
     "start_time": "2021-08-24T15:06:50.688136Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.set_active_measurements?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.758651Z",
     "start_time": "2021-08-24T15:06:50.744320Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.set_active_measurements('0.2-10')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:50.898283Z",
     "start_time": "2021-08-24T15:06:50.760849Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.view_count_spectrum();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "For profile likelihoods to be valid, there must be at least 1 [background count per channel](https://giacomov.github.io/Bias-in-profile-poisson-likelihood/). Let's do that here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:53.456059Z",
     "start_time": "2021-08-24T15:06:50.900582Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.rebin_on_background(1)\n",
    "#chandra.remove_rebinning()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:53.545479Z",
     "start_time": "2021-08-24T15:06:53.458549Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.view_count_spectrum();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Fitting \n",
    "\n",
    "Ok, we are basically ready to do a fit. But we need a model. Let's make two models, one of a black body and the other a power law. We are going to be Bayesians for now, but remember, there is little difference between the interface for the two approaches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### blackbody model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:06:53.554520Z",
     "start_time": "2021-08-24T15:06:53.547651Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bb = Blackbody()\n",
    "\n",
    "\n",
    "#priors\n",
    "bb.K.prior = Log_uniform_prior(lower_bound = 1e-2, upper_bound = 20)\n",
    "bb.kT.prior = Truncated_gaussian(mu= 1, sigma=2, lower_bound=0, upper_bound=10)\n",
    "\n",
    "\n",
    "# you can also try with the XSPEC blackbody!\n",
    "\n",
    "# bb = XS_bbody()\n",
    "# #priors\n",
    "# bb.norm.prior = Log_uniform_prior(lower_bound = 1e-2, upper_bound = 20)\n",
    "# bb.kt.prior = Truncated_gaussian(mu= 1, sigma=2, lower_bound=0, upper_bound=10)\n",
    "\n",
    "# source\n",
    "ps_bb = PointSource(\"white_drawf_bb\", 0, 0, spectral_shape=bb)\n",
    "\n",
    "# model\n",
    "model_bb = Model(ps_bb)\n",
    "\n",
    "model_bb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bayes_bb = BayesianAnalysis(model_bb, DataList(chandra))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Let's do some prior predicitive checks to make sure our priors are acceptable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from twopc import compute_priorpc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "\n",
    "ppc = compute_priorpc(bayes_bb,500, \"priorpc.h5\", overwrite=True, return_ppc=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig = ppc.chandra.plot();\n",
    "ax = fig.get_axes()[0]\n",
    "\n",
    "\n",
    "ax.set_ylim(1e-10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ppc = compute_priorpc(bayes_bb,500, \"priorpc.h5\",\n",
    "                      parameters=[\"white_drawf_bb.spectrum.main.Blackbody.K\"],\n",
    "                      overwrite=True, return_ppc=True)\n",
    "\n",
    "\n",
    "fig = ppc.chandra.plot();\n",
    "ax = fig.get_axes()[0]\n",
    "\n",
    "\n",
    "ax.set_ylim(1e-10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:09.765583Z",
     "start_time": "2021-08-24T15:06:53.556530Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# let's use ultranest this time\n",
    "bayes_bb.set_sampler(\"ultranest\")\n",
    "\n",
    "bayes_bb.sampler.setup(min_num_live_points=400, use_mlfriends=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:09.765583Z",
     "start_time": "2021-08-24T15:06:53.556530Z"
    },
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "_ = bayes_bb.sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:09.887216Z",
     "start_time": "2021-08-24T15:07:09.767665Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bayes_bb.results.corner_plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:10.481689Z",
     "start_time": "2021-08-24T15:07:09.888903Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "display_spectrum_model_counts(bayes_bb,\n",
    "                              min_rate=50,\n",
    "                              show_background=True);\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "we can have a look at the configuration options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:10.489674Z",
     "start_time": "2021-08-24T15:07:10.483585Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "show_configuration('plugins')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:10.496149Z",
     "start_time": "2021-08-24T15:07:10.491719Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "threeML_config.plugins.ogip.fit_plot.data_color = 'limegreen'\n",
    "threeML_config.plugins.ogip.fit_plot.model_color = '#FF5AFD'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:10.947124Z",
     "start_time": "2021-08-24T15:07:10.498568Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "display_spectrum_model_counts(bayes_bb,\n",
    "                              min_rate=10,\n",
    "                              step=True,\n",
    "                              show_background=True);\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### power law model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:10.956433Z",
     "start_time": "2021-08-24T15:07:10.949261Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plaw = Powerlaw()\n",
    "\n",
    "plaw.K.prior = Log_uniform_prior(lower_bound = 1e-2, upper_bound = 10) \n",
    "\n",
    "plaw.index.prior = Gaussian(mu=-1, sigma=2)\n",
    "plaw.index.bounds = (None, None)\n",
    "\n",
    "\n",
    "# source\n",
    "ps_pl = PointSource(\"white_drawf_pl\", 0, 0, spectral_shape=plaw)\n",
    "\n",
    "# model\n",
    "model_pl = Model(ps_pl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:24.934267Z",
     "start_time": "2021-08-24T15:07:10.958928Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bayes_pl = BayesianAnalysis(model_pl, DataList(chandra))\n",
    "\n",
    "# let's use ultranest this time\n",
    "bayes_pl.set_sampler(\"ultranest\")\n",
    "\n",
    "bayes_pl.sampler.setup(min_num_live_points=500)\n",
    "\n",
    "_ = bayes_pl.sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:25.343217Z",
     "start_time": "2021-08-24T15:07:24.936038Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "display_spectrum_model_counts(bayes_pl,\n",
    "                              min_rate=10,\n",
    "                              show_background=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:27.994561Z",
     "start_time": "2021-08-24T15:07:25.345333Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plot_spectra(bayes_bb.results, bayes_pl.results,\n",
    "             flux_unit=\"erg/(cm2 s keV)\",\n",
    "             ene_min=1*u.keV, ene_max=10*u.keV);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-18T14:10:29.722507Z",
     "start_time": "2021-08-18T14:10:29.647068Z"
    },
    "heading_collapsed": true
   },
   "source": [
    "## posterior predictive checks (PPC)\n",
    "\n",
    "Let's use an external package to 3ML (but built with its tools box!) to compute posterior predictive checks. PPCs are model checking tool that integrate over the posterior and likelihood to compute the probability of new data from the observed data. We can compute this via simulating new data from the likelihood for sampled points from our posterior. A more detailed explanation for this can be found [here](https://academic.oup.com/mnras/article/490/1/927/5570608)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:28.076290Z",
     "start_time": "2021-08-24T15:07:27.996902Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from twopc import compute_postpc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:07:57.214592Z",
     "start_time": "2021-08-24T15:07:28.079084Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ppc_bb = compute_postpc(bayes_bb,\n",
    "                  bayes_bb.results,\n",
    "                  n_sims=500,\n",
    "                  file_name=\"ppc_bb.h5\",\n",
    "                  return_ppc=True, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:00.078395Z",
     "start_time": "2021-08-24T15:07:57.230726Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ppc_bb.chandra.plot(bkg_subtract=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:00.203275Z",
     "start_time": "2021-08-24T15:08:00.084189Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ppc_bb.chandra.plot_qq(channel_energies=None);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:30.285034Z",
     "start_time": "2021-08-24T15:08:00.205986Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ppc_pl = compute_postpc(bayes_pl,\n",
    "                  bayes_pl.results,\n",
    "                  n_sims=500,\n",
    "                  file_name=\"ppc_pl.h5\",\n",
    "                  return_ppc=True, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:33.038616Z",
     "start_time": "2021-08-24T15:08:30.289176Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ppc_pl.chandra.plot(bkg_subtract=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:33.166479Z",
     "start_time": "2021-08-24T15:08:33.042824Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ppc_pl.chandra.plot_qq(channel_energies=[1,2]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Background modeling\n",
    "\n",
    "While profile likelihoods are a great tool when we have no good model for the background, if we can model the background, we should. It is possible to introduce models for the background in 3ML by extracting a plugin with the background spectrum only from a normal OGIPLike plugin. \n",
    "\n",
    "\n",
    "<img src=\"https://upload.wikimedia.org/wikipedia/commons/f/f1/Artist_impression_of_a_supermassive_black_hole_at_the_centre_of_a_galaxy.jpg\" alt=\"drawing\" width=\"400\" align=\"center\"/>\n",
    "\n",
    "This time we look at some simulated data where we are trying to find a spectral line on top of a continuum. The source is a blackhole accreation disk. First, let's load the data and have a look\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:33.573805Z",
     "start_time": "2021-08-24T15:08:33.168574Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra = OGIPLike(name=\"chandra\", \n",
    "                   observation=\"c_data/obs_bkg_demo.pha\", \n",
    "                   background=\"c_data/obs_bkg_demo_bak.pha\",\n",
    "                   response=\"c_data/acis.rmf\",\n",
    "                   arf_file=\"c_data/acis.arf\",\n",
    "                   spectrum_number=1 )\n",
    "\n",
    "\n",
    "chandra.set_active_measurements(\"1-10\")\n",
    "\n",
    "chandra.rebin_on_background(1)\n",
    "\n",
    "fig = chandra.view_count_spectrum(scale_background=True);\n",
    "\n",
    "ax =fig.get_axes()[0]\n",
    "\n",
    "ax.set_xlim(1,10)\n",
    "ax.set_xscale('linear')\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now let's build our model and fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:33.586982Z",
     "start_time": "2021-08-24T15:08:33.576266Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "cpl = Powerlaw(piv=5) + Gaussian()\n",
    "\n",
    "#priors\n",
    "cpl.K_1.prior = Log_uniform_prior(lower_bound = 1e-4, upper_bound = 1)\n",
    "#cpl.xc_1.prior = Truncated_gaussian(mu= 5, sigma=3, lower_bound=1, upper_bound=12)\n",
    "cpl.index_1.prior = Truncated_gaussian(mu= -2, sigma=1, lower_bound=-np.inf, upper_bound=0)\n",
    "\n",
    "cpl.F_2.prior = Log_uniform_prior(lower_bound = 1e-3, upper_bound = 1)\n",
    "cpl.mu_2.prior = Truncated_gaussian(mu= 6.4, sigma=2, lower_bound=1, upper_bound=10)\n",
    "cpl.sigma_2.prior = Log_uniform_prior(lower_bound = 1e-3, upper_bound = 1)\n",
    "# source\n",
    "ps_blackhole = PointSource(\"blackhole\", 0, 0, spectral_shape=cpl)\n",
    "\n",
    "# model\n",
    "model_blackhole = Model(ps_blackhole)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:33.596228Z",
     "start_time": "2021-08-24T15:08:33.589168Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bayes = BayesianAnalysis(model_blackhole, DataList(chandra))\n",
    "\n",
    "# let's use ultranest this time\n",
    "bayes.set_sampler(\"multinest\")\n",
    "\n",
    "bayes.sampler.setup(n_live_points=500,verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:47.866090Z",
     "start_time": "2021-08-24T15:08:33.598129Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_ = bayes.sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.250624Z",
     "start_time": "2021-08-24T15:08:47.868015Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig = chandra.display_model(step=False, min_rate=5,\n",
    "                               data_color=\"#C39BD3\", model_color=\"#76D7C4\",\n",
    "                               \n",
    "                               )\n",
    "\n",
    "ax = fig.get_axes()[0]\n",
    "\n",
    "ax.set_xlim(1,10)\n",
    "\n",
    "ax.set_xscale('linear')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Building a background plugin\n",
    "\n",
    "To model the background, We first must extract the background as a second plugin with its own Poisson likelihood.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.337589Z",
     "start_time": "2021-08-24T15:08:48.252601Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.remove_rebinning()\n",
    "\n",
    "bkg_plugin = SpectrumLike.from_background(\"bkg\", chandra)\n",
    "\n",
    "bkg_plugin.view_count_spectrum();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We then create a model for the background and set that model to the background plugin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.366131Z",
     "start_time": "2021-08-24T15:08:48.339506Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bkg_spectrum = Powerlaw(K=1.e0,index=-1.2, piv=5) + Gaussian(F=1e0, mu=1, sigma=.2) \n",
    "\n",
    "bkg_spectrum.K_1.prior = Log_uniform_prior(lower_bound=1e-2, upper_bound=1e2)\n",
    "bkg_spectrum.index_1.prior = Truncated_gaussian(mu = -1.2, sigma=0.5, lower_bound=-np.inf, upper_bound=0)\n",
    "\n",
    "bkg_spectrum.F_2.prior = Log_uniform_prior(lower_bound=1e-2, upper_bound=1e2)\n",
    "bkg_spectrum.mu_2.prior = Truncated_gaussian(mu = 3, sigma=1, lower_bound=0, upper_bound=10)\n",
    "bkg_spectrum.sigma_2.fix = True\n",
    "\n",
    "\n",
    "bkg_src = PointSource(\"bkg\", 0,0, spectral_shape=bkg_spectrum)\n",
    "\n",
    "bkg_model = Model(bkg_src)\n",
    "\n",
    "\n",
    "bkg_plugin.set_model(bkg_model)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now when we create a plugin, we pass the background plugin rather than a file name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.626799Z",
     "start_time": "2021-08-24T15:08:48.368214Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra_bkg = OGIPLike(\"chandra\",\n",
    "                       observation = \"c_data/obs_bkg_demo.pha\",\n",
    "                       response = \"c_data/acis.rmf\", \n",
    "                       arf_file = \"c_data/acis.arf\" ,\n",
    "                       background = bkg_plugin,\n",
    "                       spectrum_number=1\n",
    "                        \n",
    "                      \n",
    "                      )\n",
    "\n",
    "\n",
    "chandra_bkg.set_active_measurements(\"1-10\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.779227Z",
     "start_time": "2021-08-24T15:08:48.629648Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig = chandra_bkg.view_count_spectrum();\n",
    "\n",
    "ax =fig.get_axes()[0]\n",
    "\n",
    "ax.set_xlim(1,10)\n",
    "ax.set_xscale('linear')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The background model's parameters are now nuissance parameters of the plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.790428Z",
     "start_time": "2021-08-24T15:08:48.781631Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra_bkg.nuisance_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Fitting\n",
    "\n",
    "Now let's fit. We will switch to multinest for speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:08:48.808120Z",
     "start_time": "2021-08-24T15:08:48.792960Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "new_model = clone_model(model_blackhole)\n",
    "\n",
    "bayes_bkg = BayesianAnalysis(new_model, DataList(chandra_bkg))\n",
    "\n",
    "bayes_bkg.set_sampler(\"multinest\")\n",
    "\n",
    "bayes_bkg.sampler.setup(n_live_points=500,verbose=False )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:09:48.196681Z",
     "start_time": "2021-08-24T15:08:48.810498Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_ = bayes_bkg.sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "So we recover uncertainties for both our background and the source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:09:48.205226Z",
     "start_time": "2021-08-24T15:09:48.198764Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "threeML_config.bayesian.corner_style.cmap = \"magma\"\n",
    "threeML_config.bayesian.corner_style.extremes=\"none\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:09:50.769387Z",
     "start_time": "2021-08-24T15:09:48.207790Z"
    },
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bayes_bkg.results.corner_plot(color=\"#F6FF74\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:09:51.466736Z",
     "start_time": "2021-08-24T15:09:50.771936Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# over plot the joint background and source fits\n",
    "fig = chandra_bkg.display_model(step=False, min_rate=1,\n",
    "                               data_color=\"#C39BD3\", model_color=\"#76D7C4\",\n",
    "                               \n",
    "                               )\n",
    "\n",
    "_ = bkg_plugin.display_model(\n",
    "    data_color=\"#85C1E9\", model_color=\"#F7DC6F\",\n",
    "    model_subplot=fig.axes,\n",
    "    step=False,\n",
    "    min_rate=1\n",
    ")\n",
    "\n",
    "ax = fig.get_axes()[0]\n",
    "\n",
    "ax.set_xlim(1,10)\n",
    "\n",
    "ax.set_xscale('linear')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Additional tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The OGIP plugin (or any plugin) is not just for fitting, it can be used as a generic interface between models and instruments for building pipelines. \n",
    "* Plugins and models are serializable meaning they can be farmed out to multi-processing\n",
    "* Most plugins can simulate data from their likelihoods (complex instruments still need some work here)\n",
    "\n",
    "Let's try this out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Simulating fake data\n",
    "\n",
    "The OGIP plugins can simulate fake data from thier likelihoods. If you already have a model set and/or fitted, this is a very simple task:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "new_plugin = chandra.get_simulated_dataset(\"sim_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "new_plugin.view_count_spectrum();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The data are simulated with the current model's current parameters. Thus, you can setup a simulation by setting a model and setting its current parameters.\n",
    "Let's try this out by making a new model.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plaw = Powerlaw(K=1e-1, index=-2)\n",
    "\n",
    "spectrum = TbAbs(NH=21, redshift=0) * plaw\n",
    "\n",
    "\n",
    "x_grid = np.linspace(0.1,10,1000)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "\n",
    "ax.semilogy(x_grid, spectrum(x_grid))\n",
    "ax.semilogy(x_grid, plaw(x_grid))\n",
    "\n",
    "ax.set(xlabel=\"keV\", ylabel=\"flux\",\n",
    "       ylim=[1e-5,1e-1]\n",
    "      )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model = Model(PointSource(\"synth\",0,0,spectral_shape=spectrum))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra.set_model(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sim_chandra = chandra.get_simulated_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sim_chandra.view_count_spectrum(scale_background=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra._n_synthetic_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sim_chandra.simulated_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "There is a faster way to create synthetic data by using the `from_function` interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bkg = Powerlaw(K=1e1,index=-1.5, piv=5) + Gaussian(F=1.5e1, mu=3., sigma=.2)\n",
    "\n",
    "new_data = DispersionSpectrumLike.from_function('sim',\n",
    "                                                source_function=spectrum,\n",
    "                                                response= chandra.response,\n",
    "                                                background_function=bkg,\n",
    "                                                scale_factor=1e-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Exploring samples\n",
    "\n",
    "While there are many helper functions to deal with the results of fits, you can easily construct your own pipelines. One of these are the plotting and error propagation via samples. Again, this does not change between MLE and Bayesian analysis. When performing an MLE fit, the samples are drawn from the covariance matrix of the fit.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "result = bayes_bb.results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "samples = result.samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "samples.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "\n",
    "x_grid = np.geomspace(0.1, 10, 50)\n",
    "\n",
    "for s in samples.T[::10]:\n",
    "    \n",
    "    model_bb.set_free_parameters(s)\n",
    "    \n",
    "    y = model_bb.white_drawf_bb(x_grid)\n",
    "    \n",
    "    ax.loglog(x_grid, y, color='r', alpha=0.05, lw=.7)\n",
    "    \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}