Merge pull request #84 from ojustino/horne-extract

Introducing Horne extraction

notebook_sandbox/compare_extractions.ipynb

@@ -0,0 +1,321 @@
+{
+ "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": 1,
+   "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.tile(np.arange(nrows), (ncols, 1))\n",
+    "index_arr.T"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7c5fb4b1-9c7d-4846-9963-b476070662ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "img = col_model(index_arr.T) + 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()\n",
+    "bxc_result1d_slice = bxc(img, trace, 14)\n",
+    "bxc_result1d_whole = bxc(img, trace, nrows)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "be1c3f96-efd4-4108-803f-7f1698fe2564",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hrn = HorneExtract()\n",
+    "hrn_result1d_whole = hrn(img, trace, 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()\n",
+    "hrn2_result1d_whole = hrn(img_obj, trace)"
+   ]
+  },
+  {
+   "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 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
+}