{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting 2D images with Sherpa\n",
    "\n",
    "### Introduction\n",
    "\n",
    "Sherpa is the X-ray satellite Chandra modeling and fitting application. It enables the user to construct complex models from simple definitions and fit those models to data, using a variety of statistics and optimization methods. \n",
    "The issues of constraining the source position and morphology are common in X- and Gamma-ray astronomy. \n",
    "This notebook will show you how to apply Sherpa to CTA data.\n",
    "\n",
    "Here we will set up Sherpa to fit the counts map and loading the ancillary images for subsequent use. A relevant test statistic for data with Poisson fluctuations is the one proposed by Cash (1979). The simplex (or Nelder-Mead) fitting algorithm is a good compromise between efficiency and robustness. The source fit is best performed in pixel coordinates.\n",
    "\n",
    "This tutorial has 2 important parts\n",
    "1. Generating the Maps\n",
    "2. The actual fitting with sherpa.\n",
    "\n",
    "Since sherpa deals only with 2-dim images, the first part of this tutorial shows how to prepare gammapy maps to make classical images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Necessary imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from astropy.io import fits\n",
    "import astropy.units as u\n",
    "from astropy.wcs import WCS\n",
    "from astropy.coordinates import SkyCoord\n",
    "from gammapy.data import DataStore\n",
    "from gammapy.irf import EnergyDispersion, make_mean_psf\n",
    "from gammapy.maps import WcsGeom, MapAxis, Map\n",
    "from gammapy.cube import MapMaker, PSFKernel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate the required Maps\n",
    "\n",
    "We first generate the required maps using 3 simulated runs on the Galactic center, exactly as in the [analysis_3d](analysis_3d.ipynb) tutorial.\n",
    "\n",
    "It is always advisable to make the maps on fine energy bins, and then sum them over to get an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define which data to use\n",
    "data_store = DataStore.from_dir(\"$GAMMAPY_DATA/cta-1dc/index/gps/\")\n",
    "obs_ids = [110380, 111140, 111159]\n",
    "observations = data_store.get_observations(obs_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_axis = MapAxis.from_edges(\n",
    "    np.logspace(-1, 1.0, 10), unit=\"TeV\", name=\"energy\", interp=\"log\"\n",
    ")\n",
    "geom = WcsGeom.create(\n",
    "    skydir=(0, 0),\n",
    "    binsz=0.02,\n",
    "    width=(10, 8),\n",
    "    coordsys=\"GAL\",\n",
    "    proj=\"CAR\",\n",
    "    axes=[energy_axis],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "maker = MapMaker(geom, offset_max=4.0 * u.deg)\n",
    "maps = maker.run(observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Making a PSF Map\n",
    "\n",
    "Make a PSF map and weigh it with the exposure at the source position to get a 2D PSF "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mean PSF\n",
    "src_pos = SkyCoord(0, 0, unit=\"deg\", frame=\"galactic\")\n",
    "table_psf = make_mean_psf(observations, src_pos)\n",
    "\n",
    "# PSF kernel used for the model convolution\n",
    "psf_kernel = PSFKernel.from_table_psf(table_psf, geom, max_radius=\"0.3 deg\")\n",
    "\n",
    "# get the exposure at the source position\n",
    "exposure_at_pos = maps[\"exposure\"].get_by_coord(\n",
    "    {\n",
    "        \"lon\": src_pos.l.value,\n",
    "        \"lat\": src_pos.b.value,\n",
    "        \"energy\": energy_axis.center,\n",
    "    }\n",
    ")\n",
    "\n",
    "# now compute the 2D PSF\n",
    "psf2D = psf_kernel.make_image(exposures=exposure_at_pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make 2D images from 3D ones\n",
    "\n",
    "Since sherpa image fitting works only with 2-dim images,\n",
    "we convert the generated maps to 2D images using `run_images()` and save them as fits files. The exposure is weighed with the spectrum before averaging (assumed to be a power law by default).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maps = maker.run_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Path(\"analysis_3d\").mkdir(exist_ok=True)\n",
    "\n",
    "maps[\"counts\"].write(\"analysis_3d/counts_2D.fits\", overwrite=True)\n",
    "maps[\"background\"].write(\"analysis_3d/background_2D.fits\", overwrite=True)\n",
    "maps[\"exposure\"].write(\"analysis_3d/exposure_2D.fits\", overwrite=True)\n",
    "fits.writeto(\"analysis_3d/psf_2D.fits\", psf2D.data, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read the maps and store them in a sherpa model\n",
    "\n",
    "We now have the prepared files which sherpa can read. \n",
    "This part of the notebook shows how to do image analysis using sherpa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sherpa.astro.ui as sh\n",
    "\n",
    "sh.set_stat(\"cash\")\n",
    "sh.set_method(\"simplex\")\n",
    "\n",
    "sh.load_image(\"analysis_3d/counts_2D.fits\")\n",
    "sh.set_coord(\"logical\")\n",
    "\n",
    "sh.load_table_model(\"expo\", \"analysis_3d/exposure_2D.fits\")\n",
    "sh.load_table_model(\"bkg\", \"analysis_3d/background_2D.fits\")\n",
    "sh.load_psf(\"psf\", \"analysis_3d/psf_2D.fits\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In principle one might first want to fit the background amplitude. However the background estimation method already yields the correct normalization, so we freeze the background amplitude to unity instead of adjusting it. The (smoothed) residuals from this background model are then computed and shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sh.set_full_model(bkg)\n",
    "bkg.ampl = 1\n",
    "sh.freeze(bkg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resid = Map.read(\"analysis_3d/counts_2D.fits\")\n",
    "resid.data = sh.get_data_image().y - sh.get_model_image().y\n",
    "resid_smooth = resid.smooth(width=4)\n",
    "resid_smooth.plot(add_cbar=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Find and fit the brightest source\n",
    "We then find the position of the maximum in the (smoothed) residuals map, and fit a (symmetrical) Gaussian source with that initial position:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yp, xp = np.unravel_index(\n",
    "    np.nanargmax(resid_smooth.data), resid_smooth.data.shape\n",
    ")\n",
    "ampl = resid_smooth.get_by_pix((xp, yp))[0]\n",
    "\n",
    "# creates g0 as a gauss2d instance\n",
    "sh.set_full_model(bkg + psf(sh.gauss2d.g0) * expo)\n",
    "g0.xpos, g0.ypos = xp, yp\n",
    "sh.freeze(g0.xpos, g0.ypos)  # fix the position in the initial fitting step\n",
    "\n",
    "# fix exposure amplitude so that typical exposure is of order unity\n",
    "expo.ampl = 1e-9\n",
    "sh.freeze(expo)\n",
    "sh.thaw(g0.fwhm, g0.ampl)  # in case frozen in a previous iteration\n",
    "\n",
    "g0.fwhm = 10  # give some reasonable initial values\n",
    "g0.ampl = ampl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "sh.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit all parameters of this Gaussian component, fix them and re-compute the residuals map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sh.thaw(g0.xpos, g0.ypos)\n",
    "sh.fit()\n",
    "sh.freeze(g0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resid.data = sh.get_data_image().y - sh.get_model_image().y\n",
    "resid_smooth = resid.smooth(width=3)\n",
    "resid_smooth.plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iteratively find and fit additional sources\n",
    "Instantiate additional Gaussian components, and use them to iteratively fit sources, repeating the steps performed above for component g0. (The residuals map is shown after each additional source included in the model.) This takes some time..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize components with fixed, zero amplitude\n",
    "for i in range(1, 10):\n",
    "    model = sh.create_model_component(\"gauss2d\", \"g\" + str(i))\n",
    "    model.ampl = 0\n",
    "    sh.freeze(model)\n",
    "\n",
    "gs = [g0, g1, g2]\n",
    "sh.set_full_model(bkg + psf(g0 + g1 + g2) * expo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "for i in range(1, len(gs)):\n",
    "    yp, xp = np.unravel_index(\n",
    "        np.nanargmax(resid_smooth.data), resid_smooth.data.shape\n",
    "    )\n",
    "    ampl = resid_smooth.get_by_pix((xp, yp))[0]\n",
    "    gs[i].xpos, gs[i].ypos = xp, yp\n",
    "    gs[i].fwhm = 10\n",
    "    gs[i].ampl = ampl\n",
    "\n",
    "    sh.thaw(gs[i].fwhm)\n",
    "    sh.thaw(gs[i].ampl)\n",
    "    sh.fit()\n",
    "\n",
    "    sh.thaw(gs[i].xpos)\n",
    "    sh.thaw(gs[i].ypos)\n",
    "    sh.fit()\n",
    "    sh.freeze(gs[i])\n",
    "\n",
    "    resid.data = sh.get_data_image().y - sh.get_model_image().y\n",
    "    resid_smooth = resid.smooth(width=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resid_smooth.plot(add_cbar=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating output table and Test Statistics estimation\n",
    "When adding a new source, one needs to check the significance of this new source. A frequently used method is the Test Statistics (TS). This is done by comparing the change of statistics when the source is included compared to the null hypothesis (no source ; in practice here we fix the amplitude to zero).\n",
    "\n",
    "$TS = Cstat(source) - Cstat(no source)$\n",
    "\n",
    "The criterion for a significant source detection is typically that it should improve the test statistic by at least 25 or 30. We have added only 3 sources to save time, but you should keep doing this till del(stat) is less than the required number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.stats import gaussian_fwhm_to_sigma\n",
    "from astropy.table import Table\n",
    "\n",
    "rows = []\n",
    "for g in gs:\n",
    "    ampl = g.ampl.val\n",
    "    g.ampl = 0\n",
    "    stati = sh.get_stat_info()[0].statval\n",
    "    g.ampl = ampl\n",
    "    statf = sh.get_stat_info()[0].statval\n",
    "    delstat = stati - statf\n",
    "\n",
    "    geom = resid.geom\n",
    "    # sherpa uses 1 based indexing\n",
    "    coord = geom.pix_to_coord((g.xpos.val - 1, g.ypos.val - 1))\n",
    "    pix_scale = geom.pixel_scales.mean().deg\n",
    "    sigma = g.fwhm.val * pix_scale * gaussian_fwhm_to_sigma\n",
    "    rows.append(\n",
    "        dict(delstat=delstat, glon=coord[0], glat=coord[1], sigma=sigma)\n",
    "    )\n",
    "\n",
    "table = Table(rows=rows, names=rows[0])\n",
    "for name in table.colnames:\n",
    "    table[name].format = \".5g\"\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "1. Keep adding sources till there are no more significat ones in the field. How many Gaussians do you need?\n",
    "2. Use other morphologies for the sources (eg: disk, shell) rather than only Gaussian.\n",
    "3. Compare the TS between different models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### More about sherpa\n",
    "\n",
    "These are good resources to learn more about Sherpa:\n",
    "\n",
    "* https://python4astronomers.github.io/fitting/fitting.html\n",
    "* https://github.com/DougBurke/sherpa-standalone-notebooks\n",
    "\n",
    "You could read over the examples there, and try to apply a similar analysis to this dataset here to practice.\n",
    "\n",
    "If you want a deeper understanding of how Sherpa works, then these proceedings are good introductions:\n",
    "\n",
    "* http://conference.scipy.org/proceedings/scipy2009/paper_8/full_text.pdf\n",
    "* http://conference.scipy.org/proceedings/scipy2011/pdfs/brefsdal.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}