notebook_sandbox/compare_extractions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d1f2d88b-b715-487d-8e96-0e644ed3e7d2",
   "metadata": {},
   "source": [
    "# Extraction algorithm comparison\n",
    "\n",
    "We create a synthetic 2D image, upon which we compare the results of `specreduce`'s `BoxcarExtract` and `HorneExtract` algorithms.\n",
    "\n",
    "Since we control amplitude/uncertainty/etc., we can check that the results match our expectations. Among other things, we expect the Horne extraction's signal-to-noise ratio to outperform the boxcar's when using the whole scene as the aperture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1b17df2-81fe-49e7-9083-d5ce3f278937",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc741345-720e-4189-9748-b3bb98d2d693",
   "metadata": {},
   "outputs": [],
   "source": [
    "import astropy.units as u\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from astropy.modeling import models\n",
    "from astropy.nddata import CCDData, VarianceUncertainty\n",
    "from specreduce.extract import BoxcarExtract, HorneExtract\n",
    "from specreduce.tracing import FlatTrace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "000a0fc1-90e8-4472-bc8e-42a1c1b67506",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'axes.titlesize': 18, 'axes.labelsize': 12,\n",
    "                     'legend.fontsize': 12,  'axes.grid': False,\n",
    "                     'grid.alpha': .5, 'grid.color': 'k',\n",
    "                     'axes.edgecolor': 'k'})\n",
    "np.random.seed(7) # use same random values in different sessions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13866c09-5b84-4fa9-be19-04e6aafa8ac1",
   "metadata": {},
   "source": [
    "## Create a 2D image\n",
    "\n",
    "The flux in each column will follow a Gaussian distribution that we set using `astropy`'s modeling functionality.\n",
    "\n",
    "We also add normally distributed noise throughout the image to make the extraction more difficult."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0471b5d-907c-4c23-8b24-9ac2e7a619a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nrows = 50*4\n",
    "ncols = 40*4\n",
    "sigma_pix = 4\n",
    "sigma_noise = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff677978-744f-4264-a1d5-c51a32727d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "col_model = models.Gaussian1D(amplitude=1, mean=nrows/2, stddev=sigma_pix)\n",
    "noise = np.random.normal(scale=sigma_noise, size=(nrows, ncols))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f96eb2e-b6ac-4396-88ec-f2011e335107",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_arr = np.broadcast_to(np.arange(nrows)[:, np.newaxis], (nrows, ncols))\n",
    "index_arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c5fb4b1-9c7d-4846-9963-b476070662ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = col_model(index_arr) + noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05c8c9be-e553-4ab1-8f01-1edd84d3ca25",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,8))\n",
    "im1 = ax.imshow(img, cmap='magma', origin='lower', vmin=0, vmax=1)\n",
    "ax.set_title('synthetic 2D image')\n",
    "fig.colorbar(im1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7161a77b-87ee-40ec-b780-64a1e0510c8b",
   "metadata": {},
   "source": [
    "In addition to the image, we also create variance and mask images for `HorneExtract`. The former was defined above; we use an array of zeros for the latter since this image has no \"bad\" pixels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ba40e0-8bc9-49c1-80f7-943697c43cfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "variance = np.tile(sigma_noise, img.shape)\n",
    "mask = np.zeros_like(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe8bfe43-54d8-44e0-b60e-c0ee835e4621",
   "metadata": {},
   "source": [
    "## Create a trace\n",
    "\n",
    "Here, we manually set the trace to the middle row of the 2D image's y-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb3e2f73-e863-4fed-b919-4a1991344c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "trace = FlatTrace(img, nrows/2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9341c3a4-6ca4-4460-9aa1-37f8332f4193",
   "metadata": {},
   "source": [
    "## Calculate the extractions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7103992-4462-4bd9-8805-89c21ac17ed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "bxc = BoxcarExtract(img, trace)\n",
    "bxc_result1d_slice = bxc(width=14)\n",
    "bxc_result1d_whole = bxc(width=nrows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be1c3f96-efd4-4108-803f-7f1698fe2564",
   "metadata": {},
   "outputs": [],
   "source": [
    "hrn = HorneExtract(img, trace)\n",
    "hrn_result1d_whole = hrn(variance=variance,\n",
    "                         mask=mask, unit=u.DN) # whole image is aperture"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d547548e-6ff9-4672-bd03-aa666b3affdf",
   "metadata": {},
   "source": [
    "Note that `HorneExtract` can also take an `NDData` or `CCDData` image object as an argument. These are convenient because they allow for compiling the image, the variance, a mask, and any units into a single object. \n",
    "\n",
    "Once that's created, the only other argument needed is the trace. See this example with `CCDData`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b31779-2080-4541-b0b4-6a3d08b37062",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = np.zeros_like(img)\n",
    "var_obj = VarianceUncertainty(variance)\n",
    "img_obj = CCDData(img, uncertainty=var_obj, mask=mask, unit=u.DN)\n",
    "\n",
    "hrn2 = HorneExtract(img_obj, trace)\n",
    "hrn2_result1d_whole = hrn2()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e45324-a42b-46cf-9512-04cb0be8db47",
   "metadata": {},
   "source": [
    "The results are the same either way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cbf978f-1567-4558-8789-3e8783d29be4",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.array_equal(hrn_result1d_whole.flux, hrn2_result1d_whole.flux)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a9de2c3-f318-4d9e-bbd0-21f883da4650",
   "metadata": {},
   "source": [
    "## Compare results\n",
    "The whole-image extractions come out as expected, with the Horne-extracted 1D spectrum showing a noticeably better signal-to-noise ratio than its windowless boxcar equivalent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fd7c15f-c8a0-40e2-af68-173eaaad0a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig2, ax2 = plt.subplots(figsize=(10, 8))\n",
    "\n",
    "ax2.plot(hrn_result1d_whole.flux, c='#1d1160', label='1D spectrum, Horne, whole')\n",
    "ax2.plot(bxc_result1d_whole.flux, c='cadetblue',\n",
    "         label='1D spectrum, boxcar, whole', alpha=.5)\n",
    "ax2.plot(bxc_result1d_slice.flux, c='cadetblue', linestyle='--',\n",
    "         label='1D spectrum, boxcar, slice')\n",
    "ax2.axhline(sigma_pix * np.sqrt(2*np.pi), c='#d4bd8a', linestyle='--',\n",
    "            label=r'target ($\\sigma_{spatial}$ * $\\sqrt{2\\pi}$)')\n",
    "\n",
    "ax2.legend(fontsize=12)#, loc=(1.05,.5))\n",
    "ax2.set_title('extracted 1D spectra')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44ef28f4-959f-491a-8e0c-d5477f66d7be",
   "metadata": {},
   "source": [
    "The boxcar extraction can produce a similar result when its aperture is sliced to remove edge pixels. (Of course, that comes with the cost of losing any information those pixels contained.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb606fad-34fc-4adf-b52b-06035f4d1f23",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3, ax3 = plt.subplots(figsize=(10, 4))\n",
    "\n",
    "ax3.plot(hrn_result1d_whole.flux / np.nanmax(hrn_result1d_whole.flux),\n",
    "         c='#1d1160', label='1D spectrum, Horne, whole')\n",
    "# ax3.plot(bxc_result1d_whole.flux / np.nanmax(bxc_result1d_whole.flux),\n",
    "#          c='cadetblue', label='1D spectrum, boxcar, whole')\n",
    "ax3.plot(bxc_result1d_slice.flux / np.nanmax(bxc_result1d_slice.flux),\n",
    "         c='cadetblue', label='1D spectrum, boxcar, slice', linestyle='--')\n",
    "\n",
    "ax3.legend(fontsize=12)\n",
    "ax3.set_title('extracted 1D spectra, normalized')"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}