<jupyter_start><jupyter_text>Exploring the TSGRISM mode for NIRCAM Time Series Grism observations which can be done with subarrays  Make sure that you have set the JWST_NOTEBOOK_DATA environment variable in the terminal from which you started Jupyter Notebook.The data will be read from that directory, and the pipeline should write to the current working directory, avoiding clobbers.If you would like to use your own data just substitute the locations below.<jupyter_code>notebook_dir = os.environ['JWST_NOTEBOOK_DATA']
nircam_data = notebook_dir + 'nircam/'
# plotting, the inline must come before the matplotlib import
%matplotlib inline
from matplotlib import pyplot as plt
import matplotlib.patches as patches

params = {'legend.fontsize': 6,
          'figure.figsize': (8, 8),
          'figure.dpi': 150,
         'axes.labelsize': 6,
         'axes.titlesize': 6,
         'xtick.labelsize':6,
         'ytick.labelsize':6}
plt.rcParams.update(params)


# python general
import os
import sys
import numpy as np

# astropy modules
from astropy.io import fits
from astropy.table import QTable
from astropy.wcs.utils import skycoord_to_pixel
import photutils

# jwst 
import jwst
from jwst.datamodels import image, CubeModel
from jwst.assign_wcs import nircam

print("Using jwst version: {}".format(jwst.__version__))
print(sys.version)
tsgrism_image=nircam_data + 'jw12345001001_01101_00001_nrcalong_rateints.fits'
gim=CubeModel(tsgrism_image)
gim.meta.instrument.pupil, gim.meta.instrument.filter,gim.meta.exposure.type, gim.meta.instrument.module
gim.shape
gim.wavelength
gim.meta.wcsinfo.crpix1, gim.meta.wcsinfo.crpix2, gim.meta.wcsinfo.crval1, gim.meta.wcsinfo.crval2, gim.meta.wcsinfo.v2_ref, gim.meta.wcsinfo.v3_ref, gim.meta.wcsinfo.v3yangle, gim.meta.wcsinfo.roll_ref<jupyter_output><empty_output><jupyter_text>Take a look at the image we have<jupyter_code>ys,xs=gim.shape[1:]
fig = plt.figure(figsize=(8,8), dpi=150)
ax = fig.add_subplot(1, 2, 1)
ax.set_title(tsgrism_image.split("/")[-1]+"[0]", fontsize=8)
ax.imshow(gim.data[0], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

ax2 = fig.add_subplot(1, 2, 2)
ax2.set_title(tsgrism_image.split("/")[-1]+"[1]", fontsize=8)
ax2.imshow(gim.data[1], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

fig.tight_layout()<jupyter_output><empty_output><jupyter_text>Get the reference files for use with our image<jupyter_code>step=assign_wcs.AssignWcsStep()
distortion=step.get_reference_file(gim,'distortion')
specwcs=step.get_reference_file(gim, 'specwcs')
reference_file_names = {'camera': 'N/A',
 'collimator': 'N/A',
 'disperser': 'N/A',
 'distortion': distortion,
 'filteroffset': 'N/A',
 'fore': 'N/A',
 'fpa': 'N/A',
 'ifufore': 'N/A',
 'ifupost': 'N/A',
 'ifuslicer': 'N/A',
 'msa': 'N/A',
 'ote': 'N/A',
 'regions': 'N/A',
 'specwcs': specwcs,
 'v2v3': 'N/A',
 'wavelengthrange': 'N/A'}
reference_file_names<jupyter_output><empty_output><jupyter_text>We need to do the subarray transform on either side of the grism trace. This is blurred in the grism() assign_wcs because the regular grism mode doesn't allow subarrays.<jupyter_code>print(gim.meta.subarray.name)
tsgrism_pipeline=nircam.tsgrism(gim,reference_file_names)  # create the wcs pipeline
tsgrism_pipeline
from gwcs.wcs import WCS
tswcs=WCS(tsgrism_pipeline)<jupyter_output><empty_output><jupyter_text>This will detail the available transformation frames<jupyter_code>tswcs.available_frames<jupyter_output><empty_output><jupyter_text>In grism time-series mode, the Module A Grism R is used to disperse the target's spectrum along (parallel to) detector rows. The grism is used in conjunction with one of 4 wide filters in the long wavelength channel (2.4–5.0 µm): F277W, F322W2, F356W, and F444W.<jupyter_code>tswcs  # the three inputs are x, y, order<jupyter_output><empty_output><jupyter_text>Translate pixel 100,100 in the tso grism image for source location crpix1, crpix2, and order 1<jupyter_code>tswcs(100,100,1)  # returns ra, dec, wave, order
gim.get_fits_wcs()
#gim.meta.wcsinfo.crval1, gim.meta.wcsinfo.crval2, gim.meta.wcsinfo.crpix1, gim.meta.wcsinfo.crpix2<jupyter_output><empty_output><jupyter_text>Change the location. The wavelength should change, but the ra, dec should be the same<jupyter_code>tswcs(200,200,1)<jupyter_output><empty_output><jupyter_text>This should return the same x value and CRPIX2 for y, since the transforms for NIRCAM return 0 for GRISMR in the y so that the extraction box can be chosen around that<jupyter_code>tswcs.invert(3.2472519999724887, 1.0)<jupyter_output><empty_output><jupyter_text>This should return error about and invalid order since order 3 is not available<jupyter_code>tswcs(100,100,3)
tswcs.input_frame, tswcs.output_frame
tswcs.get_transform('world','v2v3')
tswcs.get_transform('v2v3', 'world')<jupyter_output><empty_output><jupyter_text>Now go back to the image and see if we can get extract_2d to go to the correct place<jupyter_code>ys,xs=gim.shape[1:]
fig = plt.figure(figsize=(8,8), dpi=150)
ax = fig.add_subplot(1, 1, 1)
ax.set_adjustable('box-forced')
ax.set_title(tsgrism_image.split("/")[-1]+"[0]", fontsize=8)
ax.imshow(gim.data[0], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

fig.tight_layout()
from jwst.extract_2d import extract_2d_step, extract_2d
step=extract_2d_step.Extract2dStep()
reference_file_names = {'camera': 'N/A',
 'collimator': 'N/A',
 'disperser': 'N/A',
 'distortion': step.get_reference_file(gim,'distortion'),
 'filteroffset': 'N/A',
 'fore': 'N/A',
 'fpa': 'N/A',
 'ifufore': 'N/A',
 'ifupost': 'N/A',
 'ifuslicer': 'N/A',
 'msa': 'N/A',
 'ote': 'N/A',
 'regions': 'N/A',
 'specwcs':  step.get_reference_file(gim,'specwcs'),
 'v2v3': 'N/A',
 'wavelengthrange': step.get_reference_file(gim, 'wavelengthrange')}
reference_file_names<jupyter_output><empty_output><jupyter_text>Set the wavelengthrange reference file to the one we updated, but hasn't yet been accepted to CRDS<jupyter_code># run assign_wcs step on the image to attach the gwcs object we created
tso_wcs_assigned = assign_wcs.assign_wcs.load_wcs(gim, reference_files=reference_file_names)
tso_wcs_assigned.meta.instrument.filter
from jwst.datamodels import WavelengthrangeModel
wrm = WavelengthrangeModel(reference_file_names['wavelengthrange'])
print(wrm.wavelengthrange)
print()
print(wrm.waverange_selector)
import asdf
test_read=asdf.open(reference_file_names['wavelengthrange'])
test_read.tree
with WavelengthrangeModel(reference_file_names['wavelengthrange']) as f:
    if (f.meta.instrument.name != 'NIRCAM'):
        raise ValueError("Wavelengthrange reference file not for NIRCAM!")
    wavelengthrange = f.wavelengthrange
    waverange_selector = f.waverange_selector
    orders = f.order
    extract_orders = f.extract_orders
    print(f.meta.instrument.name)

print(wavelengthrange)
range_select = [(x[2], x[3]) for x in wavelengthrange if (x[0] == 1 and x[1] == 'F444W')]
lmin, lmax = range_select.pop()
print(lmin, lmax)
from jwst.datamodels import SpecwcsModel
print(reference_file_names['specwcs'])
spec=SpecwcsModel(reference_file_names['specwcs'])

spec.instance
x2d = extract_2d.extract_tso_object(tso_wcs_assigned, reference_files=reference_file_names, extract_height=64)
x2d.instance
x2d.data.shape
x2d.meta.wcs
x2d.meta.wcs(887, 35)
x2d.meta.wcs.invert(4.586513215999723, 1.0)
x2d.meta.wcsinfo.instance
x2d.ysize, x2d.xsize, x2d.shape, gim.shape
x2d.wavelength
ys,xs=gim.shape[1:]
fig = plt.figure(figsize=(10,10), dpi=150)
ax3 = fig.add_subplot(1, 3, 1)
ax3.set_title(tsgrism_image.split("/")[-1]+"[0]", fontsize=8)
ax3.imshow(gim.data[0], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

zs, ys, xs = x2d.data.shape
fig = plt.figure(figsize=(10,10), dpi=150)
ax = fig.add_subplot(1, 3, 2)
xpos, ypos = x2d.source_xpos, x2d.source_ypos
title = x2d.meta.instrument.filter+" order {0}\nx={1} y={2}".format(x2d.meta.wcsinfo.spectral_order,
                                                                    xpos,
                                                                    ypos)
ax.set_title(title, fontsize=8)
ax.imshow(x2d.data[0,:,:], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)


ax2.set_title(title, fontsize=8)
ax2.imshow(x2d.data[0,:,:], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

fig.tight_layout()
tso_wcs_assigned.data.shape
x2d.data.shape
x2d.source_xpos, x2d.source_ypos
x2d.xsize, x2d.ysize<jupyter_output><empty_output><jupyter_text>We should also be able to call the extraction and override the extract_orders for filter, this should produce a SlitModel output for just the order specified<jupyter_code>x2d_single = extract_2d.extract_tso_object(tso_wcs_assigned, reference_files=reference_file_names, extract_orders=[1])
x2d_single.meta.model_type  # should be SlitModel this time
x2d_single.data.shape # should exist and be 3D
x2d_single.meta.wcsinfo.spectral_order
x2d_single.meta.wcs
ys,xs=gim.shape[1:]
fig = plt.figure(figsize=(10,10), dpi=150)
ax3 = fig.add_subplot(1, 3, 1)
ax3.set_title(tsgrism_image.split("/")[-1]+"[0]", fontsize=8)
ax3.imshow(gim.data[0], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

zs, ys, xs = x2d_single.data.shape
fig = plt.figure(figsize=(10,10), dpi=150)
ax = fig.add_subplot(1, 3, 2)
xpos, ypos = x2d_single.source_xpos, x2d_single.source_ypos
title = x2d.meta.instrument.filter+" order {0}\nx={1} y={2}".format(x2d_single.meta.wcsinfo.spectral_order,
                                                                    xpos,
                                                                    ypos)
ax.set_title(title, fontsize=8)
ax.imshow(x2d_single.data[0,:,:], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)<jupyter_output><empty_output><jupyter_text>Check that height can be specified<jupyter_code>from jwst.extract_2d.grisms import extract_tso_object<jupyter_output><empty_output><jupyter_text>The source should always be returned at pixel 34, even if the extract height doesn't allow for that.<jupyter_code>x2d_single = extract_tso_object(tso_wcs_assigned, reference_files=reference_file_names, extract_orders=[1], extract_height=50)
x2d_single.meta.model_type  # should be SlitModel this time
x2d_single.data.shape # should exist and be 3D
x2d_single.meta.wcsinfo.spectral_order, x2d_single.ysize, x2d_single.xsize
ys,xs=gim.shape[1:]
fig = plt.figure(figsize=(10,10), dpi=150)
ax3 = fig.add_subplot(1, 3, 1)
ax3.set_title(tsgrism_image.split("/")[-1]+"[0]", fontsize=8)
ax3.imshow(gim.data[0], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

zs, ys, xs = x2d_single.data.shape
fig = plt.figure(figsize=(10,10), dpi=150)
ax = fig.add_subplot(1, 3, 2)
xpos, ypos = x2d_single.source_xpos, x2d_single.source_ypos
title = x2d.meta.instrument.filter+" order {0}\nx={1} y={2}".format(x2d_single.meta.wcsinfo.spectral_order,
                                                                    xpos,
                                                                    ypos)
ax.set_title(title, fontsize=8)
ax.imshow(x2d_single.data[0,:,:], origin='lower', extent=[0,xs,0,ys], vmin=-3, vmax=3)

x2d.write('blah.fits')  # extract2d saves a SlitModel
from jwst import datamodels
blah=datamodels.open('blah.fits')
blah  # should have been read in as a SlitModel<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13474734872114785, 'fraction_numerical': 0.03538009090098922, 'mean_word_length': 5.1897407611693325, 'pattern_counts': {'":': 0, '<': 61, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '620451', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 4492, 'n_tokens_neox': 4314, 'n_words': 868, 'path': 'nircam/NIRCAM-TSGRISM-mode.ipynb', 'repo_name': 'jaytmiller/jwst-pipeline-notebooks', 'size': '11221'}