{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pulsar analysis with Gammapy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to do a pulsar analysis with Gammapy. It's based on a Vela simulation file from the CTA DC1, which already contains a column of phases. We will produce a phasogram, a phase-resolved map and a phase-resolved spectrum of the Vela pulsar using the class PhaseBackgroundEstimator from gammapy.background.phase. \n",
    "\n",
    "The phasing in itself is not done here, and it requires specific packages like Tempo2 or PINT (https://nanograv-pint.readthedocs.io/en/latest/readme.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Opening the data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first do the imports and load the only observation containing Vela in the CTA 1DC dataset shipped with Gammapy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from regions import CircleSkyRegion\n",
    "from astropy.coordinates import SkyCoord, Angle\n",
    "import astropy.units as u\n",
    "\n",
    "from gammapy.maps import Map, WcsGeom\n",
    "from gammapy.cube import fill_map_counts\n",
    "from gammapy.data import DataStore\n",
    "from gammapy.background import PhaseBackgroundEstimator\n",
    "from gammapy.spectrum.models import PowerLaw\n",
    "from gammapy.utils.energy import EnergyBounds\n",
    "from gammapy.spectrum import (\n",
    "    SpectrumExtraction,\n",
    "    SpectrumFit,\n",
    "    FluxPointEstimator,\n",
    "    SpectrumResult,\n",
    "    SpectrumEnergyGroupMaker,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the data store (which is a subset of CTA-DC1 data):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_store = DataStore.from_dir(\"$GAMMAPY_DATA/cta-1dc/index/gps\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define obsevation ID and print events:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_obs_vela = [111630]\n",
    "obs_list_vela = data_store.get_observations(id_obs_vela)\n",
    "print(obs_list_vela[0].events)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our observation, let's select the events in 0.2° radius around the pulsar position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_target = SkyCoord(ra=128.836 * u.deg, dec=-45.176 * u.deg, frame=\"icrs\")\n",
    "on_radius = 0.2 * u.deg\n",
    "\n",
    "# Apply angular selection\n",
    "events_vela = obs_list_vela[0].events.select_sky_cone(\n",
    "    center=pos_target, radius=on_radius\n",
    ")\n",
    "print(events_vela)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the phases of the selected events in a dedicated array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phases = events_vela.table[\"PHASE\"]\n",
    "\n",
    "# Let's take a look at the first 10 phases\n",
    "phases[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Phasogram"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have the phases, we can make a phasogram. A phasogram is a histogram of phases and it works exactly like any other histogram (you can set the binning, evaluate the errors based on the counts in each bin, etc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbins = 30\n",
    "phase_min, phase_max = (0, 1)\n",
    "values, bin_edges = np.histogram(\n",
    "    phases, range=(phase_min, phase_max), bins=nbins\n",
    ")\n",
    "bin_width = (phase_max - phase_min) / nbins\n",
    "\n",
    "bin_center = (bin_edges[:-1] + bin_edges[1:]) / 2\n",
    "\n",
    "\n",
    "# Poissonian uncertainty on each bin\n",
    "values_err = np.sqrt(values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.bar(\n",
    "    x=bin_center,\n",
    "    height=values,\n",
    "    width=bin_width,\n",
    "    color=\"#d53d12\",\n",
    "    alpha=0.8,\n",
    "    edgecolor=\"black\",\n",
    "    yerr=values_err,\n",
    ")\n",
    "plt.xlim(0, 1)\n",
    "plt.xlabel(\"Phase\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.title(\"Phaseogram with angular cut of {}\".format(on_radius));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's add some fancy additions to our phasogram: a patch on the ON- and OFF-phase regions and one for the background level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate background level\n",
    "off_phase_range = (0.7, 1.0)\n",
    "on_phase_range = (0.5, 0.6)\n",
    "\n",
    "mask_off = (off_phase_range[0] < phases) & (phases < off_phase_range[1])\n",
    "\n",
    "count_bkg = mask_off.sum()\n",
    "print(\"Number of Off events: {}\".format(count_bkg))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bkg level normalized by the size of the OFF zone (0.3)\n",
    "bkg = count_bkg / nbins / (off_phase_range[1] - off_phase_range[0])\n",
    "\n",
    "# error on the background estimation\n",
    "bkg_err = (\n",
    "    np.sqrt(count_bkg) / nbins / (off_phase_range[1] - off_phase_range[0])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's redo the same plot for the basis\n",
    "plt.bar(\n",
    "    x=bin_center,\n",
    "    height=values,\n",
    "    width=bin_width,\n",
    "    color=\"#d53d12\",\n",
    "    alpha=0.8,\n",
    "    edgecolor=\"black\",\n",
    "    yerr=values_err,\n",
    ")\n",
    "\n",
    "# Plot background level\n",
    "x_bkg = np.linspace(0, 1, 50)\n",
    "\n",
    "kwargs = {\"color\": \"black\", \"alpha\": 0.5, \"ls\": \"--\", \"lw\": 2}\n",
    "\n",
    "plt.plot(x_bkg, (bkg - bkg_err) * np.ones_like(x_bkg), **kwargs)\n",
    "plt.plot(x_bkg, (bkg + bkg_err) * np.ones_like(x_bkg), **kwargs)\n",
    "\n",
    "plt.fill_between(\n",
    "    x_bkg, bkg - bkg_err, bkg + bkg_err, facecolor=\"grey\", alpha=0.5\n",
    ")  # grey area for the background level\n",
    "\n",
    "# Let's make patches for the on and off phase zones\n",
    "on_patch = plt.axvspan(\n",
    "    on_phase_range[0], on_phase_range[1], alpha=0.3, color=\"gray\", ec=\"black\"\n",
    ")\n",
    "\n",
    "off_patch = plt.axvspan(\n",
    "    off_phase_range[0],\n",
    "    off_phase_range[1],\n",
    "    alpha=0.4,\n",
    "    color=\"white\",\n",
    "    hatch=\"x\",\n",
    "    ec=\"black\",\n",
    ")\n",
    "\n",
    "# Legends \"ON\" and \"OFF\"\n",
    "plt.text(0.55, 5, \"ON\", color=\"black\", fontsize=17, ha=\"center\")\n",
    "plt.text(0.895, 5, \"OFF\", color=\"black\", fontsize=17, ha=\"center\")\n",
    "plt.xlabel(\"Phase\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.xlim(0, 1)\n",
    "plt.title(\"Phasogram with angular cut of {}\".format(on_radius));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Phase-resolved map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the phases are computed, we want to do a phase-resolved sky map : a map of the ON-phase events minus alpha times the OFF-phase events. Alpha is the ratio between the size of the ON-phase zone (here 0.1) and the OFF-phase zone (0.3).\n",
    "It's a map of the excess events in phase, which are the pulsed events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geom = WcsGeom.create(binsz=0.02 * u.deg, skydir=pos_target, width=\"5 deg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Let's create an ON-map and an OFF-map:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "on_map = Map.from_geom(geom)\n",
    "off_map = Map.from_geom(geom)\n",
    "\n",
    "events_vela_on = events_vela.select_parameter(\"PHASE\", on_phase_range)\n",
    "events_vela_off = events_vela.select_parameter(\"PHASE\", off_phase_range)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fill_map_counts(on_map, events_vela_on)\n",
    "fill_map_counts(off_map, events_vela_off)\n",
    "\n",
    "# Defining alpha as the ratio of the ON and OFF phase zones\n",
    "alpha = (on_phase_range[1] - on_phase_range[0]) / (\n",
    "    off_phase_range[1] - off_phase_range[0]\n",
    ")\n",
    "\n",
    "# Create and fill excess map\n",
    "# The pulsed events are the difference between the ON-phase count and alpha times the OFF-phase count\n",
    "excess_map = on_map - off_map * alpha\n",
    "\n",
    "# Plot excess map\n",
    "excess_map.smooth(kernel=\"gauss\", width=0.2 * u.deg).plot(add_cbar=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Phase-resolved spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also do a phase-resolved spectrum. In order to do that, there is the class PhaseBackgroundEstimator. In a phase-resolved analysis, the background is estimated in the same sky region but in the OFF-phase zone.\n",
    "\n",
    "We start by estimating the background with the class PhaseBackgroundEstimator. It takes the observations, the ON-region, and an ON- and OFF-phase zones (the same we defined for the phasogram and the phase-resolved map). It results in a gammapy.background.phase.PhaseBackgroundEstimator that serves as an input for other spectral analysis classes in Gammapy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining an on-region around the pulsar to pass it to the background estimator\n",
    "on_region = CircleSkyRegion(pos_target, on_radius)\n",
    "\n",
    "# The PhaseBackgroundEstimator uses the OFF-phase in the ON-region to estimate the background\n",
    "bkg_estimator = PhaseBackgroundEstimator(\n",
    "    observations=obs_list_vela,\n",
    "    on_region=on_region,\n",
    "    on_phase=on_phase_range,\n",
    "    off_phase=off_phase_range,\n",
    ")\n",
    "bkg_estimator.run()\n",
    "bkg_estimate = bkg_estimator.result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The rest of the analysis is the same as for a standard spectral analysis with Gammapy. All the specificity of a phase-resolved analysis is contained in the PhaseBackgroundEstimator, where the background is estimated in the ON-region OFF-phase rather than in an OFF-region.\n",
    "\n",
    "We can now extract a spectrum with the SpectrumExtraction class. It takes the reconstructed and the true energy binning. Both are expected to be a Quantity with unit energy, i.e. an array with an energy unit. EnergyBounds is a dedicated class to do it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "etrue = EnergyBounds.equal_log_spacing(0.005, 10.0, 100, unit=\"TeV\")\n",
    "ereco = EnergyBounds.equal_log_spacing(0.01, 10, 30, unit=\"TeV\")\n",
    "\n",
    "extraction = SpectrumExtraction(\n",
    "    observations=obs_list_vela,\n",
    "    bkg_estimate=bkg_estimate,\n",
    "    containment_correction=True,\n",
    "    e_true=etrue,\n",
    "    e_reco=ereco,\n",
    ")\n",
    "\n",
    "extraction.run()\n",
    "extraction.compute_energy_threshold(\n",
    "    method_lo=\"energy_bias\", bias_percent_lo=20\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's a look at the files we just created with spectrum_observation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extraction.spectrum_observations[0].peek()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll fit a model to the spectrum with SpectrumFit. First we load a power law model with an initial value for the index and the amplitude and then wo do a likelihood fit. The fit results are printed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = PowerLaw(\n",
    "    index=4, amplitude=\"1.3e-9 cm-2 s-1 TeV-1\", reference=\"0.02 TeV\"\n",
    ")\n",
    "\n",
    "fit_range = (0.04 * u.TeV, 0.4 * u.TeV)\n",
    "ebounds = EnergyBounds.equal_log_spacing(0.04, 0.4, 7, u.TeV)\n",
    "\n",
    "joint_fit = SpectrumFit(\n",
    "    obs_list=extraction.spectrum_observations, model=model, fit_range=fit_range\n",
    ")\n",
    "joint_fit.run()\n",
    "joint_result = joint_fit.result\n",
    "\n",
    "print(joint_result[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you might want to do the stacking here even if in our case there is only one observation which makes it superfluous.\n",
    "We can compute flux points by fitting the norm of the global model in energy bands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked_obs = extraction.spectrum_observations.stack()\n",
    "seg = SpectrumEnergyGroupMaker(obs=stacked_obs)\n",
    "\n",
    "seg.compute_groups_fixed(ebounds=ebounds)\n",
    "fpe = FluxPointEstimator(\n",
    "    obs=stacked_obs, groups=seg.groups, model=joint_result[0].model\n",
    ")\n",
    "flux_points = fpe.run()\n",
    "\n",
    "amplitude_ref = 0.57 * 19.4e-14 * u.Unit(\"1 / (cm2 s MeV)\")\n",
    "spec_model_true = PowerLaw(\n",
    "    index=4.5, amplitude=amplitude_ref, reference=\"20 GeV\"\n",
    ")\n",
    "\n",
    "spectrum_result = SpectrumResult(\n",
    "    points=flux_points, model=joint_result[0].model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot.\n",
    "We present here two different spectra: one for the spectral flux and one for the spectral energy density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First plot for the spectral flux\n",
    "ax0, ax1 = spectrum_result.plot(\n",
    "    energy_range=joint_fit.result[0].fit_range,\n",
    "    fig_kwargs=dict(figsize=(8, 8)),\n",
    "    point_kwargs=dict(label=\"Flux points\"),\n",
    "    fit_kwargs=dict(label=\"Gammapy fit\"),\n",
    ")\n",
    "\n",
    "ax0.set_ylim([1e-14, 1e-7])\n",
    "ax0.set_xlim([4e-2, 5e-1])\n",
    "ax1.set_ylim([-1.7, 1.7])\n",
    "\n",
    "spec_model_true.plot(\n",
    "    ax=ax0,\n",
    "    energy_range=joint_fit.result[0].fit_range,\n",
    "    label=\"Reference model\",\n",
    "    c=\"black\",\n",
    "    linestyle=\"dashed\",\n",
    ")\n",
    "\n",
    "ax0.legend(loc=\"best\")\n",
    "ax0.set_ylabel(r\"Flux [cm$^{-2}$ s$^1$ TeV$^{-1}$]\", size=14)\n",
    "ax1.set_ylabel(\"Residuals\", size=14)\n",
    "ax1.set_xlabel(\"Energy [TeV]\", size=14)\n",
    "ax1.set_xticks([5e-2, 1e-1, 3e-1])\n",
    "ax1.set_xticklabels([5e-2, 1e-1, 3e-1]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Second plot for the spectral energy flux\n",
    "\n",
    "ax0, ax1 = spectrum_result.plot(\n",
    "    energy_range=joint_fit.result[0].fit_range,\n",
    "    energy_power=2,\n",
    "    flux_unit=\"erg-1 cm-2 s-1\",\n",
    "    fig_kwargs=dict(figsize=(8, 8)),\n",
    "    point_kwargs=dict(label=\"Flux points\"),\n",
    "    fit_kwargs=dict(label=\"Gammapy fit\"),\n",
    ")\n",
    "\n",
    "spec_model_true.plot(\n",
    "    ax=ax0,\n",
    "    energy_range=[4e-2, 5e-1] * u.TeV,\n",
    "    energy_power=2,\n",
    "    flux_unit=\"erg-1 cm-2 s-1\",\n",
    "    label=\"Input model\",\n",
    "    c=\"black\",\n",
    "    linestyle=\"dashed\",\n",
    ")\n",
    "\n",
    "ax0.set_ylim([5e-15, 5e-11])\n",
    "ax0.set_xlim([4e-2, 5e-1])\n",
    "ax1.set_ylim([-1.7, 1.7])\n",
    "ax0.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial suffers a bit from the lack of statistics: there were 9 Vela observations in the CTA DC1 while there is only one here. When done on the 9 observations, the spectral analysis is much better agreement between the input model and the gammapy fit."
   ]
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}