Splitting the celestial axes of a transform causes errors

[Cadair]

Take the following setup:

import astropy.modeling.models as m
import gwcs.coordinate_frames as cf
import astropy.units as u
from gwcs import WCS

# Setup a model where the pixel axes are (lat, wave, lon)
spatial = m.Multiply(10*u.arcsec/u.pix) & m.Multiply(15*u.arcsec/u.pix)  # pretend this is a spatial model
compound = m.Linear1D(intercept=0*u.nm, slope=10*u.nm/u.pix) & spatial
forward = m.Mapping((1, 2, 0)) | compound | m.Mapping((2, 0, 1))

# Setup the output frame
celestial_frame = cf.CelestialFrame(axes_order=(0, 2), unit=(u.arcsec, u.arcsec))
spectral_frame = cf.SpectralFrame(axes_order=(1,), unit=u.nm)
output_frame = cf.CompositeFrame([spectral_frame, celestial_frame])

input_frame = cf.CoordinateFrame(3, ["PIXEL"]*3, axes_order=list(range(3)), unit=[u.pix]*3)

wcs = WCS(forward, input_frame, output_frame)
input_pixel = [0*u.pix]*3
output_world = wcs.pixel_to_world_values(*input_pixel)
output_pixel = wcs.world_to_pixel_values(*output_world)

Here we have a 3D cube where the axes of the array correspond to (lat, wave, lon) in world space.
We setup a model with the wave transform first, and then the spatial transform (ordered (lon, lat)) and then use mappings to make the inputs and outputs match the world & pixel order.

Likewise, we setup the output frame to match the output of the model, where the axes are split [1].

When this code is run, the last line errors with the following output:

Click for full traceback

UnitsError: Linear1D: Units of input 'x', arcsec (angle), could not be converted to required input units of nm (length)

---------------------------------------------------------------------------
UnitsError                                Traceback (most recent call last)
File ~/projects/visp_wcs/build_example_split.py:23
     21 input_pixel = [0*u.pix]*3
     22 output_world = wcs.pixel_to_world_values(*input_pixel)
---> 23 output_pixel = wcs.world_to_pixel_values(*output_world)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/gwcs/api.py:140, in GWCSAPIMixin.world_to_pixel_values(self, *world_arrays)
    127 """
    128 Convert world coordinates to pixel coordinates.
    129 
   (...)
    136 horizontal coordinate and ``y`` is the vertical coordinate.
    137 """
    138 world_arrays = self._add_units_input(world_arrays, self.backward_transform, self.output_frame)
--> 140 result = self.invert(*world_arrays, with_units=False)
    142 return self._remove_quantity_output(result, self.input_frame)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/gwcs/wcs.py:508, in WCS.invert(self, *args, **kwargs)
    505 try:
    506     # remove iterative inverse-specific keyword arguments:
    507     akwargs = {k: v for k, v in kwargs.items() if k not in _ITER_INV_KWARGS}
--> 508     result = self.backward_transform(*args, **akwargs)
    509 except (NotImplementedError, KeyError):
    510     result = self.numerical_inverse(*args, **kwargs, with_units=with_units)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1125, in Model.__call__(self, *args, **kwargs)
   1120 # prepare for model evaluation (overridden in CompoundModel)
   1121 evaluate, inputs, broadcasted_shapes, kwargs = self._pre_evaluate(
   1122     *args, **kwargs
   1123 )
-> 1125 outputs = self._generic_evaluate(evaluate, inputs, fill_value, with_bbox)
   1127 # post-process evaluation results (overridden in CompoundModel)
   1128 return self._post_evaluate(
   1129     inputs, outputs, broadcasted_shapes, with_bbox, **kwargs
   1130 )

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1087, in Model._generic_evaluate(self, evaluate, _inputs, fill_value, with_bbox)
   1085     outputs = bbox.evaluate(evaluate, _inputs, fill_value)
   1086 else:
-> 1087     outputs = evaluate(_inputs)
   1088 return outputs

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:3308, in CompoundModel._pre_evaluate.<locals>.evaluate(_inputs)
   3307 def evaluate(_inputs):
-> 3308     return self._evaluate(*_inputs, **kwargs)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:3351, in CompoundModel._evaluate(self, *args, **kw)
   3349 elif op == "|":
   3350     if isinstance(leftval, tuple):
-> 3351         return self.right(*leftval, **kw)
   3352     else:
   3353         return self.right(leftval, **kw)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1125, in Model.__call__(self, *args, **kwargs)
   1120 # prepare for model evaluation (overridden in CompoundModel)
   1121 evaluate, inputs, broadcasted_shapes, kwargs = self._pre_evaluate(
   1122     *args, **kwargs
   1123 )
-> 1125 outputs = self._generic_evaluate(evaluate, inputs, fill_value, with_bbox)
   1127 # post-process evaluation results (overridden in CompoundModel)
   1128 return self._post_evaluate(
   1129     inputs, outputs, broadcasted_shapes, with_bbox, **kwargs
   1130 )

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1087, in Model._generic_evaluate(self, evaluate, _inputs, fill_value, with_bbox)
   1085     outputs = bbox.evaluate(evaluate, _inputs, fill_value)
   1086 else:
-> 1087     outputs = evaluate(_inputs)
   1088 return outputs

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:3308, in CompoundModel._pre_evaluate.<locals>.evaluate(_inputs)
   3307 def evaluate(_inputs):
