{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IFU Spectral Extraction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook demonstrates spectral extraction from a [NIRSpec IFU](https://jwst-docs.stsci.edu/near-infrared-spectrograph/nirspec-observing-modes/nirspec-ifu-spectroscopy) datacube. It is a streamlined version of the [ifu_optimal](https://spacetelescope.github.io/jdat_notebooks/notebooks/ifu_optimal/ifu_optimal.html) notebook from the [JWST Data Analysis Notebooks](https://spacetelescope.github.io/jdat_notebooks/) collection.\n",
    "\n",
    "* Quicklook extraction with [Cubeviz](https://jdaviz.readthedocs.io/en/latest/cubeviz/index.html)\n",
    "* Retrieving spectra, spatial regions, and spectral regions from Cubeviz\n",
    "* Creating and manipulating specutils Spectrum1D objects\n",
    "    * Slicing\n",
    "    * Extracting a spectral sub-region\n",
    "    * Simple arithmetic\n",
    "* Aperture extraction in a notebook \n",
    "* Optimal extraction in a notebook\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "A faint (quasar) point source was simulated using the NIRSpec Instrument Performance Simulator (IPS), then run through the Spec2 stage of the [JWST pipeline](https://jwst-pipeline.readthedocs.io/en/latest/).   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* _numpy_ for array math\n",
    "* _scipy_ for ndcube gaussian smoothing\n",
    "* _specutils_ for Spectrum1D data model and cube manipulation\n",
    "* _jdaviz_ for the Cubeviz data visualization tool\n",
    "* _jwst_ for JWST data models\n",
    "* _photutils_ to define circular apertures\n",
    "* _astropy.io_ for reading and writing FITS cubes and images\n",
    "* _astropy.wcs, units, coordinates_ for defining and reading WCS\n",
    "* _astropy.stats_ for sigma_clipping\n",
    "* _astropy.utils_ for downloading files from URLs\n",
    "* _astropy.units_ for unit definitions\n",
    "* _astropy.nddata_ for uncertainty object types\n",
    "* _matplotlib_ for plotting spectra and images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Resize notebook to full width\n",
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "\n",
    "import specutils\n",
    "from specutils import Spectrum1D, SpectralRegion\n",
    "from specutils.manipulation import extract_region, spectral_slab\n",
    "from specutils.analysis import moment\n",
    "\n",
    "from jdaviz import Cubeviz\n",
    "\n",
    "from jwst import transforms\n",
    "\n",
    "from photutils import CircularAperture, SkyCircularAperture, aperture_photometry\n",
    "from photutils.detection import DAOStarFinder\n",
    "from photutils.centroids import centroid_com, centroid_1dg, centroid_2dg, centroid_quadratic\n",
    "\n",
    "from regions import PixCoord, CirclePixelRegion\n",
    "\n",
    "from astropy.io import fits\n",
    "from astropy import wcs\n",
    "from astropy.stats import sigma_clip\n",
    "from astropy.stats import sigma_clipped_stats\n",
    "from astropy.utils.data import download_file\n",
    "import astropy.units as u\n",
    "from astropy.nddata import StdDevUncertainty, VarianceUncertainty\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulated NIRSpec IFU Cube\n",
    "\n",
    "A point source (quasar) was simulated using the NIRSpec Instrument Performance Simulator (IPS), then run through the JWST Spec2 pipeline. We will use this for our science dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Download the cube data from Box (bright quasar)\n",
    "#filename = \"https://stsci.box.com/shared/static/ff0bj31acot1272x5qq2clbmto3s6v4f.fits\"\n",
    "\n",
    "# NIRSpec IFU science data cube (faint quasar)\n",
    "BoxPath = \"https://data.science.stsci.edu/redirect/JWST/jwst-data_analysis_tools/IFU_optimal_extraction/\"\n",
    "filename = BoxPath + \"NRS00001-faintQSO-F100LP-G140H-01_1_491_SE_2020-08-25T12h15m00_s3d.fits\"\n",
    "\n",
    "cube_file = download_file(filename, cache=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Science Data with Cubeviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cubeviz = Cubeviz()\n",
    "cubeviz.app"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### UI Instructions:\n",
    "#### Task 1:  Load the cube and view it\n",
    "* Load science datacube into Cubeviz using the next code cell below\n",
    "* Click the Display icon in the leftmost image viewer \n",
    "* Change the stretch to Logarithmic, 99 percentile. \n",
    "* Manually adjust the bias and contrast sliders until the faint point source at (x,y) = (17,21) becomes visible\n",
    "* Scrub through the cube using the Slice slider.\n",
    "\n",
    "#### Task 2:  Extract a spectrum from a region (subset) of the cube\n",
    "* Select a circular subset region centered on the source, using the region selection tool under the Tools icon. \n",
    "* Note that the region is pixelated and doesn't include fractional pixels.\n",
    "* Change the collapse method function to \"Sum\" in spectrum viewer: Display icon: Viewer. \n",
    "* Hide the first [SCI] layer by clicking the \"eye\" icon in Display icon: Layer. This will make the noisy, full-cube spectrum disappear.\n",
    "* Adjust the size and location of the region until you are happy with the extraction that shows up in the spectrum viewer.\n",
    "\n",
    "#### Task 3: Select a spectral region for further analysis\n",
    "* In the spectrum viewer, use the vertical and horizontal pan and zoom tools to focus on the broad H-beta line (redshifted to 1.22 microns).  To pan: drag and drop spectrum with mouse or trackpad. To zoom: scroll with mouse or trackpad.\n",
    "* Select a spectral subset region in spectrum viewer, using the region selection tool under the Tools icon, centered on a spectral feature of interest.\n",
    "* Open the Line Analysis plugin to measure various line parameters in spectral subset 2.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Cube into Cubeviz\n",
    "\n",
    "*Developer notes:* \n",
    "* *Multiple spectral_cube warnings occur when loading the cube into Cubeviz*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Developer Note--spectral_cube warnings:\n",
    "cubeviz.app.load_data(cube_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export Region from Cubeviz\n",
    "Export the region defined by the user in Cubeviz as an astropy CirclePixel Region, which has units of pixels.  Use \"try:, except:\" to catch the case where there is no Subset 1 selected in Cubeviz. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    region1 = cubeviz.app.get_subsets_from_viewer('flux-viewer')['Subset 1']\n",
    "    print(region1)\n",
    "    region1_exists = True\n",
    "    center1_xy = [region1.center.x, region1.center.y]  \n",
    "    r_pix = region1.radius\n",
    "\n",
    "except Exception:\n",
    "    print(\"There are no regions selected in the cube viewer.\")\n",
    "    region1_exists = False\n",
    "    center1_xy = [17.1, 20.]\n",
    "    r_pix = 6.0\n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieve Subset Spectrum from Cubeviz Spectrum Viewer\n",
    "Get the spectrum of user-defined Subset1 from the Spectrum Viewer as a Spectrum1D object. Trim to remove bad wavelength ranges.\n",
    "\n",
    "*Developer notes:* \n",
    "* *Can't convert Subset1 Spectrum1D wavelengths from meters to microns.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check what data are available\n",
    "#print(cubeviz.app.data_collection)\n",
    "print(cubeviz.app.get_data_from_viewer('spectrum-viewer'))\n",
    "\n",
    "# try..except in case there is no Subset1\n",
    "try:\n",
    "    spec1d_subset1_untrimmed = cubeviz.app.get_data_from_viewer('spectrum-viewer')['Subset 1']\n",
    "    print(spec1d_subset1_untrimmed)\n",
    "    \n",
    "except Exception:\n",
    "    print(\"There are no subsets selected in the spectrum viewer.\")\n",
    "    \n",
    "#Developer Note--Attribute Error trying to convert to microns\n",
    "#spec1d_subset1_untrimmed = spec1d_subset1_untrimmed.with_spectral_unit(u.um)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trim the Spectral Cube\n",
    "Read in the spectral cube as a specutils Spectrum1D object, slice it, and trim the wavelength axis using the spectral_slab convenience function.\n",
    "\n",
    "*Developer notes:* \n",
    "* *Can't print Spectrum1D read with JWST s3d format. See [specutils issue #868]( https://github.com/astropy/specutils/issues/868).*\n",
    "* *spectral slab gives warning for Subset1 spectrum: 'No Observer defined on WCS, Spectral Coord'*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load the cube as a specutils Spectrum1D object \n",
    "spec1d_untrimmed = Spectrum1D.read(cube_file, format='JWST s3d')\n",
    "wavelength_untrimmed = spec1d_untrimmed.spectral_axis.value\n",
    "\n",
    "#Developer Note--can't print the Spectrum1D.  \n",
    "#See https://github.com/astropy/specutils/issues/868\n",
    "#print(spec1d_untrimmed)\n",
    "\n",
    "#Good spectral and spatial regions\n",
    "wave_trim =[1.0*u.um,1.44*u.um]\n",
    "x_trim = [2,-1]\n",
    "y_trim = [5,-4]\n",
    "\n",
    "#Trim the spectral and spatial axes in one go\n",
    "#Spectrum1D objects can be sliced like numpy arrays...\n",
    "#...and/or use the convenience function specutils.manipulation.spectral_slab\n",
    "spec1d = spectral_slab(spec1d_untrimmed, wave_trim[0], wave_trim[1])[x_trim[0]:x_trim[1],y_trim[0]:y_trim[1],:]\n",
    "wavelength = spec1d.spectral_axis\n",
    "spec1d_len = len(wavelength.value)\n",
    "print('Trimmed spectral axis:')\n",
    "print(wavelength)\n",
    "\n",
    "# Trim the Subset1 spectrum to match\n",
    "#Developer note--\"No observer defined on WCS\":\n",
    "#spec1d_subset1 = spectral_slab(spec1d_subset1_untrimmed, wave_trim[0], wave_trim[1])\n",
    " \n",
    "#Adjust extraction region location in trimmed cube\n",
    "center_trim = PixCoord(x=center1_xy[0]-x_trim[0], y=center1_xy[1]-y_trim[0])\n",
    "center_x = center_trim.x\n",
    "center_y = center_trim.y\n",
    "region_trim = CirclePixelRegion(center=center_trim, radius=r_pix)\n",
    "print()\n",
    "print('Adjusted Extraction Region:')\n",
    "print(region_trim)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Spectrum in Expanding Circular Aperture (Cone)\n",
    "This method is appropriate for a point source PSF with width proportional to wavelength. In the next cell, we read the cube into a Spectrum1D object and perform aperture photometry on each wavelength slice using photutils.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Use photutils.aperture_photometry to measure flux in expanding or constant aperture\n",
    "aperture = 'cone'  # 'cone' or 'cylinder'\n",
    "flux_sum = []\n",
    "for idx in range(spec1d_len):\n",
    "    if aperture == 'cone':\n",
    "        r_cone = r_pix * wavelength.value[idx]/ wavelength.value[0]\n",
    "    else: r_cone = r_pix  \n",
    "    aperture_cone = CircularAperture((center_x,center_y), r=r_cone)\n",
    "    phot_table = aperture_photometry(spec1d.flux.value[:, :, idx], aperture_cone)\n",
    "    flux_sum.append(phot_table['aperture_sum'][0])\n",
    "    \n",
    "cone = Spectrum1D(flux=np.array(flux_sum)*u.MJy/u.sr, spectral_axis=spec1d.spectral_axis)\n",
    "print(cone)\n",
    "\n",
    "# Compute the RMS-summed error over the full FOV, in each spectral bin\n",
    "err_sum = np.sqrt(np.sum(spec1d_untrimmed.uncertainty.array**2, axis=(0,1)))\n",
    "print(\"Max RMS Uncertainty: \", np.amax(np.nan_to_num(err_sum)))\n",
    "\n",
    "#Compare conical extraction to Cubeviz pixelated cylindrical extraction\n",
    "f, (ax1,ax2) = plt.subplots(2, 1, figsize=(10, 8)) \n",
    "ax1.set_title(\"Spectral extractions\")\n",
    "ax1.set_xlabel(\"Observed Wavelength (microns)\")  \n",
    "ax1.set_ylabel(\"Flux Density\")\n",
    "ax1.set_xlim(0.99, 1.45)\n",
    "ax1.set_ylim(0, 0.4)\n",
    "ax1.plot(wavelength, cone.flux.value, label=\"Cone\", c='darkorange', alpha=0.5)\n",
    "try:\n",
    "    ax1.plot(spec1d_subset1.spectral_axis.value*1.0E6, spec1d_subset1.flux.value, c='r', label=\"Subset1\", alpha=0.4)\n",
    "except Exception:\n",
    "    print(\"There is no Cubeviz Subset1 spectrum to plot.\")\n",
    "ax1.legend()\n",
    "\n",
    "ax2.set_xlabel(\"Observed Wavelength (microns)\")  \n",
    "ax2.set_ylabel(\"Uncertainty\")\n",
    "ax2.plot(wavelength_untrimmed, err_sum, c='r', label=\"RMS Uncertainty\", alpha=0.4)\n",
    "ax2.set_xlim(0.99, 1.45)\n",
    "ax2.set_ylim(0, 0.1)\n",
    "ax2.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Top: Comparison of the conical and Cubeviz subset spectral extractions. \n",
    "The conical extraction captures slightly more flux at long wavelengths.\n",
    "Red-shifted Broad Balmer and narrow [O III] lines  are visible in the quasar spectra. \n",
    "\n",
    "Bottom: RMS uncertainty, sqrt(sum(sigma^2)), summed over full IFU footprint. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## WebbPSF Model PSF for Optimal Extraction\n",
    "Generate PSF model cube using WebbPSF for NIRSpec IFU, or read in precomputed PSF model cube.\n",
    "\n",
    "Caution! The WebbPSF model takes about 10 hr to run. Uncomment the following cell to do so. Otherwise, read in the precomputed WebbPSF model, which covers the full F100LP/G140H wavelength range (blue and red). For other filter/grating combinations, uncomment and run the cell below using the wavelengths from the science data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "#WebbPSF imports\n",
    "%pylab inline\n",
    "import webbpsf\n",
    "\n",
    "#WebbPSF commands used to create PSF model cube\n",
    "ns.image_mask = \"IFU\"  # Sets to 3x3 arcsec square mask\n",
    "ns = webbpsf.NIRSpec()\n",
    "wavelengths = wavelength*1.0E-6\n",
    "psfcube = ns.calc_datacube(wavelengths, fov_pixels=30, oversample=4,  add_distortion=True)\n",
    "psfcube.writeto(\"Webbpsf_ifucube.fits\")\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BoxPath = \"https://data.science.stsci.edu/redirect/JWST/jwst-data_analysis_tools/IFU_optimal_extraction/\"\n",
    "psf_filename = BoxPath+\"Webbpsf_ifucube.fits\"\n",
    "\n",
    "# Load with astropy.fits.open\n",
    "with fits.open(psf_filename, memmap=False) as hdulist:\n",
    "    psf_model = hdulist['DET_SAMP'].data\n",
    "    psf_hdr = hdulist['DET_SAMP'].header\n",
    "    hdulist.info()    \n",
    "print(psf_model.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Align Model PSF Cube with Science Data\n",
    "\n",
    "Find the (x,y) offset between the Model and simulated PSF peaks. Shift and smooth the model PSF cube to align with the simulated data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flip model PSF left-right because WebbPSF is flipped with respect to the IPS simulation.\n",
    "psf_model = psf_model[:, ::-1, :]\n",
    "\n",
    "#Measure centroids of source and PSF model\n",
    "#Developer note--centroid_1dg and centroid_2dg give errors here; centroid_quadratic is not robust:\n",
    "cube_sum = moment(spec1d, order=0)\n",
    "psf_model_sum = np.sum(psf_model, axis=0)\n",
    "x1, y1 = centroid_com(cube_sum)\n",
    "x2, y2 = centroid_com(psf_model_sum)\n",
    "shiftx = x1 - x2\n",
    "shifty = y1 - y2\n",
    "print((x1,y1))\n",
    "print((x2,y2))\n",
    "\n",
    "#Manual adjustment\n",
    "#For some reason, the centroiding is not perfect...\n",
    "shiftx = shiftx - 0.1\n",
    "shifty = shifty + 0.5\n",
    "print(\"shift = \", (shiftx, shifty))\n",
    "\n",
    "#Shift model PSF using linear interpolation\n",
    "#The x and y shifts are interchanged, per python indexing convention\n",
    "psf_model_aligned = scipy.ndimage.shift(psf_model, (0.0, shifty, shiftx), order=1, \n",
    "                                        mode='constant', cval=0.0, prefilter=True)\n",
    "\n",
    "#Smooth model PSF\n",
    "#EMSM smoothing for G140H grating\n",
    "scalerad = 0.046  # sigma (arcsec)\n",
    "pixelscale = 0.1\n",
    "scalerad_pix = scalerad / pixelscale\n",
    "smooth_psf = False\n",
    "if smooth_psf:\n",
    "    psf_model_smoothed = scipy.ndimage.filters.gaussian_filter(psf_model_aligned, \n",
    "                                                           (0.0, scalerad_pix, scalerad_pix), \n",
    "                                                           order=0, mode='reflect', cval=0.0,  \n",
    "                                                           truncate=10.0)\n",
    "else:\n",
    "    psf_model_smoothed = psf_model_aligned\n",
    "    \n",
    "    \n",
    "#PSF Spectrum1D\n",
    "#We transpose the datacube because the last axis of a Spectrum1D flux array must be the spectral axis\n",
    "#Trim the PSF model in the wavelength dimension to match the data \n",
    "spec1d_webb = Spectrum1D(flux=np.transpose(psf_model_smoothed[0:1872, :, :])*u.Jy, spectral_axis=spec1d.spectral_axis)\n",
    "\n",
    "#PSF sums over wavelength and spaxels\n",
    "psf_model_sum = np.sum(psf_model_smoothed, axis=0)\n",
    "psf_model_fnusum = np.sum(psf_model, axis=(1, 2))\n",
    "\n",
    "# Scale factor for PSF subtraction\n",
    "scalefactor = np.amax(cube_sum) / np.amax(psf_model_sum)\n",
    "\n",
    "# Sum model PSF over spaxels\n",
    "\n",
    "\n",
    "# Plot the collapsed model and data PSFs\n",
    "f, ([ax1, ax2, ax3], [ax4, ax5, ax6]) = plt.subplots(2, 3, figsize=(10, 10)) \n",
    "ms, mew = 15., 2.\n",
    "ax1.set_title(\"PSF slice sum\")\n",
    "ax1.imshow(psf_model_sum,origin = 'lower',norm=LogNorm())\n",
    "ax1.plot(x1, y1, color='red', marker='+', ms=ms, mew=mew)\n",
    "\n",
    "ax2.set_title(\"Science Data slice sum\")\n",
    "ax2.imshow(cube_sum, origin='lower', norm=LogNorm()) \n",
    "ax2.plot(x1, y1, color='red', marker='+', ms=ms, mew=mew)\n",
    "\n",
    "ax3.set_title(\"Data / PSF Ratio\")\n",
    "ax3.imshow(cube_sum / psf_model_sum, origin='lower', norm=LogNorm())\n",
    "\n",
    "ax4.set_title(\"PSF Model integrated flux\")\n",
    "ax4.plot(psf_model_fnusum)\n",
    "\n",
    "ax5.set_title(\"Data - PSF\")\n",
    "ax5.imshow(cube_sum - scalefactor * psf_model_sum, origin='lower')\n",
    "\n",
    "#Developer Note--Unit Conversion Error:\n",
    "#im6 = ax6.imshow(np.log10(np.absolute(cube_sum - scalefactor * psf_model_sum)))\n",
    "im6 = ax6.imshow(np.absolute(cube_sum - scalefactor * psf_model_sum), norm=LogNorm())\n",
    "#plt.colorbar(im6)\n",
    "ax6.set_title(\"abs(Data - PSF)\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Top row: Collapsed WebbPSF Model PSF, simulated quasar PSF, and their ratio. The ratio will be used for the optimal extraction below. The red cross symbols show the computed centroids of the two aligned PSFs.\n",
    "\n",
    "Bottom left: The integrated flux decreases from 98% to 97% of the total model flux over the wavelength range of the dataset.\n",
    "\n",
    "Bottom center, bottom right : Difference and absolute value of difference between the aligned science data and model PSFs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimal Extraction using WebbPSF Model\n",
    "\n",
    "Optimal extraction ([Horne 1986, PASP, 98, 609](https://ui.adsabs.harvard.edu/abs/1986PASP...98..609H/abstract)) weights the flux contributions to a spectrum by their signal-to-noise ratio (SNR). Dividing the simulated data by the model PSF gives an estimate of the total flux density spectrum in each spaxel. A weighted average of these estimates over all spaxels yields the optimally extracted spectrum over the cube. In the faint source limit, where the noise is background-dominated, optimal extraction inside a 3-sigma radius can increase the effective exposure time by a factor of 1.69 (Horne et al. 1986). In the bright source limit, where the noise is dominated by the Poisson statistics of the source, optimal extraction is formally identical to a straight sum over spaxels for a perfect PSF model.\n",
    "An in-flight reference star observation aligned with the science target PSF could be used in place of the WebbPSF model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mask out bad data using 3-sigma clipping in each slice\n",
    "# Developers note: Need a function to do sigma-clipping of Spectrum1D\n",
    "spec1d_norm = spec1d/spec1d_webb\n",
    "nan_cleaned = np.nan_to_num(spec1d_norm.flux.value)\n",
    "spec1d_norm_clipped = sigma_clip(nan_cleaned, sigma=3.0, maxiters=5, axis=(0, 1))\n",
    "badvoxel = np.where(spec1d_norm_clipped == 0)\n",
    "spec1d_clean = spec1d\n",
    "spec1d_clean.flux.value[badvoxel] = 0.0\n",
    "\n",
    "# Variance\n",
    "spec1d_var = Spectrum1D(flux=spec1d.uncertainty.array*u.Jy, spectral_axis=spec1d.spectral_axis)\n",
    "\n",
    "#Horne algorithm\n",
    "spec1d_optimal_weight = spec1d_webb * spec1d_webb / spec1d_var   # SNR^2\n",
    "weight_sum = np.sum(spec1d_optimal_weight.flux.value, axis=(0,1))\n",
    "spec1d_optimal_cube = spec1d_webb * spec1d_clean / spec1d_var\n",
    "spec_optimal = np.sum(spec1d_optimal_cube.flux.value, axis=(0,1)) / weight_sum\n",
    "spec1d_optimal = Spectrum1D(flux=spec_optimal*u.Jy, spectral_axis=spec1d.spectral_axis)\n",
    "\n",
    "# Plots\n",
    "f, (ax1) = plt.subplots(1, 1, figsize=(12, 6)) \n",
    "ax1.set_title(\"Optimal Extraction Comparison\")\n",
    "ax1.set_xlabel(\"Observed Wavelength (microns)\") \n",
    "ax1.set_ylabel(\"Flux Density\")\n",
    "ax1.set_ylim(0, 0.5)\n",
    "ax1.plot(wavelength, cone.flux, label=\"Conical Extraction\", alpha=0.5)\n",
    "ax1.plot(wavelength, 1.35 * spec1d_optimal.flux.value, label=\"1.35 * Optimal\")\n",
    "ax1.legend()\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimally extracted spectrum generally shows lower noise compared to the conical extraction, for relatively wide conical extraction apertures (r > 2 pixels). The PSF of the IPS-simulated dataset is not perfectly matched to the WebbPSF simulation, giving ~35% loss of flux in the optimally extracted spectrum. Using a reference PSF from IPS gives a higher fidelity optimal extraction. Ideally, one would use the in-flight PSF of a bright reference star for optimal extraction of in-flight data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img style=\"float: right;\" src=\"https://raw.githubusercontent.com/spacetelescope/notebooks/master/assets/stsci_pri_combo_mark_horizonal_white_bkgd.png\" alt=\"Space Telescope Logo\" width=\"200px\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook created by Patrick Ogle."
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}