jwst_pipeline_mos_workshop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Full JWST Pipeline Stages\n",
    "https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/main.html#pipelines\n",
    "## The end-to-end calibration of JWST data is divided into 3 main stages of processing:\n",
    "\n",
    "- ### Stage 1 consists of detector-level corrections that are performed on a group-by-group basis, followed by ramp fitting. The output of stage 1 processing is a countrate image per exposure, or per integration for some modes. Details of this pipeline can be found at:\n",
    "\n",
    "    - #### calwebb_detector1: Stage 1 Detector Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_detector1.html#calwebb-detector1\n",
    "\n",
    "- ### Stage 2 processing consists of additional instrument-level and observing-mode corrections and calibrations to produce fully calibrated exposures. The details differ for imaging and spectroscopic exposures, and there are some corrections that are unique to certain instruments or modes. Details are at:\n",
    "\n",
    "    - #### calwebb_image2: Stage 2 Imaging Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_image2.html#calwebb-image2\n",
    "\n",
    "    - #### calwebb_spec2: Stage 2 Spectroscopic Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_spec2.html#calwebb-spec2\n",
    "\n",
    "- ### Stage 3 processing consists of routines that work with multiple exposures and in most cases produce some kind of combined product. There are unique pipeline modules for stage 3 processing of imaging, spectroscopic, coronagraphic, AMI, and TSO observations. Details of each are available at:\n",
    "\n",
    "    - #### calwebb_image3: Stage 3 Imaging Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_image3.html#calwebb-image3\n",
    "\n",
    "    - #### calwebb_spec3: Stage 3 Spectroscopic Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_spec3.html#calwebb-spec3\n",
    "\n",
    "    - #### calwebb_coron3: Stage 3 Coronagraphic Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_coron3.html#calwebb-coron3\n",
    "\n",
    "    - #### calwebb_ami3: Stage 3 Aperture Masking Interferometry (AMI) Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_ami3.html#calwebb-ami3\n",
    "\n",
    "    - #### calwebb_tso3: Stage 3 Time-Series Observation (TSO) Processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_tso3.html#calwebb-tso3\n",
    "\n",
    "- ### In addition, there are several pipeline modules designed for special instrument or observing modes, including:\n",
    "\n",
    "    - #### calwebb_dark for processing dark exposures: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_dark.html#calwebb-dark\n",
    "\n",
    "    - #### calwebb_guider for calibrating FGS guide star data: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_guider.html#calwebb-guider\n",
    "\n",
    "    - #### calwebb_wfs-image3 for stage 3 WFS&C processing: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_wfs-image3.html#calwebb-wfs-image3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This Notebook\n",
    "\n",
    "- ## This notebook focuses on the different processing stages using the NIRISS grism spectroscopic mode as example.\n",
    "- ## Starting with an uncalibrated raw image (uncal) of a simulated NIRISS scene, we will review all the steps to obtain and extract 2D and 1D calibrated grism spectra of the objects in the scene.\n",
    "- ### The NIRISS scene was prepared by Cami Pacifici with MIRAGE for the JWST GTO CANUCS team.\n",
    "    - #### MIRAGE: https://www.stsci.edu/jwst/science-planning/proposal-planning-toolbox/mirage\n",
    "    - #### CANUCS: https://www.stsci.edu/jwst/science-execution/program-information?id=1208\n",
    "    http://www.exoplanetes.umontreal.ca/wp-content/uploads/2018/11/ChrisWillott.pdf\n",
    "    - #### This simulation used data products from the JWST Extragalactic Mock Catalog: https://fenrir.as.arizona.edu/jaguar/index.html\n",
    "    \n",
    "- ## Read the associated README first.\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's start!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#system general\n",
    "import os\n",
    "\n",
    "#plotting and tabling\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "import astropy\n",
    "from astropy.io import fits\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import json\n",
    "\n",
    "#pipeline\n",
    "import jwst #(warning about findstars.py)\n",
    "from jwst.pipeline import Detector1Pipeline\n",
    "from jwst.pipeline import Image2Pipeline\n",
    "from jwst.pipeline import Image3Pipeline\n",
    "from jwst.pipeline import Spec2Pipeline\n",
    "\n",
    "#associations\n",
    "from jwst.associations import asn_from_list\n",
    "from jwst.associations.lib.rules_level2b import DMSLevel2bBase\n",
    "from jwst.associations.lib.rules_level3_base import DMS_Level3_Base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Versions')\n",
    "print('astropy:',astropy.__version__) #4.2.1\n",
    "print('pipeline:',jwst.__version__) #1.1.0\n",
    "print('numpy:',np.__version__) #1.20.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the uncalibrated raw images (direct and grisms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Uncal direct image (F115W)\n",
    "uncaldirectfile = \"inputs/NIRISS_F115W_direct_dit1_uncal.fits\"\n",
    "uncal_direct = fits.open(uncaldirectfile)\n",
    "\n",
    "#Uncal grism image (F115W, C orientation)\n",
    "uncalspecfileC = \"inputs/NIRISS_F115W_WFSSC_dit1_uncal.fits\"\n",
    "uncal_specC = fits.open(uncalspecfileC)\n",
    "\n",
    "#Uncal grism image (F115W, R orientation)\n",
    "uncalspecfileR = \"inputs/NIRISS_F115W_WFSSR_dit1_uncal.fits\"\n",
    "uncal_specR = fits.open(uncalspecfileR)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Review the direct image info and header\n",
    "print(uncal_direct.info())\n",
    "uncal_direct[0].header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Review the grism_C info and header\n",
    "print(uncal_specC.info())\n",
    "uncal_specC[0].header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Review the grism_R info and header\n",
    "print(uncal_specR.info())\n",
    "uncal_specR[0].header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Look at the uncalibrated direct and grism images\n",
    "plt.figure(figsize=(8,8))\n",
    "cax = plt.imshow(np.log10(uncal_direct[1].data[0][0]))\n",
    "cbar = plt.colorbar(cax)\n",
    "plt.title(\"Uncal Image 1/13\")\n",
    "\n",
    "plt.figure(figsize=(8,8))\n",
    "cax = plt.imshow(np.log10(uncal_specC[1].data[0][0]))\n",
    "cbar = plt.colorbar(cax)\n",
    "plt.title(\"Uncal Grism_C 1/25\")\n",
    "\n",
    "plt.figure(figsize=(8,8))\n",
    "cax = plt.imshow(np.log10(uncal_specR[1].data[0][0]))\n",
    "cbar = plt.colorbar(cax)\n",
    "plt.title(\"Uncal Grism_R 1/25\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Stage 1 of the pipeline.\n",
    "- ## This stage does the detector-level corrections and ramp fitting for individual exposures.\n",
    "- ## It takes the raw (uncal) images as input and delivers a countrate image per exposure.\n",
    "- ## This requires detector calibration files. They will be accessed via direct download if your CRDS cache is set up correctly (see README)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here run Stage 1 of the pipeline for the direct image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stage 1 for the direct image\n",
    "#Warning: this step can take quite some time if it requires to fetch the reference data from CRDS (1.4 GB).\n",
    "os.system('mkdir ./outputs')\n",
    "\n",
    "stage1_direct = Detector1Pipeline.call(uncaldirectfile, save_results=True, output_dir='outputs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now inspect the newly created countrate image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Get and plot newly created rate file\n",
    "ratedirectfile = 'outputs/NIRISS_F115W_direct_dit1_rate.fits'\n",
    "rateimage_direct = fits.open(ratedirectfile)\n",
    "\n",
    "#Stage 1 also creates two additionnal files: *_rateints.fits and *_trapsfilled.fits\n",
    "# *_rateints.fits is the countrate product for individual integrations.\n",
    "# *_trapsfilled.fits is the charge traps per pixels.\n",
    "# Details: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_detector1.html\n",
    "\n",
    "print(rateimage_direct.info())\n",
    "print(rateimage_direct[1].header)\n",
    "\n",
    "#centered on an object\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_direct[1].data),clim=(-1,2))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_direct[3].data),clim=(0,10))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()\n",
    "\n",
    "#full frame\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_direct[1].data),clim=(-1,2))\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_direct[3].data),clim=(0,10))\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here run Stage 1 of the pipeline for the grism images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Perform Stage 1 for the grism_C and grism_R images\n",
    "#This step can take some time.\n",
    "\n",
    "stage1_grismC = Detector1Pipeline.call(uncalspecfileC, save_results=True, output_dir='outputs')\n",
    "stage1_grismR = Detector1Pipeline.call(uncalspecfileR, save_results=True, output_dir='outputs')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now inspect the newly created countrate images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Get and plot newly created rate files\n",
    "ratespecfileC = 'outputs/NIRISS_F115W_WFSSC_dit1_rate.fits'\n",
    "ratespecfileR = 'outputs/NIRISS_F115W_WFSSR_dit1_rate.fits'\n",
    "rateimage_specC = fits.open(ratespecfileC)\n",
    "rateimage_specR = fits.open(ratespecfileR)\n",
    "#It also creates the *_rateints.fits and *_trapsfilled.fits files for each orientation.\n",
    "\n",
    "print(rateimage_specC.info())\n",
    "print(rateimage_specR.info())\n",
    "print(rateimage_specC[1].header)\n",
    "\n",
    "#show grism_C\n",
    "#centered on an object\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specC[1].data),clim=(-1,2))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specC[3].data),clim=(0,10))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()\n",
    "\n",
    "#full frame\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specC[1].data),clim=(-1,2))\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specC[3].data),clim=(0,10))\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#show grism_R\n",
    "#centered on an object\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specR[1].data),clim=(-1,2))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specR[3].data),clim=(0,10))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()\n",
    "\n",
    "#full frame\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specR[1].data),clim=(-1,2))\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((rateimage_specR[3].data),clim=(0,10))\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's all for Stage 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Stage 2 of the pipeline on the direct image only. Note that stage 2 for the grism data can only be run after running Stage 3 on the direct image as it requires some data products from that step.\n",
    "- ## This stage does the instrument-mode calibrations for individual exposures.\n",
    "- ## It takes the countrate images from Stage 1 as input and delivers calibrated, unrectified images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here run Stage 2 of the pipeline for the direct image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Perform Stage 2 for the direct image\n",
    "#This step can take some time. Especially if it needs to download reference files from CRDS first.\n",
    "\n",
    "stage2_direct = Image2Pipeline.call(ratedirectfile, save_results=True, output_dir='outputs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now inspect the newly created calibrated image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#The newly created file (calibrated, unrectified) is NIRISS_F115W_direct_dit1_cal.fits.\n",
    "#Details: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_image2.html\n",
    "\n",
    "#Grab and view\n",
    "caldirectfile = 'outputs/NIRISS_F115W_direct_dit1_cal.fits'\n",
    "calimage_direct = fits.open(caldirectfile)\n",
    "\n",
    "print(calimage_direct.info())\n",
    "\n",
    "#plot the calibrated image\n",
    "#centered on an object\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((calimage_direct[1].data),clim=(-1,2))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((calimage_direct[3].data),clim=(0,10))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()\n",
    "\n",
    "#full frame\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((calimage_direct[1].data),clim=(-1,2))\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((calimage_direct[3].data),clim=(0,10))\n",
    "plt.title(\"DQ array\")\n",
    "plt.colorbar()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's all for Stage 2 for the direct image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Stage 3 of the pipeline on the direct image only.\n",
    "- ## Stage 3 processing for direct-imaging observations is intended for combining the calibrated data from multiple exposures (e.g. a dither or mosaic pattern) into a single rectified (distortion corrected) product. Before being combined, the exposures receive additional corrections for the purpose of astrometric alignment, background matching, and outlier rejection.\n",
    "- ## The inputs to the Stage 3 pipeline are one or more Stage 2 calibrated (“_cal”) image products. In order to process and combine multiple images, an ASN file must be used as input, listing the exposures to be processed. It is also possible to use a single “_cal” file as input to Stage 3, in which case only the resample and source_catalog steps will be applied."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here create the association file for Stage 3 of the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stage 3 requires an association file to be created first\n",
    "#We only use a single Stage 2 image here.\n",
    "association_direct_file = './NIRISS_F115W_direct_dit1_stage3.json'\n",
    "\n",
    "association_direct = asn_from_list.asn_from_list([caldirectfile], rule=DMS_Level3_Base, product_name='NIRISS_F115W_direct_dit1')\n",
    "association_direct['asn_type'] = 'Image3'\n",
    "\n",
    "with open(association_direct_file, 'w') as outfile:\n",
    "    outfile.write(association_direct.dump()[1])\n",
    "print('Saved association file for cal file to %s' % association_direct_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Read and show association file\n",
    "f = open(association_direct_file)\n",
    "data = json.load(f)\n",
    "print(data)\n",
    "\n",
    "#make copy to outputs directory\n",
    "os.system('cp '+association_direct_file+' outputs/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here run Stage 3 of the pipeline for the direct image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Perform Stage 3 of the reduction pipeline for the direct image.\n",
    "#We only use a single Stage 2 image here. Only the resample and source_catalog steps will be applied.\n",
    "\n",
    "stage3_direct = Image3Pipeline.call(association_direct_file, save_results=True, output_dir='outputs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stage 3 creates three files: *_i2d.fits, *_cat.ecsv and *_segm.fits. \n",
    "# *_i2d.fits is the rectified product.\n",
    "# *_cat.ecsv is the object catalog.\n",
    "# *_segm.fits is the segmentation image associated to the object catalog.\n",
    "#Details: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_image3.html\n",
    "\n",
    "#Get and view the rectified image.\n",
    "i2ddirectfile = 'outputs/NIRISS_F115W_direct_dit1_i2d.fits'\n",
    "i2dimage_direct = fits.open(i2ddirectfile)\n",
    "\n",
    "print(i2dimage_direct.info())\n",
    "\n",
    "#plot the resampled image\n",
    "#centered on an object\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((i2dimage_direct[1].data),clim=(-1,2))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((i2dimage_direct[3].data))\n",
    "plt.xlim(700,800)\n",
    "plt.ylim(150,250)\n",
    "plt.title(\"Weight image\")\n",
    "plt.colorbar()\n",
    "\n",
    "#full frame\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((i2dimage_direct[1].data),clim=(-1,2))\n",
    "plt.title(\"Science image\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.figure(figsize=(12,12))\n",
    "plt.imshow((i2dimage_direct[3].data))\n",
    "plt.title(\"Weight image\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here inspect the newly created object catalog."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Step 3 created 'NIRISS_F115W_direct_dit1_i2d.fits'.\n",
    "#It also created an object catlaog 'NIRISS_F115W_direct_dit1_cat.ecsv'.\n",
    "#And a segmentation image 'NIRISS_F115W_direct_dit1_segm.fits'.\n",
    "#For the purpose of this notebook, we will use the object catalog provided in the 'inputs' directory,\n",
    "#as it only contains the 'real' detections.\n",
    "#This object catalog is needed to run stage 2 on the grism data.\n",
    "\n",
    "#object_cat_file = 'outputs/NIRISS_F115W_direct_dit1_cat.ecsv' #--> Should only contains 3 spurious detections.\n",
    "object_cat_file = 'inputs/NIRISS_F115W_direct_dit1_cat_clean.ecsv' #--> This is the clean (original) catalog.\n",
    "object_cat = pd.read_table(object_cat_file, comment='#', sep='\\s+')\n",
    "object_cat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here inspect the objects on the science image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#See objects on the science image 'i2dimage_direct'\n",
    "segimage_direct = fits.open('outputs/NIRISS_F115W_direct_dit1_segm.fits')\n",
    "segimage_direct.info()\n",
    "for i in range (0, len(object_cat)):\n",
    "    #object in science image\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.imshow((i2dimage_direct[1].data),clim=(0,1.5))\n",
    "    plt.xlim(object_cat.loc[i,'xcentroid']-50,object_cat.loc[i,'xcentroid']+50)\n",
    "    plt.ylim(object_cat.loc[i,'ycentroid']-50,object_cat.loc[i,'ycentroid']+50)\n",
    "    plt.colorbar()\n",
    "    plt.scatter(object_cat['xcentroid'], object_cat['ycentroid'], s=800, facecolors='none', edgecolors='r')\n",
    "    plt.annotate(object_cat.loc[i,'id'], xy=(object_cat.loc[i,'xcentroid']-0.5, object_cat.loc[i,'ycentroid']+5), color='white')\n",
    "    #segmentation map\n",
    "    plt.figure(figsize=(8,8))\n",
    "    plt.imshow((segimage_direct[1].data),clim=(0,1.5))\n",
    "    plt.xlim(object_cat.loc[i,'xcentroid']-50,object_cat.loc[i,'xcentroid']+50)\n",
    "    plt.ylim(object_cat.loc[i,'ycentroid']-50,object_cat.loc[i,'ycentroid']+50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's all for Stage 3 for the direct image.\n",
    "Onto Stage 2 for the grism images now."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now, we can run Stage 2 of the pipeline on the grism images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here create the association files that are needed for Stage 2.\n",
    "For each grism image, the association file needs to contain the grism countrate info as well as the source catalog information. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stage 2 for grism images first requires to create association files.\n",
    "#For this, we need the object catalog from Stage 3 of the direct image processing 'object_cat_file',\n",
    "#as well as the grismC and grismR rate files from Stage 1 of grism processing ('ratespecfileC' and 'ratespecfileR'). \n",
    "association_grismC_file = './NIRISS_F115W_WFSSC_dit1_stage2.json'\n",
    "association_grismR_file = './NIRISS_F115W_WFSSR_dit1_stage2.json'\n",
    "\n",
    "#grismC\n",
    "association_grismC = asn_from_list.asn_from_list([ratespecfileC], rule=DMSLevel2bBase, product_name='NIRISS_F115W_WFSSC_dit1')\n",
    "association_grismC['asn_type'] = 'None'\n",
    "association_grismC['products'][0]['members'].append({})\n",
    "association_grismC['products'][0]['members'][1][\"expname\"] = object_cat_file\n",
    "association_grismC['products'][0]['members'][1][\"exptype\"] = \"sourcecat\"\n",
    "with open(association_grismC_file, 'w') as outfile:\n",
    "    outfile.write(association_grismC.dump()[1])\n",
    "print('Saved association file to %s' % association_grismC_file)\n",
    "\n",
    "#grismR\n",
    "association_grismR = asn_from_list.asn_from_list([ratespecfileR], rule=DMSLevel2bBase, product_name='NIRISS_F115W_WFSSR_dit1')\n",
    "association_grismR['asn_type'] = 'None'\n",
    "association_grismR['products'][0]['members'].append({})\n",
    "association_grismR['products'][0]['members'][1][\"expname\"] = object_cat_file\n",
    "association_grismR['products'][0]['members'][1][\"exptype\"] = \"sourcecat\"\n",
    "with open(association_grismR_file, 'w') as outfile:\n",
    "    outfile.write(association_grismR.dump()[1])\n",
    "print('Saved association file to %s' % association_grismR_file)\n",
    "\n",
    "#make copies to outputs directory\n",
    "os.system('cp ' + association_grismC_file + ' ' + association_grismR_file + ' outputs/')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We here run Stage 2 of the pipeline for the grism images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Run Stage 2\n",
    "#This step can take quite some time.\n",
    "os.system('cp '+object_cat_file+' .')\n",
    "\n",
    "stage2_grismC = Spec2Pipeline.call(association_grismC_file, save_results=True, output_dir='outputs')\n",
    "stage2_grismR = Spec2Pipeline.call(association_grismR_file, save_results=True, output_dir='outputs')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Two files per grism have been created at this step.\n",
    "#The clabrated 2D grism spectra: *_cal.fits\n",
    "#And the extracted 1D spectra: *_x1d.fits\n",
    "\n",
    "#See infos of the calibrated 2D spectra\n",
    "calgrismfileC = 'outputs/NIRISS_F115W_WFSSC_dit1_cal.fits'\n",
    "calgrismfileR = 'outputs/NIRISS_F115W_WFSSR_dit1_cal.fits'\n",
    "calimage_specC = fits.open(calgrismfileC)\n",
    "calimage_specR = fits.open(calgrismfileR)\n",
    "\n",
    "calimage_specC.info()\n",
    "calimage_specR.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now inspect the extracted grism spectra of the detected objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#See a few 2D extracted spectra\n",
    "ids = [1, 2, 15, 16, 19]\n",
    "print(calimage_specC[1].header)\n",
    "for i in ids:\n",
    "    spec2D_C = calimage_specC[1+7*(i-1)].data\n",
    "    spec2D_R = calimage_specR[1+7*(i-1)].data\n",
    "    #grism C\n",
    "    plt.figure(figsize=(15,15))\n",
    "    plt.subplot(2,1,1)\n",
    "    ys,xs = spec2D_C.shape\n",
    "    plt.imshow(spec2D_C, extent=[0,xs,0,ys], norm=LogNorm(vmin=1e-1,vmax=1e1))\n",
    "    plt.colorbar()\n",
    "    plt.title('ID'+str(i)+'_C')\n",
    "    #grism R\n",
    "    plt.figure(figsize=(15,15))\n",
    "    plt.subplot(2,1,2)\n",
    "    ys,xs = spec2D_R.shape\n",
    "    plt.imshow(spec2D_R, extent=[0,xs,0,ys], norm=LogNorm(vmin=1e-1,vmax=1e1))\n",
    "    plt.colorbar()\n",
    "    plt.title('ID'+str(i)+'_R')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now inspect the extracted 1D grism spectra."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#See infos of the 1D extracted images\n",
    "x1dgrismfileC = 'outputs/NIRISS_F115W_WFSSC_dit1_x1d.fits'\n",
    "x1dgrismfileR = 'outputs/NIRISS_F115W_WFSSR_dit1_x1d.fits'\n",
    "x1dimage_specC = fits.open(x1dgrismfileC)\n",
    "x1dimage_specR = fits.open(x1dgrismfileR)\n",
    "\n",
    "x1dimage_specC.info()\n",
    "x1dimage_specR.info()\n",
    "x1dimage_specC[1].header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#See a few 1D extracted spectra\n",
    "ids = [1, 2, 15, 16, 19]\n",
    "F115W_min = 10000\n",
    "F115W_max = 13000\n",
    "for i in ids:\n",
    "    spec1D_C = x1dimage_specC[i].data\n",
    "    spec1D_R = x1dimage_specR[i].data\n",
    "    #grism C\n",
    "    wave = spec1D_C['WAVELENGTH']*1e4 #angstroms\n",
    "    flux = spec1D_C['FLUX'] * 2.99792458e-5 / wave**2 #ergs/s/cm2/Ang\n",
    "    flux_err = spec1D_C['ERROR'] * 2.99792458e-5 / wave**2 #ergs/s/cm2/Ang\n",
    "    cut = (wave > F115W_min) & (wave < F115W_max)\n",
    "    plt.figure(figsize=(12,8))\n",
    "    plt.subplot(2,1,1)\n",
    "    plt.errorbar(wave[cut], flux[cut], drawstyle='steps-mid')#, yerr=flux_err[cut]\n",
    "    plt.title('ID'+str(i)+'_C')\n",
    "    plt.ylim(0, 1.5e-18)\n",
    "    plt.xlim(F115W_min, F115W_max)\n",
    "    #grism R\n",
    "    wave = spec1D_R['WAVELENGTH']*1e4 #angstroms\n",
    "    flux = spec1D_R['FLUX'] * 2.99792458e-5 / wave**2 #ergs/s/cm2/Ang\n",
    "    flux_err = spec1D_R['ERROR'] * 2.99792458e-5 / wave**2 #ergs/s/cm2/Ang\n",
    "    cut = (wave > F115W_min) & (wave < F115W_max)\n",
    "    plt.figure(figsize=(12,8))\n",
    "    plt.subplot(2,1,2)\n",
    "    plt.errorbar(wave[cut], flux[cut], drawstyle='steps-mid')#, yerr=flux_err[cut])\n",
    "    plt.title('ID'+str(i)+'_R')\n",
    "    plt.ylim(0, 1.5e-18)\n",
    "    plt.xlim(F115W_min, F115W_max)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's all for Stage 2 of the pipeline for the grism images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stage 3 for grism images combines calibrated data from multiple exposures (eg, dithering) into combined 2D and 1D spectra. Details: https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_spec3.html\n",
    "\n",
    "# This is however not part of this notebook.\n",
    "\n",
    "# All done, congrats! :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----------\n",
    "2021 May 07\n",
    "<br>Gaël Noirot\n",
    "<br>Postdoctoral Fellow\n",
    "<br>Saint Mary's University"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}