diff --git a/notebooks/NIRCam/NIRCam_WFSS_simulating_spectra/Simulating_WFSS_spectra_CRDS.ipynb b/notebooks/NIRCam/NIRCam_WFSS_simulating_spectra/Simulating_WFSS_spectra_CRDS.ipynb
new file mode 100644
index 000000000..8031e8a01
--- /dev/null
+++ b/notebooks/NIRCam/NIRCam_WFSS_simulating_spectra/Simulating_WFSS_spectra_CRDS.ipynb
@@ -0,0 +1,1443 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "01fe963e-c1b9-460a-9f34-eee3724f909e",
+ "metadata": {},
+ "source": [
+ "# WFSS Spectra Simulation for Contamination Correction"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4e68aca2-165e-4d67-9b7c-ce06ec9f0f69",
+ "metadata": {},
+ "source": [
+ "This notebook demonstrates basic techniques to simulate WFSS dispersed spectra, given source locations from an imaging mode exposure and a generalized World Coordinate System (gWCS) from an accompanying WFSS exposure.\n",
+ "\n",
+ "Such simulations, when applied to all sources in the field, are crucial to estimate the contamination of overlapping spectra, as well as to mask dispersed spectral traces when trying to estimate the background.\n",
+ "\n",
+ "This notebook builds on the simpler Box Extraction notebook where we introduced the general concepts of spectral extraction as well as the GRISMCONF module, which provides us with a low level interface to the gWCS model for WFSS modes."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a646f783-b72b-427c-95fa-df38ae1f41ad",
+ "metadata": {},
+ "source": [
+ "The simulation method is relatively straight-foward and begins with an imaging mode exposure. \n",
+ "\n",
+ "* We first determine which pixels contain the source and the signal levels in those pixels.\n",
+ "\n",
+ "* We then use the gWCS from the imaging mode file and also the WFSS observation to compute the corresponding location of each pixel in the frame of the WFSS observation. \n",
+ "\n",
+ "* Using the GRISMCONF functions and a wavelength vector, we next simulate the dispersion of the source. For each of the source pixels, we move through the wavelength vector and calculate the location where each wavelength value is dispersed to. We calculate the signal level associated with each wavelength position based on the flux of each input pixel in $F_{\\lambda}$ units ($erg/s/cm^2/A$). \n",
+ "\n",
+ "* Since the calculated coordinates of the wavelength values will not, in general, align with the detector pixel grid, we use the Sutherland-Hodgman algorithm to compute the fraction of each projected dispersed pixel that overlaps onto each detector pixel. Summing the signals from each "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1bf8bfdc-197b-4b13-8ac2-3ca575fccaad",
+ "metadata": {},
+ "source": [
+ "Author: N. Pirzkal
\n",
+ "Date created: 12 December 2024"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "256467d4-7d69-4223-b7c4-ab2489bf7a29",
+ "metadata": {},
+ "source": [
+ "## Table of Contents\n",
+ "1. [Set CRDS Path and Server](#Set-CRDS-Path-and-Server)\n",
+ "2. [Package Imports](#Package-Imports)\n",
+ "3. [Define Functions and Parameters](#Define-Functions-and-Parameters)\n",
+ "4. [Download Data](#Download-Data)\n",
+ "5. [Run Pipeline Steps](#Run-Pipeline-Steps)\n",
+ "6. [Detect Sources](#Detect-Sources)\n",
+ "7. [Simulate spectrum of one source](#Simulate-spectrum-of-one-source)\n",
+ " * [Locate source in imaging and WFSS data](#Locate-source-in-imaging-and-WFSS-data)\n",
+ " * [Get wavelength information](#Get-wavelength-information)\n",
+ " * [Simulate the dispersion of a single pixel](#Simulate-the-dispersion-of-a-single-pixel)\n",
+ " * [Disperse all the pixels for our source](#Disperse-all-the-pixels-for-our-source)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "fe743322-ce2d-4525-9fc4-d0361f8ab089",
+ "metadata": {},
+ "source": [
+ "## Set CRDS Path and Server"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "927b85b2-3d8b-43bf-8597-f1a59a55412b",
+ "metadata": {},
+ "source": [
+ "Before running the pipeline steps, we need to ensure our our CRDS environment is configured. This includes defining a CRDS cache directory in which to keep the reference files that will be used by the calibration pipeline.\n",
+ "\n",
+ "If the root directory for the local CRDS cache has not already been set, it will be created in the home directory. This needs to be done before importing the crds package or any package that has crds as a dependency."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "ae2c8be0-65cd-438b-afad-f8c0b4f55ac0",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "6a68f79c-6e9e-4200-a7a9-638a9f5d1f86",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Check whether the local CRDS cache directory has been set.\n",
+ "# If not, set it to the user home directory\n",
+ "if (os.getenv('CRDS_PATH') is None):\n",
+ " os.environ['CRDS_PATH'] = os.path.join(os.path.expanduser('~'), 'crds')\n",
+ " \n",
+ "# Check whether the CRDS server URL has been set. If not, set it.\n",
+ "if (os.getenv('CRDS_SERVER_URL') is None):\n",
+ " os.environ['CRDS_SERVER_URL'] = 'https://jwst-crds.stsci.edu'\n",
+ "\n",
+ "# Echo CRDS path and context in use\n",
+ "print('CRDS local filepath:', os.environ['CRDS_PATH'])\n",
+ "print('CRDS file server:', os.environ['CRDS_SERVER_URL'])\n",
+ "\n",
+ "# Import crds after setting up the required environment variables\n",
+ "from crds import client\n",
+ "if client.get_crds_server() != os.environ['CRDS_SERVER_URL']:\n",
+ " client.set_crds_server('https://jwst-crds.stsci.edu')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "760baeb3-e191-4e6c-a426-dc927c07edda",
+ "metadata": {},
+ "source": [
+ "## Package Imports"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e35257cb-61d9-40d0-a883-a59563649ec3",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from astropy.convolution import convolve\n",
+ "from astropy.io import fits\n",
+ "from copy import deepcopy\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "import requests\n",
+ "from scipy.sparse import coo_matrix\n",
+ "import tqdm\n",
+ "\n",
+ "import grismconf\n",
+ "from jwst import datamodels\n",
+ "from jwst.assign_wcs import AssignWcsStep\n",
+ "from jwst.flatfield import FlatFieldStep\n",
+ "from jwst.photom import PhotomStep\n",
+ "from photutils.background import Background2D, MedianBackground\n",
+ "from photutils.segmentation import make_2dgaussian_kernel\n",
+ "from photutils.segmentation import detect_sources\n",
+ "from pypolyclip import clip_multi"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "9687f0cc-e588-4cd2-a521-7975a5955cd5",
+ "metadata": {},
+ "source": [
+ "## Define Functions and Parameters"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7b97b500-97c6-4624-9d1c-81d5757b6a21",
+ "metadata": {},
+ "source": [
+ "Define a function to download a named file via the MAST API to the current directory. The function includes authentication logic, but this example uses public data, so no MAST API token is required."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a23eaad8-17ea-4a7f-8b66-51394bc82e49",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_jwst_file(name, mast_api_token=None, overwrite=False):\n",
+ " \"\"\"Retrieve a JWST data file from MAST archive.\n",
+ "\n",
+ " Parameters\n",
+ " ----------\n",
+ " name : str\n",
+ " Name of the file to download from MAST\n",
+ " \n",
+ " mast_api_token : str\n",
+ " MAST API token. Required only for proprietary data\n",
+ " \n",
+ " overwrite : bool\n",
+ " If True and the requested file already exists locally, the file will not be downloaded. IF False,\n",
+ " the file will be downloaded\n",
+ "\"\"\"\n",
+ " # If the file already exists locally, don't redownload it, unless the\n",
+ " # user has set the overwrite keyword\n",
+ " if os.path.isfile(name):\n",
+ " if not overwrite:\n",
+ " print(f'{name} already exists locally. Skipping download.')\n",
+ " return\n",
+ " else:\n",
+ " print(f'{name} exists locally. Re-downloading.')\n",
+ "\n",
+ " mast_url = \"https://mast.stsci.edu/api/v0.1/Download/file\"\n",
+ " params = dict(uri=f\"mast:JWST/product/{name}\")\n",
+ " if mast_api_token:\n",
+ " headers = dict(Authorization=f\"token {mast_api_token}\")\n",
+ " else:\n",
+ " headers = {}\n",
+ " r = requests.get(mast_url, params=params, headers=headers, stream=True)\n",
+ " r.raise_for_status()\n",
+ " with open(name, \"wb\") as fobj:\n",
+ " for chunk in r.iter_content(chunk_size=1024000):\n",
+ " fobj.write(chunk)\n",
+ " if os.path.isfile(name):\n",
+ " print(f\"{name} successfully downloaded\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1712c691-d2ea-42ef-a11f-eea80c7055f8",
+ "metadata": {},
+ "source": [
+ "Define a function that will run assign_wcs and flat fielding on an input rate file"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d3f15eeb-e5c3-4c90-a9e6-6bdbe006a9ec",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def run_pipeline_steps(filename):\n",
+ " \"\"\"Run the assign_wcs, flat field, and photom calibration steps on the given file.\n",
+ " If the file contains WFSS data, trick the pipeline to use the imaging mode flat\n",
+ " field reference file.\n",
+ " \n",
+ " Parameters\n",
+ " ----------\n",
+ " filename : str\n",
+ " Name of the input file upon which the steps will be run\n",
+ " \n",
+ " Returns\n",
+ " -------\n",
+ " filename : str\n",
+ " Name of the output file saved by the pipeline steps\n",
+ " \n",
+ " photom : jwst.datamodels.ImageModel\n",
+ " Datamodel instance containing the calibrated data\n",
+ " \"\"\"\n",
+ " assign_wcs = AssignWcsStep.call(filename)\n",
+ "\n",
+ " # In order to apply the imaging mode flat field reference file to the data,\n",
+ " # we need to trick CRDS by temporarily changing the pupil value to be CLEAR\n",
+ " reset_pupil = False\n",
+ " if 'GRISM' in assign_wcs.meta.instrument.pupil:\n",
+ " true_pupil = deepcopy(assign_wcs.meta.instrument.pupil)\n",
+ " assign_wcs.meta.instrument.pupil = 'CLEAR'\n",
+ " reset_pupil = True\n",
+ "\n",
+ " # Run the flat field step\n",
+ " flat = FlatFieldStep.call(assign_wcs, save_results=True)\n",
+ "\n",
+ " # Run the photom step to populate the name of the WFSS sensitivity \n",
+ " photom = PhotomStep.call(flat, save_results=True)\n",
+ "\n",
+ " # Set the pupil back to the original value now that flat fielding is complete\n",
+ " if reset_pupil:\n",
+ " photom.meta.instrument.pupil = true_pupil\n",
+ " photom.save(photom.meta.filename)\n",
+ " \n",
+ " # Return the name of the output file, as well as the datamodel\n",
+ " return photom.meta.filename, photom"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "462b1d4e-cb98-4493-b236-b9f4a2429850",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def show_2d_spec(data, xlim=(200, 240), ylim=(1705, 1730), vmin=0, vmax=3):\n",
+ " \"\"\"Show a 2D image with colorbar. Designed to display the 2D real and\n",
+ " simulated spectra\n",
+ "\n",
+ " Parameters\n",
+ " ----------\n",
+ " data : numpy.ndimage\n",
+ " 2D image\n",
+ " xlim : tup\n",
+ " 2-tuple of beginning and ending x-coordinates for the display\n",
+ " ylim : tup\n",
+ " 2-tuple of beginning and ending y-coordinates for the display\n",
+ " vmin : float\n",
+ " Signal corresponding to minimum display scale\n",
+ " vmax : float\n",
+ " Signal corresponding to maximum display scale\n",
+ " \"\"\"\n",
+ " fig, ax = plt.subplots(1, 1, figsize=(15, 3))\n",
+ " cax = ax.imshow(data, origin=\"lower\", aspect='auto', cmap='viridis', vmin=vmin, vmax=vmax)\n",
+ " ax.set_xlim(xlim[0], xlim[1]) # change to (0, 700) to see the entire spectrum)\n",
+ " ax.set_ylim(ylim[0], ylim[1])\n",
+ " plt.xticks(range(xlim[0], xlim[1], 5))\n",
+ " plt.xlabel(\"Dispersion coordinate (pixel)\")\n",
+ " plt.ylabel(\"Cross dispersion coordinate (pixel)\")\n",
+ " colorbar = fig.colorbar(cax, orientation='vertical', pad=0.01)\n",
+ " colorbar.ax.set_ylabel('Signal', labelpad=10, rotation=270)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f3501dc2-fd64-4383-a8c3-55ed4292b348",
+ "metadata": {},
+ "source": [
+ "## Download Data"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1d29305f-719a-4fed-bce2-032da1c9d66f",
+ "metadata": {},
+ "source": [
+ "We start with a simple pair of imaging and wfss files. These were manually selected and point at the same field on the sky, and use the same NIRCam module, channel, and cross filter."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "df6e71c6-4ccf-4e46-baf8-6da51585c70a",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# First, download the imaging and WFSS files from MAST\n",
+ "imaging_file = \"jw01076109001_02102_00001_nrcalong_cal.fits\"\n",
+ "wfss_file = \"jw01076109001_02101_00001_nrcalong_rate.fits\"\n",
+ "get_jwst_file(imaging_file)\n",
+ "get_jwst_file(wfss_file)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4e09f3f1-1ba9-4c0b-994c-b21a44b5ecd4",
+ "metadata": {},
+ "source": [
+ "## Run Pipeline Steps"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e6e32b22-40e2-4158-aea0-4d1126e470e4",
+ "metadata": {},
+ "source": [
+ "We want to assign a WCS, apply a flat-field, and flux calibrate the WFSS data. For this, we use the run_pipeline_steps() function defined above. This will call the appropriate pipeline steps. We apply the imaging mode flat field to the WFSS file since the flat fields are not wavelength dependent."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "03c23902-0d4f-439f-8bff-4f9ade6864df",
+ "metadata": {
+ "scrolled": true
+ },
+ "outputs": [],
+ "source": [
+ "# Run AssignWcsStep, FlatFieldStep, and PhotomStep on the WFSS rate file\n",
+ "wfss_file, wfss_data = run_pipeline_steps(wfss_file)\n",
+ "\n",
+ "# Extract the WFSS pixel data from the datamodel instance\n",
+ "wfss_data = wfss_data.data"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "fb9917aa-cd94-478e-b630-038f8222eff1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(f\"Pipeline-processed WFSS file is {wfss_file}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cb9657fa-5887-4b02-8d03-a812c1c7c4bb",
+ "metadata": {},
+ "source": [
+ "Read some information from the imaging data and WFSS data. We need to know which module, channel, cross filter, and grism we are looking at. We also need to find the values needed to convert the surface brightness units of the imaging cal files into units of $erg/s/cm^2/A$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e77d61b5-6f43-42ec-b8d4-d53de60bc6b5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "img_hdr0 = fits.getheader(imaging_file, 0)\n",
+ "img_hdr1 = fits.getheader(imaging_file, 1)\n",
+ "wfss_hdr = fits.getheader(wfss_file)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3ee73119-514e-448c-9d4d-014cc0352ec5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "FILTER = img_hdr0[\"FILTER\"]\n",
+ "MODULE = img_hdr0[\"MODULE\"]\n",
+ "PUPIL = img_hdr0[\"PUPIL\"]\n",
+ "PIXAR_SR = img_hdr1[\"PIXAR_SR\"]\n",
+ "print(f\"IMAGING FILTER: {FILTER}, MODULE: {MODULE}, PUPIL: {PUPIL}, Sise of pixel: {PIXAR_SR} steradians\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e8ac7611-7d83-4a65-a1d2-3ed5dabc52b3",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "FILTER = wfss_hdr[\"FILTER\"]\n",
+ "MODULE = wfss_hdr[\"MODULE\"]\n",
+ "PUPIL = wfss_hdr[\"PUPIL\"]\n",
+ "print(f\"WFSS FILTER: {FILTER}, MODULE: {MODULE}, PUPIL: {PUPIL}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "130b67f2-8b90-4ce6-848e-cd3b80e1b4ae",
+ "metadata": {},
+ "source": [
+ "Note that in this simple example, the imaging and WFSS data use the same cross filter and are of course using the same NIRCam module"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "6ecf1fea-da1c-40df-958a-ca5a7fbe3516",
+ "metadata": {},
+ "source": [
+ "Compute the conversion between pixel values in the imaging data, which are in MJy/sr, and $F_{\\lambda}$ units of $erg/s/cm^2/A$ (per pixel). Multiplying the values in our cal imaging file by this value (called PHOTFLAM) is all we will need to determine the $F_{\\lambda}$ values of each of the pixels in an object detected in our imaging data"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e72beddf-3b68-4553-aa0c-119b72d23ed5",
+ "metadata": {},
+ "source": [
+ "NIRCam filter pivot wavelengths are listd in Tables 2 and 3 on the NIRCam Filters J"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "0b813f77-6b2a-4d02-be05-9ba1126db450",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "Pivot_wavelength = 44010 # Angstroms\n",
+ "PHOTFLAM = 1e6 * PIXAR_SR / 3.3356e4 / Pivot_wavelength**2\n",
+ "print(f'PHOTFLAM is {PHOTFLAM}')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b8ca3871-5725-45d1-882e-3b9729eb231d",
+ "metadata": {},
+ "source": [
+ "## Detect Sources"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "6fe48690-846d-4dab-891e-66b57a9ce55f",
+ "metadata": {},
+ "source": [
+ "Since we want to disperse each pixel comprising a given source, we first need to get a segmentation map of all the objects in the image. We follow the general method presented in photutils' detect_sources() function documentation to create the segmentation map. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b7108ee3",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "imaging_data = fits.getdata(imaging_file)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "19ef7f11",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Create a 2D background model of the scene, and subtract\n",
+ "bkg_estimator = MedianBackground()\n",
+ "bkg = Background2D(imaging_data, (50, 50), filter_size=(21, 31), bkg_estimator=bkg_estimator)\n",
+ "imaging_data -= bkg.background "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "6f74bb26",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Use the RMS of the background to set the threshold value for source detection\n",
+ "threshold = 50 * bkg.background_rms"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c76bf196",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Convolve the image with a 2D Gaussian kernel\n",
+ "kernel = make_2dgaussian_kernel(3.0, size=5)\n",
+ "convolved_data = convolve(imaging_data, kernel)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "1f1d9dfd",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Detect the sources in the imaging data\n",
+ "segment_map = detect_sources(convolved_data, threshold, npixels=10)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f853775b-a005-4eac-8152-c40efe4a57ef",
+ "metadata": {},
+ "source": [
+ "Below we see that the segmentation map shows a large number of sources spread across the detector."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "337cecc5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Show the segmentation map\n",
+ "plt.imshow(segment_map, origin=\"lower\", vmin=1, vmax=2)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "61bda00d-ac3e-4f04-bbdf-e2897b0a34ae",
+ "metadata": {},
+ "source": [
+ "## Simulate spectrum of one source"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8a09951e-1dde-4eca-b4bf-c3f8e513f603",
+ "metadata": {},
+ "source": [
+ "Here, we show how to simulate the dispersion of only one source. In order to simulate a full WFSS observation, what we show here needs to be done for every sourece in the field. Simulating all the dispersed spectra is also a way to mask out spectra when estimating the dispersed background level during subsequent extraction and it also allows for an estimate of the amount of spectral contamination by overlapping spectra."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8058f79b-886c-4049-8be8-53e460cd7365",
+ "metadata": {},
+ "source": [
+ "For this example, we choose the source at coordinates (x,y) = (405,1465) and show how to simulate its spectrum.\n",
+ "\n",
+ "We first get its segmentation map ID and create a list of all pixels associated with the source in the imaging data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4b3bc709",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xd, yd = 405, 1465\n",
+ "ID = segment_map.data[yd, xd]\n",
+ "print(f\"Object ID is: {ID}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b121eb67-d9d5-4a90-ba68-9c38a0ee0115",
+ "metadata": {},
+ "source": [
+ "### Locate source in imaging and WFSS data"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8d711fdd-6e86-4b96-adfd-92334a535c16",
+ "metadata": {},
+ "source": [
+ "Working with the imaging mode file, create a list of coordinates for all of the source's pixels, along with their flux values (in MJy/SR)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "8b87237f",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ok = segment_map.data == ID\n",
+ "yds, xds = np.nonzero(ok)\n",
+ "cds = imaging_data[ok]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f338ef1d-2602-4f26-aa07-a156372c4506",
+ "metadata": {},
+ "source": [
+ "Show the source, as well as its segmentation map, from the imaging data. On the left we see that the source appears to be a point source. The segmentation map shows a collection of roughly 6x6 pixels that have been identified as part of this source."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f249aad6",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "min_x = np.min(xds)\n",
+ "max_x = np.max(xds)\n",
+ "min_y = np.min(yds)\n",
+ "max_y = np.max(yds)\n",
+ "\n",
+ "fig, axs = plt.subplots(1, 2, figsize=(15, 5))\n",
+ "axs[0].imshow(imaging_data[min_y:max_y + 1, min_x:max_x + 1], origin=\"lower\")\n",
+ "axs[1].imshow(segment_map.data[min_y:max_y + 1, min_x:max_x + 1], origin=\"lower\")\n",
+ "axs[0].set_title(\"Source\")\n",
+ "axs[1].set_title(\"Segmentation map\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "aa4510e5-fc21-4d35-ae74-13e2a7e09ebf",
+ "metadata": {},
+ "source": [
+ "All the information we have for this source is within the reference frame of the imaging data. But we need to know where the flux for this source is in the WFSS observation. This is handled using the gWCS of both imaging and WFSS observations. With these, we translate the location of each source pixel in the imaging data into the corresponding pixel location in the WFSS data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "88469751-28c4-4532-8822-5bdaff3f0817",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Open the imaging file using the JWST datamodels, and retrieve the WCS information\n",
+ "imaging_wcs = datamodels.open(imaging_file)\n",
+ "imaging_to_world = imaging_wcs.meta.wcs.get_transform('detector', 'world')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "febd5bcb-f7f4-4fb0-b5d4-0a468d1ba403",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Get the WCS information associated with the WFSS file\n",
+ "wfss_wcs = datamodels.open(wfss_file)\n",
+ "wfss_to_pix = wfss_wcs.meta.wcs.get_transform('world', 'detector')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e5a6fd6e-d7ca-4c44-ac3a-86c76652759b",
+ "metadata": {},
+ "source": [
+ "For reference, below we see the pixel coordinates in the imaging data. The pixels corresponding to the source range between x values of 402 and 408, and y values of 1461 and 1468."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "781ed43e-5aa5-4e68-a3f4-2e9b424b5d54",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.scatter(xds, yds)\n",
+ "plt.xlabel(\"Imaging columns\")\n",
+ "plt.ylabel(\"Imaging rows\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ca5e3d4f-476e-4dfb-83ed-6602552af95e",
+ "metadata": {},
+ "source": [
+ "Compute the R.A. and Dec of each of the source's imaging mode input pixels"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b2bc8e3d-ea0e-45df-9f0f-8816688d7ef7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ras, decs = imaging_to_world(xds, yds)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "91cb852a-b3e7-44f0-9079-3452abe9cb0e",
+ "metadata": {},
+ "source": [
+ "Now compute the pixel coordinates of where the **undispersed** source would be in the WFSS data. In this case the gWCS requires an input wavelength and spectral order. These same values are returned by the translation function. In this case we choose a wavelength of 3.56 microns and a spectral order of 1. The pixel coordinates returned in this case are not dependent upon wavelength or spectral order, so any values can be used."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "6ca0aec5-0927-4a08-bba2-25a1311e9eb7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xs, ys, xxx, yyy = wfss_to_pix(ras, decs, 3.56, 1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "9f489e69-237c-4aab-9d50-d817619918b7",
+ "metadata": {},
+ "source": [
+ "Show the pixel location of the **undispersed** source in the WFSS data. Note that the location is significantly different than that in the imaging data. In this case, the source is located between x values of 117 to 123, and y values of 1740 and 1746."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b012a3be-bfae-4b72-a253-5c1177baba47",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.scatter(xs, ys)\n",
+ "plt.xlabel(\"WFSS columns\")\n",
+ "plt.ylabel(\"WFSS rows\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "79ebb828-c5a7-4b5c-a540-8ac4fdaf9bd1",
+ "metadata": {},
+ "source": [
+ "### Get wavelength information"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "df18225c-a3ab-42b6-bcf4-0573cbb48147",
+ "metadata": {},
+ "source": [
+ "When simulating this dispersed spectrum, we need to consider the wavelength of the light that is being dispersed. So each of the pixels above will be numericaly dispersed across a range of discrete wavelengths. \n",
+ "\n",
+ "To get the needed wavelength information, we initialize a grismconf Config object. This contains the information and polynomials describing the dispersion of the disperser as well as the corresponding inverse sensitivity curve."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "1ddfcad5-eb41-4692-acfc-b12916083b8e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "C = grismconf.Config(wfss_file)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "26dc8950-cc75-45c1-8427-9c2e5b40c3d9",
+ "metadata": {},
+ "source": [
+ "Show the inverse sentivity, which includes the wavelength range and shape of the sensitivity. This is defined in units of DN/s per $F_{\\lambda}$ ($erg/s/cm^2/A$). The curve shows significant sensitivity between about 3.85 microns, and 5.05 microns."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "10891802-5e3c-46be-9d0e-b347503fbbd4",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.plot(C.SENS_data[\"+1\"][0], C.SENS_data[\"+1\"][1])\n",
+ "plt.grid()\n",
+ "plt.xlabel(\"Wavelength (micron)\")\n",
+ "plt.ylabel(r\"DN/s per $erg/s/cm^2/A$\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7e95f63d-5a27-46f6-ba16-2319f9a57c13",
+ "metadata": {},
+ "source": [
+ "We use the grimconf configuration object to quickly get the wavelength range that corresponds to the disperser. This is present in the WRANGE attribute."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "9b4bf6bd-300a-4ab7-92cc-0a636b479e92",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "wmin = C.WRANGE[\"+1\"][0]\n",
+ "wmax = C.WRANGE[\"+1\"][1]\n",
+ "\n",
+ "print(f\"The wavelength range to consider is {wmin} to {wmax} microns\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "aa61092f-28c5-4a9a-93a5-5e90652cf346",
+ "metadata": {},
+ "source": [
+ "Compute the dispersion in units of wavelength per pixel. Grismconf can give us the derivative of the dispersion in units of wavelength as well as in units of pixels with respect to the t parameter. Details of the $t$ parameter are given in the Box Extraction Notebook. As noted there, $t$ is a normalized parameter, whre values of $t = 0$ and $t = 1$ correspond to the blue and red edges of a dispersed spectrum.\n",
+ "\n",
+ "While the dispersion varies slightly across the area of the detector covered by a spectrum, we use the dispersion at a $t$ value of 0.5, which corresponds to the middle of the spectral range."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b7da4676-c269-41d8-ab4e-ed3db2754a8d",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "dlam = C.DDISPL(\"+1\", 1000, 1000, 0.5) / C.DDISPX(\"+1\", 1000, 1000, 0.5)\n",
+ "print(f\"Dispersion is {dlam * 10000} Angstroms per pixel\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1f2a7e3f-e9a0-4abf-af05-1c5e8b2b4dc7",
+ "metadata": {},
+ "source": [
+ "Next, create an array of wavelength values at which the simulated spectrum will be calculated. For this, we must pick a wavelength step. Ideally this step should be smaller than the native dispersion of the grism. We therefore pick half of the dispersion value calculated above."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "16875987",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "dlam = dlam / 2\n",
+ "lams = np.arange(wmin, wmax, dlam)\n",
+ "\n",
+ "print(f\"We are using {len(lams)} values of wavelength\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "78719155-0e60-4d11-a1c6-10274e51969e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(f'First and last elements of the wavelength array: {\"%.5f\" % lams[0]} microns, {\"%.5f\" % lams[-1]} microns')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e7aa62c1-758e-45eb-aee8-b89b3ffd0f16",
+ "metadata": {},
+ "source": [
+ "### Simulate the dispersion of a single pixel"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b1b5c887-8adb-4a12-8b47-7099f71cd2b5",
+ "metadata": {},
+ "source": [
+ "With the wavelength information in hand, we can create a simulated dispersion of each object pixel.\n",
+ "\n",
+ "Below, we show the process for a single pixel. We choose a pixel relatively close to the center of the source."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d64845db-f4ee-41c0-a9d9-35d26f3b4fa9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "i = 22\n",
+ "print(f\"Use pixel (x, y) = ({xs[i]}, {ys[i]}) for single pixel dispersion\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "bc70a326-c437-4ca2-954d-406b890d1c6e",
+ "metadata": {},
+ "source": [
+ "We start by computing the $t$ values corresponding to the wavelengths (lams) we are considering. Refer to the Box Extraction notenook for additional background details on the $t$ value."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "02ea112e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ts = C.INVDISPL(\"+1\", xs[i], ys[i], lams)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ada59693-4d4e-4bd9-975e-7c3956e927ea",
+ "metadata": {},
+ "source": [
+ "Next, we create an array of polygons representing the locations of the dispersed signal coming from our selected input pixel.\n",
+ "\n",
+ "The cell below computes the array of coordinates, in the WFSS observation, of the bottom left corner of our pixel."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d6ef5fa4",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xgsA = C.DISPX(\"+1\", xs[i], ys[i], ts) + xs[i]\n",
+ "ygsA = C.DISPY(\"+1\", xs[i], ys[i], ts) + ys[i]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2f729847-2f83-411b-8afb-1bdebb9db2d2",
+ "metadata": {},
+ "source": [
+ "The following three cells compute the locations of the other three corners:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e4740945",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xgsB = C.DISPX(\"+1\", xs[i] + 1, ys[i], ts) + xs[i] + 1\n",
+ "ygsB = C.DISPY(\"+1\", xs[i] + 1, ys[i], ts) + ys[i]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b59e1d1c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xgsC = C.DISPX(\"+1\", xs[i] + 1, ys[i] + 1, ts) + xs[i] + 1\n",
+ "ygsC = C.DISPY(\"+1\", xs[i] + 1, ys[i] + 1, ts) + ys[i] + 1"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "0b0ab3d3",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xgsD = C.DISPX(\"+1\", xs[i], ys[i] + 1, ts) + xs[i]\n",
+ "ygsD = C.DISPY(\"+1\", xs[i], ys[i] + 1, ts) + ys[i] + 1"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "46c060d3-ab2a-4c8c-abd1-995be9910fed",
+ "metadata": {},
+ "source": [
+ "Re-organize things a little to contain a list of polygon corners, which are used by the pypolyclip module to compute their overlap with the pixel coordinates of the WFSS observation. While were are looking at a single input source pixel, we are computing this at many different wavelength values so the resultant is a list of many pixels/polygons to project onto our WFSS rectilinear pixel grid."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "79d009b7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "pxs = [[xgsA[ii], xgsB[ii], xgsC[ii], xgsD[ii]] for ii in range(len(xgsA))]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "7d3a1dd5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "pys = [[ygsA[ii], ygsB[ii], ygsC[ii], ygsD[ii]] for ii in range(len(ygsA))]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "5690d67f-1715-44ba-8b30-9a3f3641ff1c",
+ "metadata": {},
+ "source": [
+ "Below we create a figure showing the resulting locations of dispersed pixels. We dispersed the single input pixel using an array of wavelengths that is oversampled by a factor of 2 relative to the native dispersion of the grism. That wavelength array is translated into an array of pixel positions which overlay the WFSS detector grid. \n",
+ "\n",
+ "For clarity, we zoom in on a 40 pixel wide area in the dispersion direction. This shows that our single imaging mode pixel will be dispersed along a nearly horizontal line (seen as the rainbow colored boxes) in this area. Zooming out on this plot, by changing the xlim values in the cell below, will show the entire set of pixels corresponding to the dispersed input pixel. For consistency, we show similar 40 pixel wide plots for the other figures in the next two sections."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "94e5dc02-7d55-42f7-90bc-f9f40fd9750a",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(1, 1, figsize=(15, 3))\n",
+ "for i in tqdm.tqdm(range(len(pxs))):\n",
+ " tx = pxs[i]\n",
+ " tx.append(pxs[i][0])\n",
+ " ty = pys[i]\n",
+ " ty.append(pys[i][0])\n",
+ " plt.plot(tx, ty)\n",
+ "\n",
+ "plt.xticks(range(0, len(pxs)))\n",
+ "\n",
+ "plt.xlim(200, 240) # change to (0, 700) to see the entire spectrum)\n",
+ "plt.xlabel(\"WFSS columns\")\n",
+ "plt.ylabel(\"WFSS Rows\")\n",
+ "plt.grid()\n",
+ "plt.xlabel(\"Dispersion coordinate (pixel)\")\n",
+ "plt.ylabel(\"Cross ispersion coordinate (pixel)\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "fb5bce23-369d-4fb5-98a9-228723c4017c",
+ "metadata": {},
+ "source": [
+ "We can now use the pypolyclip.clip_multi to compute how much of each dispersed pixel (colored boxes above) falls onto each pixel in the WFSS image pixel grid (shown as the gray grid above). Details about this is available on the pypolyclip page."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f5793e51",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xc, yc, area, slices = clip_multi(pxs, pys, [2048, 2048])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cd9659b1-be21-4e66-9e00-69bcd61267a3",
+ "metadata": {},
+ "source": [
+ "In the following section, we will repeat this workflow on all pixels, and use the area information to scale the flux values for all output pixels."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f1d4297c-9f94-4a96-a314-d22fb39b41a5",
+ "metadata": {},
+ "source": [
+ "### Disperse all the pixels for our source"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1eaff90b-c69a-41d8-8740-e3b28da8939e",
+ "metadata": {},
+ "source": [
+ "Note that the figure above shows a single source pixel being dispersed. For a full source, each of the input source pixels should be similarly dispersed, resulting in multiple dispersed pixels contributing to the final counts in each of the detector pixels in the WFSS image. We will also need to compute and attribute the proper flux, in DN/s, to each of the WFSS detector pixels from each input pixel. This is done below for our selected object.\n",
+ "\n",
+ "We must keep track of all the information such as the wavelength and fraction of the original imaging pixel flux that falls onto the WFSS simulated pixel array for each imaging pixel."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "800f8c58",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xcs = []\n",
+ "ycs = []\n",
+ "alams = []\n",
+ "flams = []\n",
+ "\n",
+ "all_pxs = []\n",
+ "all_pys = []\n",
+ "all_flams = []\n",
+ "all_counts = []\n",
+ "\n",
+ "# Loop over the input source pixels in the WFSS reference frame.\n",
+ "for i in tqdm.tqdm(range(len(xs))):\n",
+ " # Use the imaging flux in each of these pixels to compute the input DN/s and flam units\n",
+ " counts = cds[i]\n",
+ " flam = counts * PHOTFLAM\n",
+ "\n",
+ " # Disperse this pixel using len(lams) wavelength. This results in len(lams) projected\n",
+ " # pixels contributing to the final WFSS data\n",
+ " ts = C.INVDISPL(\"+1\", xs[i], ys[i], lams)\n",
+ " xgsA = C.DISPX(\"+1\", xs[i], ys[i], ts) + xs[i]\n",
+ " ygsA = C.DISPY(\"+1\", xs[i], ys[i], ts) + ys[i]\n",
+ " xgsB = C.DISPX(\"+1\", xs[i] + 1, ys[i], ts) + xs[i] + 1\n",
+ " ygsB = C.DISPY(\"+1\", xs[i] + 1, ys[i], ts) + ys[i]\n",
+ " xgsC = C.DISPX(\"+1\", xs[i] + 1, ys[i] + 1, ts) + xs[i] + 1\n",
+ " ygsC = C.DISPY(\"+1\", xs[i] + 1, ys[i] + 1, ts) + ys[i] + 1\n",
+ " xgsD = C.DISPX(\"+1\", xs[i], ys[i]+1, ts) + xs[i]\n",
+ " ygsD = C.DISPY(\"+1\", xs[i], ys[i]+1, ts) + ys[i] + 1\n",
+ "\n",
+ " # Use the corners of the dispersed pixels, and compute the WFSS pixels which they\n",
+ " # overlap, and by how much\n",
+ " pxs = [[xgsA[ii], xgsB[ii], xgsC[ii], xgsD[ii]] for ii in range(len(xgsA))]\n",
+ " pys = [[ygsA[ii], ygsB[ii], ygsC[ii], ygsD[ii]] for ii in range(len(ygsA))]\n",
+ " xc, yc, area, slices = clip_multi(pxs, pys, [2048, 2048])\n",
+ "\n",
+ " # Book keeping to track the wavelength of each of the areas being projected into\n",
+ " # the WFSS pixel grid\n",
+ " tlams = np.zeros(len(xc))\n",
+ " for i in range(len(slices)):\n",
+ " tlams[slices[i]] = lams[i]\n",
+ "\n",
+ " # Store the flux, wavelength, and where they should end up on the WFSS pixel grid.\n",
+ " # Note the values in xcs and ycs are not unique\n",
+ " xcs.extend(xc.tolist())\n",
+ " ycs.extend(yc.tolist())\n",
+ " flams.extend((flam * area).tolist())\n",
+ " alams.extend(tlams.tolist())\n",
+ "\n",
+ " # Save for plotting later. Only used for plot below.\n",
+ " all_pxs.append(pxs)\n",
+ " all_pys.append(pys)\n",
+ " all_flams.append(flam)\n",
+ " all_counts.append(flam * C.SENS[\"+1\"](tlams) * dlam * 10000)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "249a3c5d-c7f1-4aa7-b5cb-e7426eae762b",
+ "metadata": {},
+ "source": [
+ "At this point, we have a list of WFSS pixels (xcs, ycs), the flux falling on these pixels (flams, in $F_{\\lambda}$ units), and the wavelength of the light contained in them (alams).\n",
+ "In our simulation, we do not want to project flux units but rather DN/s, so we convert the input flams values into DN/s (using the reverse relation we used in the Box Extraction notebook when we performed the inverse operation to convert extracted DN/s into $F_{\\lambda}$ flux units)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3edd128c-3fda-44bb-9c48-2fa414bac8e2",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Note: the factor of 10000 below accounts for dlam being in microns, while we\n",
+ "# want Angstroms since the inverse sensitivity is defined per Angstrom.\n",
+ "s = C.SENS[\"+1\"](alams)\n",
+ "counts = flams * s * dlam * 10000"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "67b1241c-4c84-4608-9b5c-9868693f3a90",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(f\"There are {len(counts)} dispersed bits of pixels to combine into a final WFSS pixel grid\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "76d9941e-e949-4af3-8f0a-6066dea3f7ff",
+ "metadata": {},
+ "source": [
+ "We now have a large list of DN/s values and where they should be added onto our simulated WFSS observationn in order to simulate the full dispersed spectrum of our source.\n",
+ "There are duplicate entries in the (xcs,ycs) coordinate list as different wavelengths get mixed by the object's \"self-contamination\"."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8c99799e-2afe-4fff-9d0f-86ff7d0f9e7d",
+ "metadata": {},
+ "source": [
+ "The following plot shows the dispersed input pixels, using blue outlines, projected onto the final WFSS pixels, which are shown as the gray grid. The dispersed pixels are shaded in black proportionally to their flux (in DN/s)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "34c89144-6321-4b99-b59f-295b3819173b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(1, 1, figsize=(15, 3))\n",
+ "for i in tqdm.tqdm(range(len(all_pxs))):\n",
+ " for j in range(len(all_pxs[i])):\n",
+ " \n",
+ " tx = all_pxs[i][j][:]\n",
+ " tx.append(tx[0])\n",
+ " ty = all_pys[i][j][:]\n",
+ " ty.append(ty[0])\n",
+ " plt.plot(tx, ty, color='b', alpha=0.02)\n",
+ " c = all_counts[i]\n",
+ " c[c < 0] = 0\n",
+ " plt.fill(tx, ty, color='k', alpha=c[j])\n",
+ "plt.grid()\n",
+ "plt.xlim(200, 240) # change to (0, 700) to see the entire spectrum)\n",
+ "plt.xticks(range(200, 240))\n",
+ "plt.xlabel(\"Dispersion coordinate (pixel)\")\n",
+ "plt.ylabel(\"Cross dispersion coordinate (pixel)\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cf6eb1dc-49e6-4d24-8d6b-745772f5a0aa",
+ "metadata": {},
+ "source": [
+ "To quickly combine all of these counts at each of their WFSS pixel location, we can use scipy.coo_matrix which is fast and efficient:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5a67abeb-e623-4509-9eb3-a0c1c5db954a",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "xcs = np.array(xcs)\n",
+ "ycs = np.array(ycs)\n",
+ "\n",
+ "# Ignore counts and coordinates that are outside of the detector\n",
+ "ok = (xcs >= 0) & (xcs < 2048) & (ycs >= 0) & (ycs < 2048) \n",
+ "simulated = coo_matrix((counts[ok], (ycs[ok], xcs[ok])), shape=(2048, 2048)).toarray()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ed4dd1c1-4416-459b-86c2-1bca5cb45d15",
+ "metadata": {},
+ "source": [
+ "Show the resulting 2D simulated spectrum for this source. We see a nearly horizontal trace across the field of view. Increasing the x range of the plot to (0, 700) will reveal the full trace."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "8393033d-f186-4c46-a825-5f873cfc17c9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Change xlim to (0, 700) to see the entire spectrum)\n",
+ "show_2d_spec(simulated, xlim=(200, 241), ylim=(1705, 1730), vmin=0, vmax=4)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "3600044a-778a-459d-b748-ff93f269723c",
+ "metadata": {},
+ "source": [
+ "Show the real data at the same scale. This trace also appears as a nearly horizontal line. The signal level in the real data is slightly higher than that of the simulated data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c5aa8e40-ecae-4fee-801e-a79ceba0ac3c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Change xlim to (0, 700) to see the entire spectrum)\n",
+ "show_2d_spec(wfss_data, xlim=(200, 241), ylim=(1705, 1730), vmin=0, vmax=4)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cd7d09ae-a8de-41e5-beab-f625eba05343",
+ "metadata": {},
+ "source": [
+ "A good way to check our simulated image is to subtract the simulation from our data. The figure below shows the simulated minus real difference image. The difference in peak level of the trace is visible. The alignment of the simulated and real traces is also apparent as the width of the difference is constant from left to right."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "72a711b7-22ab-4419-bc40-733c2ad88a50",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Change xlim to (0, 700) to see the entire spectrum)\n",
+ "show_2d_spec(wfss_data - simulated, xlim=(200, 241), ylim=(1705, 1730), vmin=0, vmax=4)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "790e3fc8-db68-487a-a87c-26ed7cd65d5c",
+ "metadata": {},
+ "source": [
+ "Sum the simulated spectrum and also the real data in the dispersion direction to check if the trace profiles are well aligned. The figure below shows the summed real data (orange) and simulation (blue line). The left/right alignment of the two peaks shows that the location of the simulated trace is very close to that of the real trace."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "53cbe509",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.plot(np.nansum(simulated, axis=-1), color='blue')\n",
+ "plt.plot(np.nansum(wfss_data, axis=-1), color='orange')\n",
+ "plt.xlim(1705, 1730)\n",
+ "plt.ylim(0, 3000)\n",
+ "plt.grid()\n",
+ "plt.ylabel(\"Summed Counts (DN/s\")\n",
+ "plt.xlabel(\"Cross dispersion coordinate (pixel)\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2f0f47dd-e5c9-430b-b72b-46b285b1315e",
+ "metadata": {},
+ "source": [
+ "While the simulation appears to be astrometrically correct, we have not accounted for the dispersed background, causing the signal in the simulation to be too low.\n",
+ "This should be done using a model of the dispersed background but here, for simplicity and because we are looking at a region where we know the dispersed background is relatively featureless and flat, we can just use the simulation to create a mask and then compute the background level per pixel.\n",
+ "Below we mask pixels in the simulated image above a signal level of 0.001 DN/s. This masks the trace. The remaining pixels, visible below in black, will be used to calculate the background level."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e826f429-d74f-4d10-bcf2-64cc56ea0c32",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Masked pixels, set to NaN, appear white in the display\n",
+ "mask = simulated > 0.0001\n",
+ "tmp = simulated * 1.\n",
+ "tmp[mask] = np.nan\n",
+ "# Change xlim to (0, 700) to see the entire spectrum)\n",
+ "show_2d_spec(tmp, xlim=(200, 241), ylim=(1705, 1730), vmin=0, vmax=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8c109bc6-11cc-481f-9c42-311662f4115f",
+ "metadata": {},
+ "source": [
+ "Compute the median dispersed background level."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d53aeb0c-7e12-4762-8bd2-a6db7085894c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bck_level = np.nanmedian(wfss_data[~mask])\n",
+ "print(f\"Background extimated to be {bck_level} DN/s per pixel\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2a930c02-dff4-46ee-9f8b-be3d10e4336a",
+ "metadata": {},
+ "source": [
+ "Now plot the collapsed simulated data (blue) and background-subtracted real data (orange). The well-aligned peaks show more clearly that the simulated and real traces are well-aligned. The peak of the simulated data is lower than that of the real data, showing that the simulation slightly under-estimates the signal in the real trace."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d9347308-734f-4873-85cf-bf2dc8a7db94",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.plot(np.nansum(wfss_data - bck_level, axis=-1), color='orange')\n",
+ "plt.xlim(1705, 1730)\n",
+ "plt.ylim(-100, 2500)\n",
+ "plt.plot(np.nansum(simulated, axis=-1), color='blue')\n",
+ "plt.grid()\n",
+ "plt.ylabel(\"Summed Counts (DN/s\")\n",
+ "plt.xlabel(\"Cross dispersion coordinate (pixel)\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7cc1c7c2-9494-41b5-89b3-9bbf62a75a71",
+ "metadata": {},
+ "source": [
+ "After subtracting the background, the simulated spectrum is a good match to the real data. The 2D simulated spectrum can now be used to subtract contamination of overlapping spectra, as well as to mask dispersed spectral traces when trying to estimate the background."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "d5425e23-86f2-40ac-ad17-4a89c699b53a",
+ "metadata": {},
+ "source": [
+ ""
+ ]
+ }
+ ],
+ "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.11.10"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/NIRCam/NIRCam_WFSS_simulating_spectra/requirements.txt b/notebooks/NIRCam/NIRCam_WFSS_simulating_spectra/requirements.txt
new file mode 100644
index 000000000..1198bfa9e
--- /dev/null
+++ b/notebooks/NIRCam/NIRCam_WFSS_simulating_spectra/requirements.txt
@@ -0,0 +1,11 @@
+astropy >= 6.1.3
+grismconf >= 1.51
+jupyter >= 1.1.1
+jwst >= 1.16.0
+matplotlib >= 3.9.2
+numpy == 1.26.4
+photutils >= 2.0.2
+pypolyclip >= 1.1.0
+requests >= 2.32.3
+scipy >= 1.14.1
+tqdm >= 4.67.1