Update tests to always define forward transform and a full coordinate frame for GWCS

Author: WilliamJamieson

jwst/assign_wcs/tests/test_wfss.py

@@ -1,7 +1,10 @@
 import numpy as np
 import pytest
+from astropy import units as u
+from astropy.coordinates import ICRS
 from astropy.modeling.models import Const1D, Mapping, Polynomial1D, Polynomial2D, Shift
 from gwcs import wcs
+from gwcs.coordinate_frames import CelestialFrame, CompositeFrame, Frame2D, SpectralFrame
 from gwcs.wcstools import grid_from_bounding_box
 from numpy.testing import assert_allclose
 from stdatamodels.jwst.datamodels import SlitModel
@@ -21,7 +24,16 @@ def create_nircam_slit(model, x0, y0, order):
         | (Shift(xmin) & Shift(ymin) & Const1D(x0) & Const1D(y0) & Const1D(order))
         | model
     )
-    wcsobj = wcs.WCS([("det", model), ("world", None)])
+    detector = Frame2D(name="detector")
+    # Note there is a "non-coordinate" output from the model.
+    world = CompositeFrame(
+        [
+            CelestialFrame(name="sky", axes_order=(0, 1), reference_frame=ICRS()),
+            SpectralFrame(name="spectral", axes_order=(2,), unit=(u.um,)),
+        ],
+        name="world",
+    )
+    wcsobj = wcs.WCS([(detector, model), (world, None)])
     wcsobj.bounding_box = ((20, 25), (800, 805))
     slit = SlitModel()
     slit.meta.wcs = wcsobj
@@ -39,7 +51,16 @@ def create_niriss_slit(model, x0, y0, order):
         | (Shift(xmin) & Shift(ymin) & Const1D(x0) & Const1D(y0) & Const1D(order))
         | model
     )
-    wcsobj = wcs.WCS([("det", model), ("world", None)])
+    detector = Frame2D(name="detector")
+    # Note there is a "non-coordinate" output from the model.
+    world = CompositeFrame(
+        [
+            CelestialFrame(name="sky", axes_order=(0, 1), reference_frame=ICRS()),
+            SpectralFrame(name="spectral", axes_order=(2,), unit=(u.um,)),
+        ],
+        name="world",
+    )
+    wcsobj = wcs.WCS([(detector, model), (world, None)])
     slit = SlitModel()
     slit.meta.wcs = wcsobj
     slit.source_xpos = x0

jwst/datamodels/tests/test_library.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 import stdatamodels.jwst.datamodels
+from astropy.modeling.models import Identity
 from astropy.time import Time
 from gwcs import coordinate_frames as cf
 from stdatamodels.jwst.datamodels import ImageModel
@@ -274,7 +275,9 @@ def test_read_meta_from_open_model(example_asn_path, flatten):
         model.int_list = [1, 2, 3, 4]
         model.nested_list = [[{"key": "value"}], [{"key2": "value2"}]] * 2
         model.meta.wcs = gwcs.WCS(
-            input_frame=cf.Frame2D(name="input"), output_frame=cf.Frame2D(name="output")
+            input_frame=cf.Frame2D(name="input"),
+            output_frame=cf.Frame2D(name="output"),
+            forward_transform=Identity(2),
         )
         model.unsupported_type = set([1, 2, 3])
         meta = _read_meta_from_open_model(model, flatten)

jwst/extract_1d/tests/test_source_location.py

@@ -1,6 +1,7 @@
 import gwcs
 import numpy as np
 import pytest
+from astropy import units as u
 from astropy.modeling.models import Identity, Scale
 
 from jwst.extract_1d import source_location as sl
