{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe7880aa",
   "metadata": {},
   "source": [
    "# Rectification\n",
    "\n",
    "This notebook demonstrates how to rectify a detector readout from a spectroscopic simulation using Scopesim. \"Rectification\" means the transformation of a spectral trace from the detector onto a rectangular pixel grid of wavelength and spatial position along the slit. Wavelength calibration and rectification are major tasks of the spectroscopic data-reduction pipeline. For convenience, Scopesim includes functionality to perform these tasks by reversing the *known* mapping that was used for the simulation, resulting in easily analysable 2D spectra that include all the noise and background components but neglect the uncertainties of a wavelength calibration as it would be performed during the reduction of real data. \n",
    "Rectification is demonstrated on a METIS long-slit simulation, but the procedure applies to MICADO spectroscopic simulations as well (but more expensive to simulate and rectify). METIS IFU simulations have to be treated differently. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7e93cd-e320-431e-80e6-8686abf88c75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from astropy import units as u\n",
    "from astropy.io import fits\n",
    "from astropy.wcs import WCS\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from synphot import SourceSpectrum, Empirical1D\n",
    "from scopesim_templates.micado import flatlamp\n",
    "import scopesim as sim\n",
    "\n",
    "sim.bug_report()\n",
    "\n",
    "# Edit this path if you have a custom install directory, otherwise comment it out.\n",
    "sim.link_irdb(\"../../../../\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b2a57b",
   "metadata": {},
   "source": [
    "If you haven’t got the instrument packages yet, uncomment the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40add2a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#sim.download_packages([\"METIS\", \"ELT\", \"Armazones\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66118fe8",
   "metadata": {},
   "source": [
    "## Creation of a source - lamp with equally spaced lines"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd09d911",
   "metadata": {},
   "source": [
    "As an example, we use a calibration lamp with equally spaced and equally strong emission lines, covering the L band. The line list is turned into a spectrum by placing a narrow Gaussian at each line position. To simulate the lamp, we (ab)use the `flatlamp` function and replace the default spectrum (a black body) by the line spectrum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd445b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = np.arange(2.8, 4.2, 0.1)\n",
    "\n",
    "wave = np.linspace(2.8, 4.2, 4096)\n",
    "flux = np.zeros_like(wave)\n",
    "sigma = 0.0005\n",
    "for line in lines:\n",
    "    flux += 0.0003 * np.exp(-(wave - line)**2 / (2 * sigma**2))\n",
    "\n",
    "spec = SourceSpectrum(Empirical1D, points=wave*u.um, lookup_table=flux)\n",
    "\n",
    "src_linelamp = flatlamp()\n",
    "src_linelamp.fields[0].spectra[0] = spec     # NB: Do not try to set src_linelamp.spectra[0], this has no effect. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0adc8671",
   "metadata": {},
   "source": [
    "## Simulation of an observation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d84af6f1",
   "metadata": {},
   "source": [
    "We use METIS in the L-band long-slit spectroscopic mode, using a fairly narrow slit. We explicitely request the realistic spectral mapping with non-linear dispersion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44adab0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmds = sim.UserCommands(use_instrument=\"METIS\", set_modes=[\"lss_l\"],\n",
    "                       properties={\"!OBS.trace_file\": \"TRACE_LSS_L.fits\",\n",
    "                                   \"!OBS.slit\": \"B-28_6\"})\n",
    "\n",
    "metis = sim.OpticalTrain(cmds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5732f705",
   "metadata": {},
   "source": [
    "We exclude atmospheric emission (and absorption) and the telescope optics as is appropriate for a calibration-lamp observation. As the source fills the slit homogeneously a PSF convolution should have no effect on the result. Excluding PSF convolution cuts down significantly on computation time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d61f1c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "metis[\"skycalc_atmosphere\"].include = False\n",
    "metis[\"telescope_reflection\"].include = False\n",
    "metis[\"psf\"].include = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccafc4c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "metis.observe(src_linelamp, update=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fd86db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "readout = metis.readout(exptime=1)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54b7719d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(readout[1].data, origin=\"lower\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c035912b",
   "metadata": {},
   "source": [
    "## Rectification of the spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43f5145d",
   "metadata": {},
   "source": [
    "The non-linearity in the dispersion in METIS is small and not readily apparent.Still, rectification is necessary to arrive at a 2D spectrum with well-defined wavelength and spatial coordinates. The method to use is `rectify_traces` and belongs to the `SpectralTraceList` effect, which is accessible in the METIS `OpticalTrain` as `\"spectral_traces\"` (in MICADO it would be `\"micado_spectral_traces\"`. Currently, it is necessary to specify the spatial extent of the slit when calling the method. The long slit in METIS has a length of 8 arcsec and extends from -4 arcsec to +4 arcsec."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17c9c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "tracelist = metis[\"spectral_traces\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1649e402",
   "metadata": {},
   "outputs": [],
   "source": [
    "rectified = tracelist.rectify_traces(readout, -4.0, 4.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49469101",
   "metadata": {},
   "source": [
    "`rectified` is a HDU list with one extension for each spectral trace - for METIS there's only one trace, for MICADO there would be several. Each extension has a WCS that translates pixel coordinates into wavelength and position along the slit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df3b5a30",
   "metadata": {},
   "outputs": [],
   "source": [
    "wcs = WCS(rectified[1])\n",
    "naxis1 = rectified[1].header[\"NAXIS1\"]\n",
    "naxis2 = rectified[1].header[\"NAXIS2\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbd79196",
   "metadata": {},
   "outputs": [],
   "source": [
    "lam = (wcs.all_pix2world(np.arange(naxis1), 800, 0)[0] * u.Unit(wcs.wcs.cunit[0])).to(u.um).value\n",
    "xi = (wcs.all_pix2world(1000, np.arange(naxis2), 0)[1] * u.Unit(wcs.wcs.cunit[1])).to(u.arcsec).value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9827b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(rectified[1].data, origin=\"lower\", extent=[lam[0], lam[-1], xi[0], xi[-1]])\n",
    "plt.gca().set_aspect(\"auto\")\n",
    "plt.xlabel(r\"Wavelength [$\\mu$m]\")\n",
    "plt.ylabel(r\"Spatial position along slit [arcsec]\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b8c5c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "i1, i2 = 120, 620\n",
    "plt.figure(figsize=(15, 7))\n",
    "plt.plot(lam[i1:i2], rectified[1].data[800, i1:i2], label=\"single row\")\n",
    "plt.plot(lam[i1:i2], rectified[1].data.mean(axis=0)[i1:i2], label=\"average\")\n",
    "plt.legend()\n",
    "plt.xlabel(r\"Wavelength [$\\mu$m]\");"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}