cookbooks/WFC3/UVIS_fullframe_basic_extraction/wfc3-basic-uvis-fullframe-hstaxe-cookbook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d5ca3109",
   "metadata": {},
   "source": [
    "<a id=\"top\"></a>\n",
    "# HSTaXe Cookbook: Extraction for WFC3/UVIS Full Frame G280 Data\n",
    "\n",
    "This notebook contains a step-by-step guide for performing a basic spectral extraction with HSTaXe for G280 full frame data from WFC3/UVIS. <br>\n",
    "The original source for this notebook is the \"cookbook\" folder on the [spacetelescope/hstaxe](https://github.com/spacetelescope/hstaxe) GitHub repository. \n",
    "\n",
    "***\n",
    "## Learning Goals\n",
    "In this tutorial, you will:\n",
    "\n",
    "- Organize input data\n",
    "- Set up HSTaXe and prepare data for extraction\n",
    "- Learn how to handle different types of local background subtraction\n",
    "- Extract 1-D spectra \n",
    "- Display the results\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "[1. Introduction](#intro) <br>\n",
    "[2. Imports](#import) <br>\n",
    "[3. Setup](#setup) <br>\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[3.1 Verify Matching WCS Information](#wcs) <br>\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[3.2 Drizzling Input Data](#drizzle) <br>\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[3.3 Creating a Catalog with SExtractor](#catalog) <br>\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[3.4 Copy Catalog and Rename Mag Column](#copycat)<br>\n",
    "[4. Running HSTaXe](#axe) <br>\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[4.1. Outputs](#out) <br>\n",
    "[5. Conclusions](#conclusions) <br>\n",
    "[6. About this Notebook](#about) <br>\n",
    "[7. Citations](#cite) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ad558de",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1. Introduction <a id=\"intro\"></a>\n",
    "\n",
    "[HSTaXe](https://hstaxe.readthedocs.io/en/latest/index.html) is a Python package that provides spectral extraction processes for HST data. **Please be aware that running this notebook requires creating a conda environment from the provided `.yml` file in this notebook's [github repository](https://github.com/spacetelescope/hstaxe/tree/main/cookbooks).** For more details about creating the necessary environment see the notebook's README file.\n",
    "\n",
    "Below, we show the workflow for a basic specral extraction using WFC3/UVIS full frame G280 grism data. **The example data we use in this notebook are available [here](https://stsci.box.com/s/x0ndoduqkqd1ys1u9fgsz32urizsjdks).** If you would like to use this notebook with the example data, please dowload all required data in the `example_data` subdirectory from the link above, and store it within the same parent directory as this notebook. Once you have the example data directory, this notebook is intended to run continuously without needing to edit any of the cells. \n",
    "\n",
    "In addition to the grism and direct image data, **this notebook also requires configuration files for HSTaXe, which can be downloaded [here](https://stsci.box.com/s/ojbj3v4twkpkqt6tj9hwg1xkzbhrugm4).** Please download the folder titled `WFC3_UVIS_conf`. This configuration directory should be stored in the  same parent directory as this notebook, and later we will be copy them to the `CONF` subdirectory created by HSTaXe.\n",
    "\n",
    "**If you plan to use your own data with this notebook, please be aware you will be required to create an input source catalog with SExtractor.** Information regarding SExtractor including installation instructions are available [here](https://sextractor.readthedocs.io/en/latest/Installing.html). In addition to installing SExtractor, you must run the software with aXe specific configuration files. **These aXe-SExtractor configuration files can be downloaded [here](https://stsci.box.com/s/3npry36gu7ocfnuxgzwr5syj4i8r7hy8).** Once SExtractor is installed, create a `sextractor` directory in the same parent directory as this notebook, and place configuration files inside. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a91e02da",
   "metadata": {},
   "source": [
    "## 2. Imports <a id=\"import\"></a>\n",
    "\n",
    "For this workflow, we will import the following modules:\n",
    "\n",
    "- *os*, *glob* and *shutil*, for file handling\n",
    "- *numpy* for array handling\n",
    "- *astropy.io.fits* for FITS file handling\n",
    "- *matplotlib.pyplot* for plotting\n",
    "- *ginga.util.zscale* for calculating zscale limits for image display\n",
    "- *astrodrizzle* for creating input image mosaics\n",
    "- *hstaxe.axetasks* for performing the spectral extraction\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3e99ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import glob\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import shutil\n",
    "from astropy.io import fits, ascii\n",
    "from astropy.table import Table\n",
    "from ginga.util import zscale\n",
    "\n",
    "from drizzlepac import astrodrizzle\n",
    "from hstaxe import axetasks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3915d574",
   "metadata": {},
   "source": [
    "## 3. Setup <a id=\"setup\"></a>\n",
    "\n",
    "We'll start our basic extraction workflow by organizing our input data using the `example_data` directory downloaded from the link in the [Introduction](#intro). \n",
    "\n",
    "First, we save the working directory for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159c7dec",
   "metadata": {},
   "outputs": [],
   "source": [
    "cwd = os.getcwd()\n",
    "print(f'The current directory is: {cwd}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "672b237e",
   "metadata": {},
   "source": [
    "Next, we'll create directories for our grism and direct images. **HSTaXe will modify our input images in-place, so it is crucial to retain clean versions of them in another location, which will be copied into these directories.** If running this notebook multiple times, run all the lines in the next cell to clear any existing inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b991be",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(cwd)\n",
    "if os.path.isdir('grism_ims'):\n",
    "    shutil.rmtree('grism_ims')\n",
    "if os.path.isdir('direct_ims'):\n",
    "    shutil.rmtree('direct_ims')\n",
    "os.mkdir('grism_ims')\n",
    "os.mkdir('direct_ims')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9629fddd",
   "metadata": {},
   "source": [
    "Now, let us copy the full frame UVIS images to the input directories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47783cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "# src = '/path/to/your/grism/images/*.fits'\n",
    "src = 'example_data/grism/*flc.fits'\n",
    "dst = 'grism_ims/'\n",
    "for f in glob.glob(src):\n",
    "    shutil.copy(f, dst)\n",
    "\n",
    "# src = '/path/to/your/direct/images/*.fits'\n",
    "src = 'example_data/direct/*flc.fits'\n",
    "dst = 'direct_ims/'\n",
    "for f in glob.glob(src):\n",
    "    shutil.copy(f, dst)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6cb2ed",
   "metadata": {},
   "source": [
    "## 3.1  Verify Matching WCS Information <a id=\"wcs\"></a>\n",
    "\n",
    "It is possible that the WCS in the direct and grism images differ. In this section we will use a function to process all the direct and grism images to verify that the WCS information is consistent throughout. If there is any disagreement in WCS information we call `updatewcs` with the database keyword set to False, which will roll back all the solutions to the original distortion-corrected WCS. For more information regarding HST WCS and improved absolute astrometry please see [WFC3 Instrument Science Report 2022-06](https://ui.adsabs.harvard.edu/abs/2022wfc..rept....6M/abstract) (Mack et al. 2022). For documentations on `updatewcs` please see [here](https://stwcs.readthedocs.io/en/latest/updatewcs.html).\n",
    "\n",
    "Before running `updatewcs`, we need to [set CRDS environment variables](https://hst-crds.stsci.edu/docs/cmdline_bestrefs/).  We will point to a subdirectory called `crds_cache/` using the `IREF` environment variable. The `IREF` variable is used for WFC3 reference files and different instruments use other variables, e.g., `JREF` for ACS.\n",
    "\n",
    "You have the option to permanently add these environment variables to your user profile by adding the path in your shell's configuration file. If you're using bash, you would edit the `~/.bash_profile` file with lines such as:\n",
    "```\n",
    "export CRDS_SERVER_URL=\"https://hst-crds.stsci.edu\"\n",
    "export CRDS_SERVER=\"https://hst-crds.stsci.edu\"\n",
    "export CRDS_PATH=\"$HOME/crds_cache\"\n",
    "export iref=\"${CRDS_PATH}/references/hst/wfc3/\"\n",
    "```\n",
    "\n",
    "<div class=\"alert alert-block alert-warning\" style=\"color:black\" > <b> If you already have the absolute paths set for CRDS, you may skip the code cell immediately below and proceed to defining 'check_wcs'.</b> <br> </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24b5452",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CRDS_SERVER_URL\"] = \"https://hst-crds.stsci.edu\"\n",
    "os.environ['CRDS_SERVER'] = 'https://hst-crds.stsci.edu'\n",
    "if \"CRDS_PATH\" not in os.environ.keys():\n",
    "    os.environ[\"CRDS_PATH\"] = os.path.join(os.environ[\"HOME\"],\"crds_cache\")\n",
    "if \"iref\" not in os.environ.keys():\n",
    "    os.environ[\"iref\"] = \"$HOME/crds_cache/references/hst/wfc3/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "525f1d3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_wcs(images):\n",
    "    \"\"\" A helper function to verify the active world coordinate solutions match and roll them back if they do not \n",
    "    \n",
    "    Parameter\n",
    "    ---------\n",
    "    images : list \n",
    "        a list of grism and direct images \n",
    "        \n",
    "    Return\n",
    "    ------\n",
    "    N/A\n",
    "    \"\"\"\n",
    "    \n",
    "    direct_wcs = []\n",
    "    grism_wcs = []\n",
    "\n",
    "    for f in images:\n",
    "        # get filter for image to distinguish between direct and grism\n",
    "        filt = fits.getval(f, 'FILTER')\n",
    "       \n",
    "        hdul = fits.open(f)\n",
    "        db_bool = 'WCSCORR' not in hdul\n",
    "        hdul.close()\n",
    "        \n",
    "        try:\n",
    "            # get the active solution from the file's \"SCI\" extension\n",
    "            wcsname = fits.getval(f, 'WCSNAME', ext=('SCI', 1))\n",
    "            if db_bool == True:\n",
    "                updatewcs.updatewcs(f,use_db=db_bool)\n",
    "        except KeyError:\n",
    "            updatewcs.updatewcs(f,use_db=db_bool)\n",
    "            wcsname = fits.getval(f, 'WCSNAME', ext=('SCI', 1))\n",
    "        \n",
    "        # seperate between direct and grism\n",
    "        if 'G' in filt:\n",
    "            grism_wcs.append(wcsname)\n",
    "        if 'F' in filt:\n",
    "            direct_wcs.append(wcsname)\n",
    "\n",
    "    # get the number of unique active solutions in the direct and grism images       \n",
    "    num_wcs_direct = len(set(direct_wcs))\n",
    "    num_wcs_grism = len(set(grism_wcs))\n",
    "    \n",
    "    # roll back WCS on all files if there is more than one active solution for either direct or grism images\n",
    "    if num_wcs_direct > 1 or num_wcs_grism > 1:\n",
    "        [updatewcs.updatewcs(file,use_db=False) for file in images]\n",
    "        print('WCS reset complete')\n",
    "\n",
    "    # roll back WCS on all files if the active solution for the direct images do not match the grism images\n",
    "    elif set(direct_wcs) != set(grism_wcs):\n",
    "        [updatewcs.updatewcs(file,use_db=False) for file in images]\n",
    "        print('WCS reset complete')\n",
    "\n",
    "    # do nothing if there is one unique active solution and they match\n",
    "    elif set(direct_wcs) == set(grism_wcs):\n",
    "        print(f\"No WCS update needed. All grism and direct images us WCS: {grism_wcs[0]}.\")\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b8364db",
   "metadata": {},
   "source": [
    "The cell below also calls [CRDS bestref](https://hst-crds.stsci.edu/static/users_guide/basic_use.html), which will copy the necessary reference files from CRDS over to your local machine, if you do not already have them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c918dd2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_images = glob.glob('direct_ims/i*_flc.fits')+\\\n",
    "             glob.glob('grism_ims/i*_flc.fits')\n",
    "\n",
    "for image in all_images:\n",
    "    command_line_input = f'crds bestrefs --files {image} --types IDCTAB MDRIZTAB --sync-references=1 --update-bestrefs'\n",
    "    os.system(command_line_input)\n",
    "        \n",
    "check_wcs(all_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f079201a",
   "metadata": {},
   "source": [
    "## 3.2 Drizzling the Input Data<a id=\"drizzle\"></a>\n",
    "The next step is to drizzle the grism images. We'll need a list of the image names to feed to AstroDrizzle. After that, we'll do the same for the direct images, but use the drizzled grism image as a reference, which will ensure proper registration between the data. HSTaXe will use these linked drizzle images to locate spectral traces based on the positions of sources in the direct images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41cdf196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creat list file using images in grism directory\n",
    "os.chdir('grism_ims')\n",
    "\n",
    "lis = open('grism.lis', 'w')\n",
    "for f in sorted(os.listdir('.')):\n",
    "    if os.path.splitext(f)[1]=='.fits':\n",
    "        lis.write(f)\n",
    "        lis.write('\\n')\n",
    "lis.close()\n",
    "\n",
    "!cat grism.lis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f806c612",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Drizzle grism images. If only using one input image, set blot, median, driz_cr to False\n",
    "astrodrizzle.AstroDrizzle('@grism.lis', output='grism', mdriztab=True, \n",
    "                          preserve=False, skysub=True, skymethod= 'match', final_fillval=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679312fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List file for direct images\n",
    "os.chdir(cwd)\n",
    "os.chdir('direct_ims')\n",
    "\n",
    "lis = open('direct.lis', 'w')\n",
    "for f in sorted(os.listdir('.')):\n",
    "    if os.path.splitext(f)[1]=='.fits':\n",
    "        lis.write(f)\n",
    "        lis.write('\\n')\n",
    "lis.close()\n",
    "\n",
    "!cat direct.lis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c5204b",
   "metadata": {},
   "source": [
    "Next, drizzle the direct images using the drizzled grism mosaic as a reference to ensure proper registration. Note, in the cell below we have set the `AstroDrizzle` parameters for processing only a single image. **If you have more than one direct image to drizzle together please set the parameters appropriately.** For example `driz_separate`, `driz_sep_wcs`, `median`, `blot`, and `driz_cr` should all be set to `True`. For more information please see the `AstroDrizzle` documentation [here](https://drizzlepac.readthedocs.io/en/latest/astrodrizzle.html). Finally, if your input images were flt images rather than flc images, change the extension to `drz.fits` in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92e9d4e7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ref = '../grism_ims/grism_drc.fits[1]'\n",
    "astrodrizzle.AstroDrizzle(\"@direct.lis\", output=\"direct\", build=True, in_memory=False, preserve = False, \n",
    "                          skysub=True, skymethod= 'match', driz_separate=False, median=False, blot=False, driz_cr=False,\n",
    "                          final_wcs=True, final_refimage=ref)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f9187eb",
   "metadata": {},
   "source": [
    "Your grism and direct images should now be aligned. For the WFC3/UVIS grisms, there will be some vertical and horizontal offset between the positions of sources in the direct image and their correspondents in the grism image. For more information about G280 including appearance on the detector, please see WFC3 Instrument Handbook Section 8.2 [\"Slitless Spectroscopy with the UVIS G280 Grism\"](https://hst-docs.stsci.edu/wfc3ihb/chapter-8-slitless-spectroscopy-with-wfc3/8-2-slitless-spectroscopy-with-the-uvis-g280-grism). We'll perform a quick visual check here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c4ba5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(cwd)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(20,20))\n",
    "\n",
    "d = fits.getdata('grism_ims/grism_drc.fits', 1)\n",
    "z1,z2 = zscale.zscale(d)\n",
    "im1 = axes[0].imshow(d, origin='lower',cmap='Greys_r', vmin=z1, vmax=z2)\n",
    "axes[0].set_title('grism_drc.fits',size=14)\n",
    "fig.colorbar(im1,shrink=0.4,pad=0.01)\n",
    "\n",
    "d = fits.getdata('direct_ims/direct_drc.fits', 1)\n",
    "z1,z2 = zscale.zscale(d)\n",
    "im2 = axes[1].imshow(d, origin='lower',cmap='Greys_r', vmin=z1, vmax=z2)\n",
    "axes[1].set_title('direct_drc.fits',size=14)\n",
    "fig.colorbar(im2,shrink=0.4,pad=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30810db8",
   "metadata": {},
   "source": [
    "## 3.3 Creating a Catalog with SExtractor <a id=\"catalog\"></a>\n",
    "\n",
    "This section is intended for anyone using data other than the `example_data` provided for this notebook. Since we also provide a catalog in the `example_data` directory and will not formally run SExtractor here, we want to explain the process behind using the drizzle image to create a new catalog with SExtractor. **Please refer to the links in the [Introduction](#intro) section for instructions regarding installing SExtractor and downloading the necessary aXe-SExtractor configuration files.**\n",
    "\n",
    "HSTaXe will look for a highly specific format in the catalog, and does not always give clear error messages when something within the catalog is awry. If creating a catalog yourself, please follow the next steps carefully:\n",
    "\n",
    "1. Copy the drizzled direct image into the `sextractor` directory (once created), which contains HSTaXe-appropriate configuration files for SExtractor.\n",
    "\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<u>Example code:</u>\n",
    "```python\n",
    "shutil.copy('direct_ims/direct_drz.fits', 'sextractor/')\n",
    "```\n",
    "2. With SExtractor installed, run the following command from within the `sextractor` directory that you created:\n",
    "   \n",
    "   `sex -c aXe.sex direct_drz.fits[1] -DETECT_THRESH 5 -MAG_ZEROPOINT 26.4525`\n",
    "\n",
    "    Note that the value for the `DETECT_THRESH` keyword, which sets the minimum value for pixels to be considered, may be changed appropriately for your data. Similarly, `MAG_ZEROPOINT` corresponds to the UVIS zeropoint in AB mags. When using the UVIS grism there are generally only two main direct image filters to choose from, F200LP and F300X. For photometric information and filter zeropoints please see our [photometric calibration webpage](https://www.stsci.edu/hst/instrumentation/wfc3/data-analysis/photometric-calibration/uvis-photometric-calibration).\n",
    "    \n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<u>Example code:</u>\n",
    "```python\n",
    "os.chdir('sextractor')\n",
    "detect_thresh = 50\n",
    "mag_zp = 27.336\n",
    "cl_input = f'sex -c aXe.sex direct_drz.fits[1] -DETECT_THRESH {detect_thresh} -MAG_ZEROPOINT {mag_zp}'\n",
    "os.system(cl_input)\n",
    "``` \n",
    "3. At this point you should have created a catalog. The next steps include copying the file into the `direct_ims` directory and editing the name of the `MAG_ISO` column. See [Section 2.4](#copycat) for information regarding renaming the `MAG_ISO` column. \n",
    "\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<u>Example code:</u>\n",
    "```python\n",
    "os.chdir(cwd)\n",
    "shutil.copy('sextractor/aXe.cat', 'direct_ims')\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac2e06fa",
   "metadata": {},
   "source": [
    "## 3.4 Copy Catalog and Rename Mag Column <a id=\"copycat\"></a>\n",
    "\n",
    "The catalog corresponding to the data used in this notebook is included in the `example_data` directory downloaded in the Introduction section. Start by copying the catalog into the `direct_ims` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "087eb4cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the example catalog to the direct image directory:\n",
    "os.chdir(cwd)\n",
    "shutil.copy('example_data/aXe.cat', 'direct_ims')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7455e94e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat = Table.read('direct_ims/aXe.cat', format='ascii.sextractor')\n",
    "cat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0058a1c3",
   "metadata": {},
   "source": [
    "Examine the catalog. The \"MAG_ISO\" column must be renamed to \"MAG_F####\" for the catalog to be correctly read in by HSTaXe. Where \"####\" is the pivot wavelength of the direct image filter in Å (and nm for WFC3/IR e.g. 4971 for F200LP and 1392 for F140W). Please see the [WFC3/UVIS photometic calibration webpage](https://www.stsci.edu/hst/instrumentation/wfc3/data-analysis/photometric-calibration/uvis-photometric-calibration) for pivot wavelengths, which are listed in the `PHOTPLAM` column in the accordion container. \n",
    "\n",
    "Any lines in the catalog containing clearly spurious detections, such as those with magnitudes of $\\pm$99, should also be removed. **Note**: Removing spurious detections is not apart of this notebook and will need to be done manually. \n",
    "\n",
    "Lastly, locate the lines containing the sources whose spectra you want to extract and note the line number from the NUMBER column. This will be used later to identify the BEAM number for your object in the output files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdbf28e4",
   "metadata": {},
   "source": [
    "To avoid having to edit the column information manually in the SExtractor catalog, there is a helper function below called `edit_catalog_pivot`. It takes in the SExtractor catalog, output file path/name, and the pivot wavelength value and edits the information and writes to the output file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304bda79",
   "metadata": {},
   "outputs": [],
   "source": [
    "def edit_catalog_pivot(inputfile, outputfile, pivot_wave):\n",
    "    \"\"\" Function to edit the auto-generated sextractor header/column name so aXe will run\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        inputfile : str\n",
    "            The full path to the input catalog including filename\n",
    "        outputfile : str\n",
    "            The full path to the output catalog including filename\n",
    "        pivot_wave : int or str\n",
    "            The pivot wavelength of filter used in the driect image\n",
    "            For UVIS please use 4 digits in units of Angstrom and \n",
    "            for IR please use 4 ditits in units of nanometers \n",
    "            \n",
    "        Return\n",
    "        ------\n",
    "        Nothing. But a file is written to `outputfile`\n",
    "    \"\"\"\n",
    "    # Read in the input catalog\n",
    "    with open(inputfile, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "\n",
    "    with open(outputfile, 'w') as f:\n",
    "        # Find the mag_iso row and replace with pivot wavelength\n",
    "        for line in lines:\n",
    "            line = line.replace('MAG_ISO', f'MAG_F{pivot_wave}')\n",
    "            f.write(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b04c3fe",
   "metadata": {},
   "source": [
    "If you are using your own data, edit lines 1-3 in the cell below with your file paths and pivot wavelength. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ab7f6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputfile = 'direct_ims/aXe.cat'\n",
    "outputfile = 'direct_ims/aXe_uvis_f200lp.cat'\n",
    "pivot_wave = 4971 #angstrom because UVIS \n",
    "edit_catalog_pivot(inputfile, outputfile, pivot_wave)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "665d2c88",
   "metadata": {},
   "source": [
    "Open the new catalog with the renamed column and inspect it to make sure it looks good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d38ac667",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use this cell to filter your catalog for your sources\n",
    "# In this example, we sort to identify the brightest sources\n",
    "cat = Table.read(outputfile, format='ascii.sextractor')\n",
    "cat.sort('MAG_F4971')\n",
    "cat[['NUMBER','X_IMAGE','Y_IMAGE','A_IMAGE','B_IMAGE','THETA_IMAGE','FLUX_RADIUS','MAG_F4971']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b64214",
   "metadata": {},
   "source": [
    "## 4. Running HSTaXe<a id=\"aXe\"></a>\n",
    "\n",
    "Below, we show an example of using HSTaXe to complete a basic extraction. If you are interested in seeing an example of a more advanced extraction method that utilizes a fluxcube, please see the WFC3/IR cookbook [here](https://github.com/spacetelescope/hstaxe/tree/main/cookbooks/WFC3).<br>\n",
    "\n",
    "With the catalog generated and edited, we can now move on to working with HSTaXe. We'll set up a few additional directories and environment variables that point to them, while clearing out any previous data or outputs in these directories. We'll also copy our data into the fresh `DATA` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b693432",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(cwd)\n",
    "\n",
    "for dirr in ['DATA','CONF','OUTPUT']:\n",
    "    if os.path.isdir(dirr):\n",
    "        shutil.rmtree(dirr)\n",
    "    os.mkdir(dirr)\n",
    "\n",
    "os.environ['AXE_IMAGE_PATH'] = './DATA/' \n",
    "os.environ['AXE_CONFIG_PATH'] = './CONF/' \n",
    "os.environ['AXE_OUTPUT_PATH'] = './OUTPUT/' \n",
    "\n",
    "dsrc = 'direct_ims/*flc.fits'\n",
    "gsrc = 'grism_ims/*flc.fits'\n",
    "csrc = 'WFC3_UVIS_conf/*'\n",
    "\n",
    "for files in [dsrc, gsrc, csrc]:\n",
    "    if '_ims' in files:\n",
    "        dirr = 'DATA'\n",
    "    elif 'conf' in files:\n",
    "        dirr = 'CONF'\n",
    "    for f in glob.glob(files):\n",
    "        shutil.copy(f, dirr)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8662a6e1",
   "metadata": {},
   "source": [
    "Now we'll run `iolprep` to generate our IOLs, which are object catalogs for each individual direct image, from the drizzled direct image and its catalog."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9068904",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(cwd)\n",
    "os.chdir('direct_ims')\n",
    "\n",
    "axetasks.iolprep(drizzle_image = 'direct_drc.fits',\n",
    "                input_cat = 'aXe_uvis_f200lp.cat')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c82e3e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the IOLs to the aXe DATA directory\n",
    "os.chdir(cwd)\n",
    "for f in glob.glob('direct_ims/*_?.cat'):\n",
    "    print(f)\n",
    "    shutil.copy(f, 'DATA')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e616117",
   "metadata": {},
   "source": [
    "The last step before extracting the spectra is to generate a file which contains on each line the names of a grism image, IOL name, and associated direct image. This is best done manually, to ensure that each grism image lines up with the appropriate direct image. For the example data, a file called `aXe.lis` is provided. With this list, we run `axeprep`, which prepares the individual images for spectral extraction. \n",
    "\n",
    "Note that the naming convention in `aXe.lis` is different for UVIS than IR because UVIS has two CCD chips. When using full frame UVIS data, the correct format for each line of `aXe.lis` is the name of the grism image, the object catalog for UVIS2 followed by the object catalog for UVIS1 separated by a comma, and then the corresponding direct image. IOLs for UVIS2 end with `_1.cat` while UVIS1 IOLs end with `_2.cat`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d9d6c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment below to copy and rename the example list file\n",
    "os.chdir(cwd)\n",
    "shutil.copy('example_data/aXe.lis', '.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3785041",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(cwd)\n",
    "!cat aXe.lis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2671fc59",
   "metadata": {},
   "source": [
    "With this list, we run `axeprep`, which prepares the individual images for spectral extraction. This step is also responsible for performing a global background subtraction. The first UVIS G280 sky background reference file was created in 2023 and is detailed in [WFC3 ISR 2023-06](https://www.stsci.edu/files/live/sites/www/files/home/hst/instrumentation/wfc3/documentation/instrument-science-reports-isrs/_documents/2023/WFC3-ISR-2023-06.pdf). The sky file is included in the `WFC3_UVIS_conf` folder that was discussed and downloaded in the [Introduction](#intro). To use the G280 sky file, set the `backgr` parameter to `True` and list the two FITS files under the `backims` parameter such as: `\"WFC3.UVIS.G280.SKY.CHIP2.flc.fits,WFC3.UVIS.G280.SKY.CHIP1.flc.fits\"`.\n",
    "\n",
    "Note that the configuration files are listed in the same order below in `axeprep` as they are in the `aXe.lis` file; UVIS2 followed by UVIS1. Also note that HSTaXe needs the data in units of electrons per second and UVIS data is in units of electrons. Therefore the argument `norm` in `axeprep` must be set to `True` in order to have HSTaXe divide by the image's exposure time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a044ea24",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(cwd)\n",
    "\n",
    "axetasks.axeprep(inlist=\"aXe.lis\",\n",
    "                 configs=\"WFC3.UVIS.G280.CHIP2.V2.0.conf,WFC3.UVIS.G280.CHIP1.V2.0.conf\",\n",
    "                 backgr=False,\n",
    "                 backims=None,\n",
    "                 norm=True,\n",
    "                 mfwhm=3.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa06bc7f",
   "metadata": {},
   "source": [
    "The last HSTaXe task to run is `axecore`, which performs the actual extraction and generates output files. \n",
    "\n",
    "Local background subtraction is also performed by this step, if desired. The following keywords are critical for local background:\n",
    "\n",
    "* `back`: This argument, when set to TRUE, is the flag to trigger local background subtraction.\n",
    "* `np`: Defines the number of pixels on either side of the spectral trace (beam) used to calculate the local background from.\n",
    "* `interp`: Sets the interpolation method for the local background (-1=median, 0=mean, ≥1=nth order polynomial)\n",
    "* `backfwhm`: The FWHM specifying the width of the background pixel extraction table\n",
    "\n",
    "More information on background handling, both global and local, with `HSTaXe` can be found in the documentation [here](https://hstaxe.readthedocs.io/en/latest/hstaxe/description.html#sky-background). \n",
    "\n",
    "In addition to local background subtraction, HSTaXe is also able to perform a vertical extraction. This method of extraction requires editing the `THETA_IMAGE` column to `-90.0` in the object catalog (aXe.cat) and setting two keywords in `axecore`: `orient=True` and `slitless_geom=False`. Vertical extractions have been used in the past to handle the extreme curvature of the orders. For more information on the vertical extraction method please see [WFC3 ISR 2011-18](https://ui.adsabs.harvard.edu/abs/2011wfc..rept...18R/abstract) (Rothberg et al. 2011)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d147d692",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run axecore with no local background subtraction\n",
    "axetasks.axecore('aXe.lis',\n",
    "                 configs=\"WFC3.UVIS.G280.CHIP2.V2.0.conf,WFC3.UVIS.G280.CHIP1.V2.0.conf\",\n",
    "                 fconfterm=None,\n",
    "                 extrfwhm=4.,\n",
    "                 drzfwhm=3.,\n",
    "                 orient=False,\n",
    "                 weights=True,\n",
    "                 slitless_geom=False,\n",
    "                 cont_model='gauss',\n",
    "                 sampling='drizzle',\n",
    "                 exclude=True,\n",
    "                 back=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ef080fa",
   "metadata": {},
   "source": [
    "## 4.1. Outputs<a id=\"out\"></a>\n",
    "\n",
    "Each grism input file will have several corresponding output files. For each of the G280 input FLT file, HSTaXe will create the following in the `OUTPUT/` directory:\n",
    "\n",
    "- \\<ipppssoot>_flt_2.cat          : Object catalog for the FLT file [ipppssoot]_flt.fits<br>\n",
    "- \\<ipppssoot>_flt_2.OAF          : Aperture file<br>\n",
    "- \\<ipppssoot>_flt_2.PET.fits     : The Pixel Extraction Table, containing all the unbinned information about each spectrum<br>\n",
    "- \\<ipppssoot>_flt_2.SPC.fits     : 1D extracted spectra<br>\n",
    "- \\<ipppssoot>_flt_2.CONT.fits    : Contamination estimate for eact of the spectra<br>\n",
    "- \\<ipppssoot>_flt_2_opt.SPC.fits : Optimally extracted version of the 1D spectra\n",
    "\n",
    "For now, let's take a look at the STP files, which contain 2D \"stamps\" of the extracted spectral traces; and the SPC files, which contain our 1D extracted spectra. For all of the output files, those with the number 2 in their filename e.g. `ibr502i9q_flc_2.STP.fits` corresponds to data on UVIS2 and those with the number 5 in their filename e.g. `ibr502i9q_flc_5.SPC.fits` corresponds to data on UVIS1. \n",
    "\n",
    "We'll need the line numbers from the original source catalog we generated to identify the BEAM number for the objects whose spectra we want. For the example data, we'll use the target in the middle of UVIS2, which is number 3 in the provided catalog.\n",
    "\n",
    "We'll first look at the STP files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4bdab0a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!ls -ltr OUTPUT/*_?.SPC.fits\n",
    "!ls -ltr OUTPUT/*_?.STP.fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9edb5b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "beam = '3A'\n",
    "\n",
    "for i, f in enumerate(glob.glob('OUTPUT/*2.STP.fits')):\n",
    "    with fits.open(f) as hdul:\n",
    "        d = hdul[f'BEAM_{beam}'].data\n",
    "    z1,z2 = zscale.zscale(d)\n",
    "    plt.figure(figsize=(15,10),dpi=200)\n",
    "    im = plt.imshow(d, origin='lower',vmin=z1, vmax=z2)\n",
    "    plt.title(os.path.basename(f)+f' Beam {beam}')\n",
    "    plt.xlim(50,400)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05fa703c",
   "metadata": {},
   "source": [
    "And now, the SPC files. For the example data, we show the spectrum of superluminous supernova PTF12dam from HST PID WFC3/GO [12524](https://www.stsci.edu/cgi-bin/get-proposal-info?id=12524&submit=Go&observatory=HST) (PI Quimby). For more information about the spectrum and what the blended spectral features are please see [Quimby et al. 2018](https://ui.adsabs.harvard.edu/abs/2018ApJ...855....2Q/abstract)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f2a7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "beam = '3A'\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(11,7))\n",
    "for i, f in enumerate(glob.glob('OUTPUT/*2.SPC.fits')):\n",
    "    with fits.open(f) as hdul:\n",
    "        d = hdul[f'BEAM_{beam}'].data\n",
    "        h = hdul[0].header\n",
    "        \n",
    "    wl = d['LAMBDA']\n",
    "    flux = d['FLUX']\n",
    "    error = d['FERROR']\n",
    "    xrange = (wl>1900) & (wl<3500)\n",
    "    \n",
    "    ax.errorbar(wl[xrange],flux[xrange],error[xrange],label=os.path.basename(f)[:9],capsize=1)\n",
    "\n",
    "ax.legend()\n",
    "ax.grid(alpha=0.5)\n",
    "ax.set_ylim(4e-16,2.3e-15)\n",
    "ax.set_xlabel(r'Wavelength ($\\AA$)')\n",
    "ax.set_title(f\"{h['targname']}\",size=13)\n",
    "ax.set_ylabel(r'Flux ($erg/s/cm^2/\\AA$)')        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25bb13c0",
   "metadata": {},
   "source": [
    "## 5. Conclusions <a id=\"conclusions\"></a>\n",
    "\n",
    "Thank you for walking through this spectral extraction workflow. You should now be able to perform a basic extraction on WFC3/UVIS data using HSTaXe.\n",
    "\n",
    "For additional information on the WFC3 grisms, please visit the [grism resources](https://www.stsci.edu/hst/instrumentation/wfc3/documentation/grism-resources) and [grism data analysis](https://www.stsci.edu/hst/instrumentation/wfc3/documentation/grism-resources/grism-data-analysis) webpages.\n",
    "\n",
    "Further workflow cookbooks are available on the [HSTaXe GitHub](https://github.com/spacetelescope/HSTaXe), including a more advanced IR extraction, and a UVIS extraction. For detailed information on HSTaXe, please visit the [documentation webpage](https://hstaxe.readthedocs.io/en/latest/index.html).\n",
    "\n",
    "Lastly, if you have questions regarding this notebook or using WFC3 data with HSTaXe please contact our WFC3 [Help Desk](https://stsci.service-now.com/hst).\n",
    "\n",
    "**Congratulations, you have completed the notebook.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06194a54",
   "metadata": {},
   "source": [
    "## 6. About this Notebook <a id=\"about\"></a>\n",
    "\n",
    "**Author:** Benjamin Kuhn & Aidan Pidgeon WFC3 Instrument Team\n",
    "\n",
    "**Special Thanks to:** \n",
    " - Dr. Nor Pirzkal, for creating the original workflow that was adapted into this notebook\n",
    " - Ricky O'Steen and Duy Nguyen, for their fantastic work in updating the HSTaXe module\n",
    " - Debopam Som for support in testing the HSTaXe workflow\n",
    "\n",
    "**Released:** February 13, 2023 <br>\n",
    "**Last Updated:** January 9, 2024\n",
    "\n",
    "## 7. Citations <a id=\"cite\"></a>\n",
    "\n",
    "If you use `astropy`, `drizzlepac`, `matplotlib` or `numpy` for published research, please cite the\n",
    "authors. Follow this link for more information about citing the libraries:\n",
    "\n",
    "* [Citing `astropy`](https://www.astropy.org/acknowledging.html)\n",
    "* [Citing `drizzlepac`](https://drizzlepac.readthedocs.io/en/latest/LICENSE.html)\n",
    "* [Citing `matplotlib`](https://matplotlib.org/stable/users/project/citing.html)\n",
    "* [Citing `numpy`](https://numpy.org/citing-numpy/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84647a97-a47b-4d50-9fda-d6e92205ab2c",
   "metadata": {},
   "source": [
    "[Top of Page](#top)\n",
    "<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\"/> "
   ]
  }
 ],
 "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.10.9"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}