@@ -288,7 +289,20 @@ def test_nod_pair_location_miri(mock_miri_lrs_fs, dispaxis):
 
     # mock v2v3 transform
     model.meta.wcs = gwcs.WCS(
-        [model.meta.wcs.pipeline[0], ("v2v3", Identity(3)), model.meta.wcs.pipeline[1]]
+        [
+            model.meta.wcs.pipeline[0],
+            (
+                gwcs.CompositeFrame(
+                    [
+                        gwcs.Frame2D(name="frame", axes_order=(0, 1)),
+                        gwcs.SpectralFrame(name="lambda", axes_order=(2,), unit=(u.um,)),
+                    ],
+                    name="v2v3",
+                ),
+                Identity(3),
+            ),
+            model.meta.wcs.pipeline[1],
+        ]
     )
     model.meta.wcsinfo.v3yangle = 1.0
     model.meta.wcsinfo.v2_ref = 1.0

jwst/lib/tests/test_wcs_utils.py

@@ -99,7 +99,17 @@ def test_get_wavelengths_soss():
 
     # mock a SOSS wcs
     soss_transform = NirissSOSSModel([1], [model.meta.wcs])
-    new_wcs = wcs.WCS([("detector", soss_transform), ("world", None)])
+    detector = cf.Frame2D(name="detector")
+    world = cf.CompositeFrame(
+        [
+            cf.CelestialFrame(name="icrs", axes_order=(0, 1), reference_frame=coord.ICRS()),
+            cf.SpectralFrame(
+                name="spectral", axes_order=(2,), unit=(u.micron,), axes_names=("wavelength",)
+            ),
+        ],
+        name="world",
+    )
+    new_wcs = wcs.WCS([(detector, soss_transform), (world, None)])
     model.meta.wcs = new_wcs
 
     # calculate what the wavelength array should be

jwst/pixel_replace/tests/test_pixel_replace.py

@@ -1,6 +1,7 @@
 import os
 from glob import glob
 
+import astropy.units as u
 import gwcs
 import numpy as np
 import pytest
@@ -128,8 +129,14 @@ def miri_mrs():
 
     # Mock a wcs that just returns 1 for alpha, beta, lam
     transform = Mapping((0, 1, 1), n_inputs=2) | Const1D(1) & Const1D(1) & Const1D(1)
-    model.meta.wcs = gwcs.WCS([("detector", transform), ("alpha_beta", None)])
-
+    output_frame = gwcs.CompositeFrame(
+        [
+            gwcs.Frame2D(name="alpha_beta_spatial", axes_order=(0, 1), unit=(u.arcsec, u.arcsec)),
+            gwcs.SpectralFrame(name="lam", axes_order=(2,), unit=(u.um,)),
+        ],
+        name="alpha_beta",
+    )
+    model.meta.wcs = gwcs.WCS([(gwcs.Frame2D(name="detector"), transform), (output_frame, None)])
     return model, bad_idx
 
 

jwst/residual_fringe/tests/test_residual_fringe.py

@@ -1,6 +1,7 @@
 import gwcs
 import numpy as np
 import pytest
+from astropy import units as u
 from astropy.modeling.models import Const1D, Identity, Mapping
 from stdatamodels.jwst import datamodels
 
@@ -260,7 +261,21 @@ def test_get_wavemap():
 
     # Mock a WCS that returns 1 for wavelengths
     transform = Mapping((0, 1, 1), n_inputs=2) | Identity(2) & Const1D(1)
-    model.meta.wcs = gwcs.WCS([("detector", transform), ("world", None)])
+    model.meta.wcs = gwcs.WCS(
+        [
+            (gwcs.Frame2D(name="detector"), transform),
+            (
+                gwcs.CompositeFrame(
+                    [
+                        gwcs.Frame2D(name="sky0", axes_order=(0, 1)),
+                        gwcs.SpectralFrame(name="spectral", axes_order=(2,), unit=u.micron),
+                    ],
+                    name="world",
+                ),
+                None,
+            ),
+        ]
+    )
 
     rf = ResidualFringeCorrection(model, "N/A", "N/A", None)
     wavemap = rf._get_wave_map()