assign_wcs accepts CubeModel for TSGrism [spacetelescope#1710]

jwst/assign_wcs/nircam.py

@@ -7,7 +7,10 @@
 
 from . import pointing
 from .util import not_implemented_mode, subarray_transform
-from ..datamodels import ImageModel, NIRCAMGrismModel, DistortionModel
+from ..datamodels import (ImageModel,
+                          NIRCAMGrismModel,
+                          DistortionModel,
+                          CubeModel)
 from ..transforms.models import (NIRCAMForwardRowGrismDispersion,
                                  NIRCAMForwardColumnGrismDispersion,
                                  NIRCAMBackwardGrismDispersion)
@@ -66,13 +69,12 @@ def imaging_distortion(input_model, reference_files):
     return transform
 
 
-
 def tsgrism(input_model, reference_files):
     """Create WCS pipeline for a NIRCAM Time Series Grism observation.
 
     Parameters
     ----------
-    input_model: jwst.datamodels.ImagingModel
+    input_model: jwst.datamodels.ImagingModel, jwst.datamodels.CubeModel
         The input datamodel, derived from datamodels
     reference_files: dict
         Dictionary {reftype: reference file name}.
@@ -89,20 +91,19 @@ def tsgrism(input_model, reference_files):
     """
 
     # The input is the grism image
-    if not isinstance(input_model, ImageModel):
-        raise TypeError('The input data model must be an ImageModel.')
+    if not isinstance(input_model, (ImageModel, CubeModel)):
+        raise TypeError('The input data model must be an ImageModel or CubeModel.')
 
     # make sure this is a grism image
     if "NRC_TSGRISM" != input_model.meta.exposure.type:
         raise TypeError('The input exposure is not a NIRCAM time series grism')
 
-    if input_model.meta.instrument.module  != "A":
+    if input_model.meta.instrument.module != "A":
         raise ValueError('TSGRISM mode only supports module A')
 
     if input_model.meta.instrument.pupil != "GRISMR":
         raise ValueError('TSGRIM mode only supports GRISMR')
 
-
     gdetector = cf.Frame2D(name='grism_detector', axes_order=(0, 1), unit=(u.pix, u.pix))
     detector = cf.Frame2D(name='full_detector', axes_order=(0, 1), unit=(u.pix, u.pix))
     v2v3 = cf.Frame2D(name='v2v3', axes_order=(0, 1), unit=(u.deg, u.deg))
@@ -111,7 +112,6 @@ def tsgrism(input_model, reference_files):
     # get the shift to full frame coordinates
     subarray2full = subarray_transform(input_model)
 
-
     # translate the x,y detector-in to x,y detector out coordinates
     # Get the disperser parameters which are defined as a model for each
     # spectral order
@@ -139,7 +139,6 @@ def tsgrism(input_model, reference_files):
                                                     xmodels=dispx,
                                                     ymodels=dispy)
 
-
     # input into the forward transform is x,y,x0,y0,order
     #
     sub2direct = (subarray2full & Identity(3)) | det2det
@@ -152,16 +151,16 @@ def tsgrism(input_model, reference_files):
                 (v2v3, tel2sky),
                 (world, None)]
 
-
     return pipeline
 
+
 def grism(input_model, reference_files):
     """
     Create the WCS pipeline for a NIRCAM grism observation.
 
     Parameters
     ----------
-    input_model: jwst.datamodels.ImagingModel
+    input_model: jwst.datamodels.ImagingModel, CubeModel
         The input datamodel, derived from datamodels
     reference_files: dict
         Dictionary {reftype: reference file name}.
@@ -237,7 +236,6 @@ def grism(input_model, reference_files):
     if "NRC_GRISM" not in input_model.meta.exposure.type:
             raise TypeError('The input exposure is not a NIRCAM grism')
 
-
     # Create the empty detector as a 2D coordinate frame in pixel units
     gdetector = cf.Frame2D(name='grism_detector',
                            axes_order=(0, 1),
@@ -273,7 +271,6 @@ def grism(input_model, reference_files):
                                                     xmodels=dispx,
                                                     ymodels=dispy)
 
-
     # create the pipeline to construct a WCS object for the whole image
     # which can translate ra,dec to image frame reference pixels
     # it also needs to be part of the grism image wcs pipeline to