{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spectrum simulation with Gammapy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook explains how to use the functions and classes in [gammapy.spectrum](https://docs.gammapy.org/dev/spectrum/index.html) in order to simulate and fit spectra.\n",
    "\n",
    "First, we will simulate and fit a pure power law without any background. Than we will add a power law shaped background component. Finally, we will see how to simulate and fit a user defined model. For all scenarios a toy detector will be simulated. For an example using real CTA IRFs, checkout [this notebook](https://github.com/gammapy/gammapy/blob/master/tutorials/spectrum_simulation_cta.ipynb).\n",
    "\n",
    "The following clases will be used:\n",
    "\n",
    "* [gammapy.irf.EffectiveAreaTable](https://docs.gammapy.org/dev/api/gammapy.irf.EffectiveAreaTable.html)\n",
    "* [gammapy.irf.EnergyDispersion](https://docs.gammapy.org/dev/api/gammapy.irf.EnergyDispersion)\n",
    "* [gammapy.spectrum.SpectrumObservation](https://docs.gammapy.org/dev/api/gammapy.spectrum.SpectrumObservation.html)\n",
    "* [gammapy.spectrum.SpectrumSimulation](https://docs.gammapy.org/dev/api/gammapy.spectrum.SpectrumSimulation.html)\n",
    "* [gammapy.spectrum.SpectrumFit](https://docs.gammapy.org/dev/api/gammapy.spectrum.SpectrumFit.html)\n",
    "* [gammapy.spectrum.models.PowerLaw](https://docs.gammapy.org/dev/api/gammapy.spectrum.models.PowerLaw.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Same procedure as in every script ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import astropy.units as u\n",
    "from gammapy.irf import EnergyDispersion, EffectiveAreaTable\n",
    "from gammapy.spectrum import SpectrumSimulation, SpectrumFit\n",
    "from gammapy.spectrum.models import PowerLaw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create detector\n",
    "\n",
    "For the sake of self consistency of this tutorial, we will simulate a simple detector. For a real application you would want to replace this part of the code with loading the IRFs or your detector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e_true = np.logspace(-2, 2.5, 109) * u.TeV\n",
    "e_reco = np.logspace(-2, 2, 79) * u.TeV\n",
    "\n",
    "edisp = EnergyDispersion.from_gauss(\n",
    "    e_true=e_true, e_reco=e_reco, sigma=0.2, bias=0\n",
    ")\n",
    "aeff = EffectiveAreaTable.from_parametrization(energy=e_true)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 6))\n",
    "edisp.plot_matrix(ax=axes[0])\n",
    "aeff.plot(ax=axes[1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Power law\n",
    "\n",
    "In this section we will simulate one observation using a power law model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwl = PowerLaw(\n",
    "    index=2.3, amplitude=1e-11 * u.Unit(\"cm-2 s-1 TeV-1\"), reference=1 * u.TeV\n",
    ")\n",
    "print(pwl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "livetime = 2 * u.h\n",
    "sim = SpectrumSimulation(\n",
    "    aeff=aeff, edisp=edisp, source_model=pwl, livetime=livetime\n",
    ")\n",
    "sim.simulate_obs(seed=2309, obs_id=1)\n",
    "print(sim.obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = SpectrumFit(obs_list=sim.obs, model=pwl.copy(), stat=\"cash\")\n",
    "fit.fit_range = [1, 10] * u.TeV\n",
    "fit.run()\n",
    "print(fit.result[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Include background\n",
    "\n",
    "In this section we will include a background component. Furthermore, we will also simulate more than one observation and fit each one individuallt in order to get average fit results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bkg_model = PowerLaw(\n",
    "    index=2.5, amplitude=1e-11 * u.Unit(\"cm-2 s-1 TeV-1\"), reference=1 * u.TeV\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "n_obs = 30\n",
    "seeds = np.arange(n_obs)\n",
    "\n",
    "sim = SpectrumSimulation(\n",
    "    aeff=aeff,\n",
    "    edisp=edisp,\n",
    "    source_model=pwl,\n",
    "    livetime=livetime,\n",
    "    background_model=bkg_model,\n",
    "    alpha=0.2,\n",
    ")\n",
    "\n",
    "sim.run(seeds)\n",
    "print(sim.result)\n",
    "print(sim.result[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before moving on to the fit let's have a look at the simulated observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_on = [obs.total_stats.n_on for obs in sim.result]\n",
    "n_off = [obs.total_stats.n_off for obs in sim.result]\n",
    "excess = [obs.total_stats.excess for obs in sim.result]\n",
    "\n",
    "fix, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "axes[0].hist(n_on)\n",
    "axes[0].set_xlabel(\"n_on\")\n",
    "axes[1].hist(n_off)\n",
    "axes[1].set_xlabel(\"n_off\")\n",
    "axes[2].hist(excess)\n",
    "axes[2].set_xlabel(\"excess\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "results = []\n",
    "for obs in sim.result:\n",
    "    fit = SpectrumFit(obs, pwl.copy(), stat=\"wstat\")\n",
    "    fit.optimize()\n",
    "    results.append(\n",
    "        {\n",
    "            \"index\": fit.result[0].model.parameters[\"index\"].value,\n",
    "            \"amplitude\": fit.result[0].model.parameters[\"amplitude\"].value,\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = np.array([_[\"index\"] for _ in results])\n",
    "plt.hist(index, bins=10)\n",
    "print(\"spectral index: {:.2f} +/- {:.2f}\".format(index.mean(), index.std()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "* Fit a pure power law and the user define model to the observation you just simulated. You can start with the user defined model described in the [spectrum_models.ipynb](https://github.com/gammapy/gammapy/blob/master/tutorials/spectrum_models.ipynb) notebook.\n",
    "* Vary the observation lifetime and see when you can distinguish the two models (Hint: You get the final likelihood of a fit from `fit.result[0].statval`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's next\n",
    "\n",
    "In this tutorial we learnd how to simulate and fit data using a toy detector. Go to [gammapy.spectrum](https://docs.gammapy.org/dev/spectrum/index.html) to see what else you can do with gammapy."
   ]
  }
 ],
 "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.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}