3_advanced_topics.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",
    "# Advanced Topics\n",
    "\n",
    "Here we will cover a few more examples beyond just x-ray analsysis to get an idea of what is possible in 3ML.\n",
    "\n",
    "\n",
    " \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Joint Fits\n",
    "\n",
    "We have learned to do x-ray analysis with 3ML, but one of the main purposes of 3ML is to allow joint fits with different types of data. To demostrate this, we will have a look at some simulated optical and x-ray data from a GRB afterglow observered by Chandra and the 7-Band GROND instrument.\n",
    "\n",
    "<img src=\"https://room.eu.com/images/contents/GRBurst.jpg\" alt=\"drawing\" width=\"400\" align=\"left\"/>\n",
    "<img src=\"https://cdn.eso.org/images/screen/eso0730a.jpg\" alt=\"drawing\" width=\"400\" align=\"right\"/>\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:20:59.125038Z",
     "start_time": "2021-08-24T15:20:55.178471Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from threeML import *\n",
    "\n",
    "\n",
    "# get xspec models\n",
    "#from astromodels.xspec import *\n",
    "update_logging_level(\"INFO\")\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",
    "jtplot.style(context='notebook', fscale=1, ticks=True, grid=False)\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "data_path = Path(\"joint_data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Photometric Data\n",
    "\n",
    "Photometric data is handled in 3ML via the `PhotometricLike` plugin. We use the `speclite` package to interface to a variety of filters that come preloaded in 3ML. This is updated all the time, but it is already pretty extensive. If you want to use the filters, you need to load them manually:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T15:21:12.234757Z",
     "start_time": "2021-08-24T15:20:59.127075Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import speclite.filters as spec_filters\n",
    "\n",
    "from threeML.utils.photometry import PhotometericObservation\n",
    "# The filter library takes a while to load so you must import it explicitly.\n",
    "from threeML.utils.photometry import get_photometric_filter_library\n",
    "\n",
    "threeML_filter_library = get_photometric_filter_library()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Let's get the filters for GROND and have a look"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:07:25.965059Z",
     "start_time": "2021-08-24T16:07:25.962692Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond_filters = threeML_filter_library.LaSilla.GROND"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:07:26.571855Z",
     "start_time": "2021-08-24T16:07:26.520954Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "spec_filters.plot_filters(grond_filters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The shapes of the filters are important. Many times, astronomers try to convert magnitude to counts via multiplying by the average filter response. This causes the errors to become inaccurate. 3ML works directly in magnitude space via *synthetic photometry*, i.e., forward folding the photon model into the filter response to provide more accurate fits. \n",
    "\n",
    "\n",
    "There are several ways to load data into the plugin.\n",
    "\n",
    "1) via keywords where data are passed as tuples of magnitudes and uncertainties with the keyword being the band name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:18:54.256527Z",
     "start_time": "2021-08-24T16:18:54.253463Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond_observation = PhotometericObservation.from_kwargs(g=(21.5,0.02),\n",
    "                                                        r=(22.,0.03),\n",
    "                                                        i=(21.8, 0.01),\n",
    "                                                        z=(21.2, 0.01),\n",
    "                                                        J=(19.6, 0.1),\n",
    "                                                        H=(18.6, 0.01),\n",
    "                                                        K=(18.0, 0.1))\n",
    "\n",
    "grond_observation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "2) via an HDF5 file type in 3ML that stores the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:22:59.815080Z",
     "start_time": "2021-08-24T16:22:59.809932Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond_observation = PhotometericObservation.from_hdf5(\"joint_data/grond_data.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can then load these observation into the plugin, or we could use the helper functions to do this directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:44:49.949408Z",
     "start_time": "2021-08-24T16:44:49.913233Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond = PhotometryLike.from_file(\"grond\",grond_filters,file_name=\"joint_data/grond_data.h5\",)\n",
    "grond.plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can see what data are in the plugin:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:03.256615Z",
     "start_time": "2021-08-24T16:45:03.252675Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond.band_J"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "and we can turn the bands on and off for the fit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:05.121074Z",
     "start_time": "2021-08-24T16:45:05.118690Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond.band_J.off = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:05.323910Z",
     "start_time": "2021-08-24T16:45:05.319925Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond.band_J"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:05.843708Z",
     "start_time": "2021-08-24T16:45:05.841257Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond.band_J.on = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:06.231538Z",
     "start_time": "2021-08-24T16:45:06.227594Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "grond.band_J"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### X-ray data\n",
    "\n",
    "Now we can load our x-ray data as we always do:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:07.306798Z",
     "start_time": "2021-08-24T16:45:07.024242Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chandra = OGIPLike(\"chandra\",observation=\"joint_data/obs_demo.pha\", spectrum_number=1)\n",
    "\n",
    "chandra.set_active_measurements(\"1-10\")\n",
    "\n",
    "chandra.rebin_on_background(1)\n",
    "\n",
    "chandra.view_count_spectrum();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Performing a joint fit\n",
    "\n",
    "To do a joint fit, everything is the same as before except that now our DataList is a little bigger\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:45:13.817802Z",
     "start_time": "2021-08-24T16:45:13.815350Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = DataList(grond, chandra)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "This time we will use maximum likelihood for the fun of it. GRB afterglows are typically characterized as broken power laws to approximate their synchrotron radiation. Let's build a model for that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:48:31.961436Z",
     "start_time": "2021-08-24T16:48:31.957747Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bpl = Broken_powerlaw(K=1, piv=1)\n",
    "\n",
    "# set some reasonable bounds on the break energy\n",
    "bpl.xb.bounds = (0, None)\n",
    "\n",
    "ps = PointSource(\"grb\", 0,0, spectral_shape=bpl)\n",
    "\n",
    "\n",
    "model = Model(ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:48:32.400049Z",
     "start_time": "2021-08-24T16:48:32.390981Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "jl = JointLikelihood(model, data)\n",
    "\n",
    "# let's use the ROOT minimizer\n",
    "jl.set_minimizer(\"ROOT\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:48:33.042936Z",
     "start_time": "2021-08-24T16:48:32.647951Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_ = jl.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Because these data live in different spaces, they have different x-axes. Thus, we have to plot the data fits with different functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:50:19.225677Z",
     "start_time": "2021-08-24T16:50:18.850768Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "display_spectrum_model_counts(jl, min_rate=10);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:50:23.121488Z",
     "start_time": "2021-08-24T16:50:22.994911Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "display_photometry_model_magnitudes(jl);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Of course, we can always look at the model space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:56:09.528576Z",
     "start_time": "2021-08-24T16:56:08.222533Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plot_spectra(jl.results, ene_min =25000 * u.Angstrom, ene_max = 1.239842* u.Angstrom, energy_unit=u.Angstrom  , flux_unit=\"erg/s/cm2/keV\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:55:03.453867Z",
     "start_time": "2021-08-24T16:55:03.448977Z"
    },
    "hidden": true
   },
   "source": [
    "But maybe it is a power law?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:58:09.972804Z",
     "start_time": "2021-08-24T16:58:09.969175Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pl = Powerlaw(K=1, piv=1)\n",
    "\n",
    "ps_pl = PointSource(\"grb_pl\", 0,0, spectral_shape=pl)\n",
    "model_pl = Model(ps_pl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:58:26.693516Z",
     "start_time": "2021-08-24T16:58:26.683307Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "jl_pl = JointLikelihood(model_pl, data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:58:44.033067Z",
     "start_time": "2021-08-24T16:58:43.857506Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "jl_pl.set_minimizer(\"ROOT\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:58:44.033067Z",
     "start_time": "2021-08-24T16:58:43.857506Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_ = jl_pl.fit(quiet=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-25T04:15:14.887494Z",
     "start_time": "2021-08-25T04:15:12.001711Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plot_spectra(jl.results, jl_pl.results ,ene_min =25000 * u.Angstrom, ene_max = 1.239842* u.Angstrom, energy_unit=u.Angstrom  , flux_unit=\"erg/s/cm2/keV\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "#### Comparing models with the likelihood ratio test. \n",
    "\n",
    "We have a rather complex likelihood and we really do not want to trust Wilks' theorem to give us the true \"significance\" of the more advanced model. Thus, we need to calibrate our LRT via monte carlo. Once of the best papers on statistical calibration is [this one](https://arxiv.org/abs/1803.08393) by Michael Betancourt. 3ML provides some syntax magic to do this for you.\n",
    "\n",
    "**NOTE:** This operation can be very slow on the binder environment. I recommend only running it locally. The `with parallel_computation` command can start small cluster on your computer and perform the operations in parallel.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T16:59:25.214074Z",
     "start_time": "2021-08-24T16:59:25.210976Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrt = LikelihoodRatioTest(jl_pl,jl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T17:14:26.114815Z",
     "start_time": "2021-08-24T17:05:17.369208Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "with parallel_computation(start_cluster=True):\n",
    "\n",
    "    lrt.by_mc(100, continue_on_failure=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can see the results of the LRT by plotting them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-24T21:09:34.577726Z",
     "start_time": "2021-08-24T21:09:34.437975Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrt.plot_TS_distribution(bins=100, ec=\"k\", fc=\"white\", lw=1.2)\n",
    "_ = plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Complex Models\n",
    "\n",
    "\n",
    "### Generating some synthetic data\n",
    "\n",
    "![alt text](http://aasnova.org/wp-content/uploads/2016/03/fig16.jpg)\n",
    "\n",
    "Let's say we have two galactic x-ray sources, some accreting compact binaries perhaps? We observe them at two different times. These sources (imaginary) sources emit a blackbody which is theorized to always be at the same temperature, but perhaps at different flux levels.\n",
    "\n",
    "\n",
    "Lets simulate one of these sources:\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from threeML.utils.OGIP.response import OGIPResponse\n",
    "from threeML.io.package_data import get_path_of_data_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "np.random.seed(1234)\n",
    "\n",
    "# we will use a demo response\n",
    "response_1 = OGIPResponse(get_path_of_data_file(\"datasets/ogip_powerlaw.rsp\"))\n",
    "\n",
    "\n",
    "source_function_1 = Blackbody(K=5e-8, kT=500.0)\n",
    "background_function_1 = Powerlaw(K=1, index=-1.5, piv=1.0e3)\n",
    "\n",
    "\n",
    "spectrum_generator_1 = DispersionSpectrumLike.from_function(\n",
    "    \"s1\",\n",
    "    source_function=source_function_1,\n",
    "    background_function=background_function_1,\n",
    "    response=response_1,\n",
    ")\n",
    "\n",
    "fig = spectrum_generator_1.view_count_spectrum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now let's simulate the other source, but this one has an extra feature! There is a power law component in addition to the blackbody. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "response_2 = OGIPResponse(get_path_of_data_file(\"datasets/ogip_powerlaw.rsp\"))\n",
    "\n",
    "\n",
    "source_function_2 = Blackbody(K=1e-7, kT=500.0) + Powerlaw_flux(\n",
    "    F=2e2, index=-1.5, a=10, b=500\n",
    ")\n",
    "background_function_2 = Powerlaw(K=1, index=-1.5, piv=1.0e3)\n",
    "\n",
    "\n",
    "spectrum_generator_2 = DispersionSpectrumLike.from_function(\n",
    "    \"s2\",\n",
    "    source_function=source_function_2,\n",
    "    background_function=background_function_2,\n",
    "    response=response_2,\n",
    ")\n",
    "\n",
    "fig = spectrum_generator_2.view_count_spectrum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Make the model\n",
    "\n",
    "Now let's make the model we will use to fit the data. First, let's make the spectral function for source_1 and set priors on the parameters.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "spectrum_1 = Blackbody()\n",
    "\n",
    "spectrum_1.K.prior = Log_normal(mu=np.log(1e-7), sigma=1)\n",
    "spectrum_1.kT.prior = Log_normal(mu=np.log(300), sigma=2)\n",
    "\n",
    "ps1 = PointSource(\"src1\", ra=1, dec=20, spectral_shape=spectrum_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We will do the same for the other source but also include the power law component "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "spectrum_2 = Blackbody() + Powerlaw_flux(\n",
    "    a=10, b=500\n",
    ")  # a,b are the bounds for the flux for this model\n",
    "\n",
    "spectrum_2.K_1.prior = Log_normal(mu=np.log(1e-6), sigma=1)\n",
    "spectrum_2.kT_1.prior = Log_normal(mu=np.log(300), sigma=2)\n",
    "\n",
    "spectrum_2.F_2.prior = Log_normal(mu=np.log(1e2), sigma=1)\n",
    "spectrum_2.F_2.bounds = (None, None)\n",
    "\n",
    "spectrum_2.index_2.prior = Gaussian(mu=-1.0, sigma=1)\n",
    "spectrum_2.index_2.bounds = (None, None)\n",
    "\n",
    "ps2 = PointSource(\"src2\", ra=2, dec=-10, spectral_shape=spectrum_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model = Model(ps1, ps2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Linking parameters\n",
    "\n",
    "We hypothesized that both sources should have the a same blackbody temperature. We can impose this by linking the temperatures. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model.link(\n",
    "    model.src1.spectrum.main.Blackbody.kT, model.src2.spectrum.main.composite.kT_1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "we could also link the parameters with an arbitrary function rather than directly. Check out the [astromodels documentation](https://astromodels.readthedocs.io/en/latest/Model_tutorial.html#linking-parameters) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Assigning sources to plugins\n",
    "\n",
    "Now, if we simply passed out model to the BayesianAnalysis or JointLikelihood objects, it would sum the point source spectra together and apply both sources to all data. \n",
    "\n",
    "This is not what we want. Many plugins have the ability to be assigned directly to a source. Let's do that here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "spectrum_generator_1.assign_to_source(\"src1\")\n",
    "\n",
    "spectrum_generator_2.assign_to_source(\"src2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we simply make our our data list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = DataList(spectrum_generator_1, spectrum_generator_2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Fitting the data\n",
    "\n",
    "Now we fit the data as we normally would. We use Bayesian analysis here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ba = BayesianAnalysis(model, data)\n",
    "ba.set_sampler(\"multinest\")\n",
    "ba.sampler.setup(n_live_points=500)\n",
    "#ba.sampler.setup(frac_remain=0.5)\n",
    "_ = ba.sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Let's examine the fits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig = display_spectrum_model_counts(ba)\n",
    "ax = fig.get_axes()[0]\n",
    "ax.set_ylim(1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Lets grab the result. Remember, we can save the results to disk, so all of the following operations can be run at a later time without having to redo all the above steps!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "result = ba.results\n",
    "fig = result.corner_plot()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Computing fluxes\n",
    "\n",
    "Now we will compute fluxes. We can compute them in many different units, over any energy range also specified in any units. \n",
    "\n",
    "The flux is computed by integrating the function over energy. By default, a fast trapezoid method is used. If you need more accuracy, you can change the method in the configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "silence_warnings()\n",
    "\n",
    "# we could also use 'quad' here, but it is slower\n",
    "\n",
    "threeML_config.point_source.integrate_flux_method = \"trapz\"\n",
    "\n",
    "result.get_flux(ene_min=1 * u.keV, ene_max=1 * u.MeV, flux_unit=\"erg/cm2/s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We see that a pandas dataframe is returned with all the information. We could change the confidence region for the uncertainties if we desire. However, we could also sum the source fluxes! 3ML will take care of propagating the uncertainties (for any of these operations). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "result.get_flux(\n",
    "    ene_min=1 * u.keV,\n",
    "    ene_max=1 * u.MeV,\n",
    "    confidence_level=0.95,\n",
    "    sum_sources=True,\n",
    "    flux_unit=\"erg/cm2/s\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can get the fluxes of individual components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "result.get_flux(\n",
    "    ene_min=10 * u.keV, ene_max=0.5 * u.MeV, use_components=True, flux_unit=\"1/(cm2 s)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "As well as choose which component to compute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "result.get_flux(ene_min=10*u.keV, \n",
    "                ene_max = 0.5*u.MeV,\n",
    "                use_components=True,\n",
    "                components_to_use =[\"Blackbody\"],\n",
    "                \n",
    "                flux_unit=\"erg2/(cm2 s)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Finally, the returned flux object is a pandas table and can be manipulated as such:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "flux = result.get_flux(ene_min=1 * u.keV, ene_max=1 * u.MeV, flux_unit=\"J2/m2/yr\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "flux[\"flux\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "flux[\"flux\"][\"src1: total\"]\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
}