Bugfix for WCS handing non-coordinate inputs/outputs correctly (#706)

Author: WilliamJamieson

changes/706.bugfix.rst

@@ -0,0 +1 @@
+Bugfix for supporting non-coordinate inputs to transforms supporting quantities.

gwcs/coordinate_frames/_base.py

@@ -1,4 +1,5 @@
 import abc
+from itertools import zip_longest
 
 import numpy as np
 from astropy import units as u
@@ -136,8 +137,11 @@ def add_units(
             return u.Quantity(arrays, self.unit[0])
 
         return tuple(
-            u.Quantity(array, unit=unit)
-            for array, unit in zip(arrays, self.unit, strict=True)
+            array if unit is None else u.Quantity(array, unit=unit)
+            # zip_longest is used here to support "non-coordinate" inputs/outputs
+            #   This implicitly assumes that the "non-coordinate" inputs/outputs
+            #   are tacked onto the end of the tuple of "coordinate" inputs/outputs.
+            for array, unit in zip_longest(arrays, self.unit)
         )
 
     def remove_units(
@@ -151,5 +155,8 @@ def remove_units(
 
         return tuple(
             array.to_value(unit) if isinstance(array, u.Quantity) else array
-            for array, unit in zip(arrays, self.unit, strict=True)
+            # zip_longest is used here to support "non-coordinate" inputs/outputs
+            #   This implicitly assumes that the "non-coordinate" inputs/outputs
+            #   are tacked onto the end of the tuple of "coordinate" inputs/outputs.
+            for array, unit in zip_longest(arrays, self.unit)
         )

gwcs/tests/test_wcs.py

@@ -1925,3 +1925,39 @@ def test_fitswcs_transform(lon, lat):
     assert_allclose(fwcs(0, 0), (lon, lat), atol=1e-14)
     assert_allclose(fwcs(0, lat)[0], lon + 180, atol=1e-14)
     assert_allclose(fwcs.inverse(lon, lat), (0, 0), atol=1e-14)
+
+
+def test_non_coordinate_input():
+    """Test that non-coordinate inputs are handled correctly."""
+    gw = wcs.WCS(
+        [
+            (
+                cf.Frame2D(
+                    name="detector",
+                    axes_names=("x", "y"),
+                    unit=(u.deg, u.deg),
+                    axes_order=(0, 1),
+                ),
+                models.Identity(3),
+            ),
+            (
+                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",
+                ),
+                None,
+            ),
+        ]
+    )
+    assert gw(1, 2, 3) == (1, 2, 3)
+    assert gw(1 * u.deg, 2 * u.deg, 3) == (1 * u.deg, 2 * u.deg, 3 * u.micron)