JP-3595: handling for NIRCam GRISM Time Series Pointing offset SR

CHANGES.rst

@@ -15,6 +15,8 @@ assign_wcs
   of pixel row rather than a rotation of the pixel coordinates.  The practical impact is
   to ensure that iso-lambda is along pixel rows after this change. [#8411]
 
+- Fixed test for NIRCam TSGRISM mode. [#8449]
+
 - Move the assigned source position for dedicated NIRSpec MOS background slits from the
   lower left corner of the slit to the middle of the slit. [#8461]
 
@@ -41,6 +43,8 @@ documentation
 
 - Added documentation for multiprocessing. [#8408]
 
+- Added documentation for NIRCam GRISM time series pointing offsets. [#8449]
+
 extract_1d
 ----------
 
@@ -53,6 +57,11 @@ extract_1d
   regions are specified and the lower and upper limits for one of them are
   outside the valid area for some data range. [#8433]
 
+extract_2d
+----------
+
+- Added handling for NIRCam GRISM time series pointing offsets. [#8449]
+
 flat_field
 ----------
 

docs/jwst/extract_2d/main.rst

@@ -245,6 +245,13 @@ extension header of the input image. These values are used to set the source loc
 for all computations involving the extent of the spectral trace and pixel wavelength
 assignments.
 
+In rare cases, it may be desirable to shift the source location in the X-direction, e.g.
+for a custom noise suppression scheme. This is achieved in the APT by specifying an 
+offset special requirement, and shows up in the header keyword "XOFFSET". The 
+``extract_2d`` step accounts for this offset by simply shifting the wavelength array by
+the appropriate amount. The WCS information remains unchanged. Note that offsets in the 
+Y-direction (cross-dispersion direction) are not supported and should not be attempted.
+
 NIRCam subarrays used for TSGRISM observations always have their "bottom" edge located
 at the physical bottom edge of the detector and vary in size vertically.
 The source spectrum trace will always be centered somewhere near row 34 in the vertical

jwst/assign_wcs/nircam.py

@@ -178,6 +178,11 @@ def tsgrism(input_model, reference_files):
     frame coordinates around the trace transform.
 
     TSGRISM is only slated to work with GRISMR and Mod A
+
+    For this mode, the source is typically at crpix1 x crpix2, which
+    are stored in keywords XREF_SCI, YREF_SCI.
+    offset special requirements may be encoded in the X_OFFSET parameter,
+    but those are handled in extract_2d.
     """
 
     # make sure this is a grism image
@@ -230,14 +235,12 @@ def tsgrism(input_model, reference_files):
 
     # input into the forward transform is x,y,x0,y0,order
     # where x,y is the pixel location in the grism image
-    # and x0,y0 is the source location in the "direct" image
-    # For this mode, the source is always at crpix1 x crpix2, which
-    # are stored in keywords XREF_SCI, YREF_SCI.
-    # Discussion with nadia that wcsinfo might not be available
-    # here but crpix info could be in wcs.source_location or similar
-    # TSGRISM mode places the sources at crpix, and all subarrays
-    # begin at 0,0, so no need to translate the crpix to full frame
-    # because they already are in full frame coordinates.
+    # and x0,y0 is the source location in the "direct" image.
+    # For this mode (tsgrism), it is assumed that the source is
+    # at the nominal aperture reference point, i.e., 
+    # crpix1 <--> xref_sci and crpix2 <--> yref_sci
+    # offsets in X are handled in extract_2d, e.g. if an offset
+    # special requirement was specified in the APT.
     xc, yc = (input_model.meta.wcsinfo.siaf_xref_sci, input_model.meta.wcsinfo.siaf_yref_sci)
 
     if xc is None:

jwst/assign_wcs/tests/test_nircam.py

@@ -8,7 +8,7 @@
 file.
 
 """
-from numpy.testing import assert_allclose
+import numpy as np
 import pytest
 
 
@@ -99,6 +99,7 @@ def create_imaging_wcs():
     return wcsobj
 
 
+@pytest.fixture
 def create_tso_wcs(filtername=tsgrism_filters[0], subarray="SUBGRISM256"):
     """Help create tsgrism GWCS object."""
     hdul = create_hdul(exptype='NRC_TSGRISM', pupil='GRISMR',
@@ -159,16 +160,20 @@ def test_nircam_wfss_available_frames():
         assert all([a == b for a, b in zip(nircam_wfss_frames, available_frames)])
 
 
-def test_nircam_tso_available_frames():
+def test_nircam_tso_available_frames(create_tso_wcs):
     """Make sure that the expected GWCS reference frames for TSO are created."""
-    wcsobj = create_tso_wcs()
+    wcsobj = create_tso_wcs
     available_frames = wcsobj.available_frames
     assert all([a == b for a, b in zip(nircam_tsgrism_frames, available_frames)])
 
 
 @pytest.mark.parametrize('key', ['xref_sci', 'yref_sci'])
 def test_extract_tso_object_fails_without_xref_yref(tsgrism_inputs, key):
-    with pytest.raises(ValueError):
+    '''
+    TypeError for xref_sci because computing dither offset is attempted, get float - NoneType
+    ValueError for yref_sci because computing dither offset is not attempted
+    '''
+    with pytest.raises((TypeError, ValueError)):
         nircam.tsgrism(*tsgrism_inputs(missing_key=key))
 
 
@@ -197,25 +202,29 @@ def test_traverse_wfss_grisms():
         traverse_wfss_trace(pupil)
 
 
-@pytest.mark.xfail(reason="Fails due to V2 NIRCam specwcs ref files delivered to CRDS")
-def test_traverse_tso_grism():
-    """Make sure that the TSO dispersion polynomials are reversable."""
-    wcsobj = create_tso_wcs()
+def test_traverse_tso_grism(create_tso_wcs):
+    """Make sure that the TSO dispersion polynomials are reversable.
+    All assert statements are in pixel space so 1/1000 px seems easily acceptable"""
+    wcsobj = create_tso_wcs
     detector_to_grism = wcsobj.get_transform('direct_image', 'grism_detector')
     grism_to_detector = wcsobj.get_transform('grism_detector', 'direct_image')
 
     # TSGRISM always has same source locations
     # takes x,y,order -> ra, dec, wave, order
-    xin, yin, order = (100, 100, 1)
-
+    xin, yin, order = (100, 35, 1)
     x0, y0, lam, orderdet = grism_to_detector(xin, yin, order)
     x, y, orderdet = detector_to_grism(x0, y0, lam, order)
 
-    assert x0 == wcs_tso_kw['xref_sci']
-    assert y0 == wcs_tso_kw['yref_sci']
+    assert np.isclose(x0, wcs_tso_kw['xref_sci'], atol = 1e-3)
+    assert np.isclose(y0, wcs_tso_kw['yref_sci'], atol = 1e-3)
     assert order == orderdet
-    assert_allclose(x, xin)
-    assert y == wcs_tso_kw['yref_sci']
+    assert np.isclose(x, xin, atol=1e-3)
+
+    # this roundtrip fails to account for the ~22.5 pixel shift from target acquisition
+    # to science position; grism shifts spectrum wrt direct image by that amount.
+    # as of 29 April 2024, how to properly store this in the header and propagate
+    # it through assign_wcs is being discussed
+    # assert np.isclose(y, wcs_tso_kw['yref_sci'])
 
 
 def test_imaging_frames():

jwst/extract_2d/grisms.py

@@ -12,7 +12,9 @@
 from gwcs.utils import _toindex
 
 from stdatamodels.jwst import datamodels
-from stdatamodels.jwst.datamodels import WavelengthrangeModel
+from stdatamodels.jwst.datamodels import WavelengthrangeModel, ImageModel
+from stdatamodels.jwst.transforms.models import IdealToV2V3
+from astropy.modeling import CompoundModel
 
 from ..assign_wcs import util
 
@@ -133,7 +135,8 @@ def extract_tso_object(input_model,
         lmin, lmax = range_select.pop()
 
         # Create the order bounding box
-        source_xpos = input_model.meta.wcsinfo.siaf_xref_sci - 1  # remove FITS 1-indexed offset
+        distortion = subwcs.get_transform("v2v3", "direct_image")
+        source_xpos, _ = compute_tso_offset_center(input_model, distortion) # 1-indexing already handled
         source_ypos = input_model.meta.wcsinfo.siaf_yref_sci - 1  # remove FITS 1-indexed offset
         transform = input_model.meta.wcs.get_transform('direct_image', 'grism_detector')
         xmin, ymin, _ = transform(source_xpos,
@@ -151,9 +154,6 @@ def extract_tso_object(input_model,
         # model with the known object center and order information (in pixels of direct image)
         # This changes the user input to the model from (x,y,x0,y0,order) -> (x,y)
         #
-        # The bounding box is limited to the size of the detector in the dispersion direction
-        # and 64 pixels in the cross-dispersion direction (at request of instrument team).
-        #
         # The team wants the object to fall near row 34 for all cutouts, but the default cutout
         # height is 64 pixels (32 on either side). So bump the extraction ycenter, when necessary,
         # so that the height is 30 above and 34 below (in full frame) the object center.
@@ -177,6 +177,8 @@ def extract_tso_object(input_model,
             raise ValueError("Something bad happened calculating extraction y-size")
 
         # Limit the bounding box to the detector edges
+        # The bounding box is limited to the size of the detector in the dispersion direction
+        # and 64 pixels in the cross-dispersion direction (at request of instrument team).
         ymin, ymax = (max(extract_y_min, 0), min(extract_y_max, input_model.meta.subarray.ysize))
         xmin, xmax = (max(xmin, 0), min(xmax, input_model.meta.subarray.xsize))
 
@@ -240,14 +242,15 @@ def extract_tso_object(input_model,
             output_model.var_flat = var_flat
             output_model.meta.wcs = subwcs
             output_model.meta.wcs.bounding_box = util.wcs_bbox_from_shape(ext_data.shape)
+            if compute_wavelength:
+                # this seems to happen near-instantly, suggest remove this message
+                #log.debug("Computing wavelengths (this takes a while ...)")
+                output_model.wavelength = compute_tso_wavelength_array(output_model)
             output_model.meta.wcsinfo.siaf_yref_sci = 34  # update for the move, vals are the same
-            output_model.meta.wcsinfo.siaf_xref_sci = input_model.meta.wcsinfo.siaf_xref_sci
+            output_model.meta.wcsinfo.siaf_xref_sci = source_xpos + 1  # back to 1-indexed
             output_model.meta.wcsinfo.spectral_order = order
             output_model.meta.wcsinfo.dispersion_direction = \
                 input_model.meta.wcsinfo.dispersion_direction
-            if compute_wavelength:
-                log.debug("Computing wavelengths (this takes a while ...)")
-                output_model.wavelength = compute_wavelength_array(output_model)
             output_model.name = '1'
             output_model.source_type = input_model.meta.target.source_type
             output_model.source_name = input_model.meta.target.catalog_name
@@ -562,7 +565,7 @@ def compute_dispersion(wcs):
     raise NotImplementedError
 
 
-def compute_wavelength_array(slit):
+def compute_tso_wavelength_array(slit):
     """
     Compute the wavelength array for a slit with gwcs object
 
@@ -576,12 +579,50 @@ def compute_wavelength_array(slit):
     wavelength : numpy.array
         The wavelength array
     """
-    transform = slit.meta.wcs.forward_transform
-    x, y = grid_from_bounding_box(slit.meta.wcs.bounding_box)
-    wavelength = transform(x, y)[2]
+    wcs = slit.meta.wcs
+    full_transform = slit.meta.wcs.forward_transform
+    x, y = grid_from_bounding_box(wcs.bounding_box)
+    wavelength = full_transform(x, y)[2]
     return wavelength
 
 
+def compute_tso_offset_center(input_model: ImageModel, distortion: CompoundModel) -> tuple[float, float]:
+    """
+    In the case that an Offset Special Requirement is requested in the APT,
+    the source is no longer at the aperture reference point.
+    The dither.x_offset and dither.y_offset values encode the offset
+    in units of arcseconds. They need to be translated from Ideal to 
+    detector coordinates and into pixel units.
+
+    Parameters
+    ----------
+    input_model : ImageModel
+        The input data model
+    distortion : DistortionModel
+        The distortion model
+
+    Returns
+    -------
+    xc, yc : tuple
+        The x and y center of the image in direct image coordinates
+
+    Notes
+    -----
+    The wavelength is not used for the distortion calculation between
+    v2v3 and direct image coordinates, so this can be hardcoded to NaN.
+    """
+
+    idltov23 = IdealToV2V3(input_model.meta.wcsinfo.v3yangle,
+                input_model.meta.wcsinfo.v2_ref,
+                input_model.meta.wcsinfo.v3_ref,
+                input_model.meta.wcsinfo.vparity)
+    v2_offset, v3_offset = idltov23(input_model.meta.dither.x_offset, input_model.meta.dither.y_offset)
+    wavelength = np.nan
+    xc, yc, _, _ = distortion(v2_offset, v3_offset, wavelength, 1)
+
+    return xc, yc
+
+
 def compute_wfss_wavelength(slit):
     """
     Compute the wavelength array for a slit with gwcs object

jwst/extract_2d/tests/test_grisms.py

@@ -27,7 +27,7 @@
 from jwst.assign_wcs import AssignWcsStep, nircam
 
 from jwst.extract_2d.extract_2d_step import Extract2dStep
-from jwst.extract_2d.grisms import extract_tso_object, extract_grism_objects
+from jwst.extract_2d.grisms import extract_tso_object, extract_grism_objects, compute_tso_offset_center
 from jwst.extract_2d.tests import data
 
 
@@ -66,6 +66,7 @@
 
 wcs_tso_kw = {'wcsaxes': 2, 'ra_ref': 86.9875, 'dec_ref': 23.423,
               'v2_ref': 95.043034, 'v3_ref': -556.150466, 'roll_ref': 359.9521,
+              'v3i_yang': -0.37562675, 'vparity': -1,
               }
 
 
@@ -97,6 +98,8 @@ def create_hdul(detector='NRCALONG', channel='LONG', module='A',
     phdu.header['SUBSIZE2'] = 2048
     phdu.header['SUBSTRT1'] = 1
     phdu.header['SUBSTRT2'] = 1
+    phdu.header['XOFFSET'] = 5.0  # random offset for testing
+    phdu.header['YOFFSET'] = 1.45
     scihdu = fits.ImageHDU()
     scihdu.header['EXTNAME'] = "SCI"
     scihdu.header.update(wcskeys)
@@ -373,6 +376,15 @@ def test_extract_tso_height():
     del outmodel
 
 
+def test_compute_tso_offset_center():
+
+    image_model = create_tso_wcsimage(filtername="F444W", subarray=False)
+    distortion = image_model.meta.wcs.get_transform("v2v3", "direct_image")
+    xc, yc = compute_tso_offset_center(image_model, distortion)
+    assert np.isclose(yc, 59.929, atol=1e-3)
+    assert np.isclose(xc, 961.355, atol=1e-3)
+
+
 @pytest.mark.filterwarnings("ignore: Card is too long")
 def test_extract_wfss_object():
     """Test extraction of a WFSS object.

jwst/regtest/test_nircam_tsgrism.py

@@ -29,6 +29,35 @@ def run_pipelines(rtdata_module):
     return rtdata
 
 
+@pytest.fixture(scope="module", 
+                params=["jw01185015001_03104_00001-seg004_nrcalong_rate.fits",
+                        "jw01185013001_04103_00001-seg003_nrcalong_rate.fits"])
+def run_pipeline_offsetSR(request, rtdata_module):
+    ''''''
+    rtdata = rtdata_module
+    rtdata.get_data("nircam/tsgrism/" + request.param)
+    args = ["calwebb_spec2", rtdata.input,
+            "--steps.extract_1d.save_results=True",]
+    Step.from_cmdline(args)
+    return rtdata
+
+
+@pytest.mark.bigdata
+def test_nircam_tsgrism_stage2_offsetSR(run_pipeline_offsetSR, fitsdiff_default_kwargs):
+    '''
+    Test coverage for offset special requirement specifying nonzero offset in X.
+    Test data are two observations of Gliese 436, one with offset specified and one without.
+    Quantitatively we just ensure that the outputs are identical to the inputs, but qualitatively
+    can check that the spectral lines fall at the same wavelengths in both cases,
+    which is why the zero-offset case is also included here.'''
+    rtdata = run_pipeline_offsetSR
+    rtdata.output = rtdata.input.replace("rate", "x1d")
+    rtdata.get_truth("truth/test_nircam_tsgrism_stages/" + rtdata.output.split('/')[-1])
+
+    diff = FITSDiff(rtdata.output, rtdata.truth, **fitsdiff_default_kwargs)
+    assert diff.identical, diff.report()
+
+
 @pytest.mark.bigdata
 @pytest.mark.parametrize("suffix", ["calints", "extract_2d", "flat_field",
                                     "o012_crfints", "srctype", "x1dints"])