-> 3308     return self._evaluate(*_inputs, **kwargs)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:3336, in CompoundModel._evaluate(self, *args, **kw)
   3334 if op != "fix_inputs":
   3335     if op != "&":
-> 3336         leftval = self.left(*args, **kw)
   3337         if op != "|":
   3338             rightval = self.right(*args, **kw)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1125, in Model.__call__(self, *args, **kwargs)
   1120 # prepare for model evaluation (overridden in CompoundModel)
   1121 evaluate, inputs, broadcasted_shapes, kwargs = self._pre_evaluate(
   1122     *args, **kwargs
   1123 )
-> 1125 outputs = self._generic_evaluate(evaluate, inputs, fill_value, with_bbox)
   1127 # post-process evaluation results (overridden in CompoundModel)
   1128 return self._post_evaluate(
   1129     inputs, outputs, broadcasted_shapes, with_bbox, **kwargs
   1130 )

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1087, in Model._generic_evaluate(self, evaluate, _inputs, fill_value, with_bbox)
   1085     outputs = bbox.evaluate(evaluate, _inputs, fill_value)
   1086 else:
-> 1087     outputs = evaluate(_inputs)
   1088 return outputs

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:3308, in CompoundModel._pre_evaluate.<locals>.evaluate(_inputs)
   3307 def evaluate(_inputs):
-> 3308     return self._evaluate(*_inputs, **kwargs)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:3343, in CompoundModel._evaluate(self, *args, **kw)
   3340         rightval = None
   3342 else:
-> 3343     leftval = self.left(*(args[: self.left.n_inputs]), **kw)
   3344     rightval = self.right(*(args[self.left.n_inputs :]), **kw)
   3346 if op != "|":

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:418, in __call__(self, model_set_axis, with_bounding_box, fill_value, equivalencies, inputs_map, *inputs, **new_inputs)
    409 args = ("self",)
    410 kwargs = {
    411     "model_set_axis": None,
    412     "with_bounding_box": False,
   (...)
    415     "inputs_map": None,
    416 }
--> 418 new_call = make_function_with_signature(
    419     __call__, args, kwargs, varargs="inputs", varkwargs="new_inputs"
    420 )
    422 # The following makes it look like __call__
    423 # was defined in the class
    424 update_wrapper(new_call, cls)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:395, in _ModelMeta._handle_special_methods.<locals>.__call__(self, *inputs, **kwargs)
    393 def __call__(self, *inputs, **kwargs):
    394     """Evaluate this model on the supplied inputs."""
--> 395     return super(cls, self).__call__(*inputs, **kwargs)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:1121, in Model.__call__(self, *args, **kwargs)
   1118 fill_value = kwargs.pop("fill_value", np.nan)
   1120 # prepare for model evaluation (overridden in CompoundModel)
-> 1121 evaluate, inputs, broadcasted_shapes, kwargs = self._pre_evaluate(
   1122     *args, **kwargs
   1123 )
   1125 outputs = self._generic_evaluate(evaluate, inputs, fill_value, with_bbox)
   1127 # post-process evaluation results (overridden in CompoundModel)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:969, in Model._pre_evaluate(self, *args, **kwargs)
    965 """
    966 Model specific input setup that needs to occur prior to model evaluation.
    967 """
    968 # Broadcast inputs into common size
--> 969 inputs, broadcasted_shapes = self.prepare_inputs(*args, **kwargs)
    971 # Setup actual model evaluation method
    972 parameters = self._param_sets(raw=True, units=True)

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:2091, in Model.prepare_inputs(self, model_set_axis, equivalencies, *inputs, **kwargs)
   2087 self._validate_input_shapes(inputs, self.inputs, model_set_axis)
   2089 inputs_map = kwargs.get("inputs_map", None)
-> 2091 inputs = self._validate_input_units(inputs, equivalencies, inputs_map)
   2093 # The input formatting required for single models versus a multiple
   2094 # model set are different enough that they've been split into separate
   2095 # subroutines
   2096 if self._n_models == 1:

File ~/.virtualenvs/visp_wcs/lib/python3.11/site-packages/astropy/modeling/core.py:2169, in Model._validate_input_units(self, inputs, equivalencies, inputs_map)
   2161             raise UnitsError(
   2162                 f"{name}: Units of input '{self.inputs[i]}', "
   2163                 f"{inputs[i].unit} ({inputs[i].unit.physical_type}),"
   (...)
   2166                 "input"
   2167             )
   2168         else:
-> 2169             raise UnitsError(
   2170                 f"{name}: Units of input '{self.inputs[i]}', "
   2171                 f"{inputs[i].unit} ({inputs[i].unit.physical_type}),"
   2172                 " could not be "
   2173                 "converted to required input"
   2174                 f" units of {input_unit} ({input_unit.physical_type})"
   2175             )
   2176 else:
   2177     # If we allow dimensionless input, we add the units to the
   2178     # input values without conversion, otherwise we raise an
   2179     # exception.
   2181     if (
   2182         not self.input_units_allow_dimensionless[input_name]
   2183         and input_unit is not dimensionless_unscaled
   2184         and input_unit is not None
   2185     ):

UnitsError: Linear1D: Units of input 'x', arcsec (angle), could not be converted to required input units of nm (length)

---

[1] Should CelestialFrame always have lon first like the models do, or should it be able to have lon and lat in any order?

[Cadair]

As far as I got with investigating this was noticing that CompoundFrame.coordinate_to_quantity uses axes_order to order the inputs which are then passed through to the model.