spacetelescope / gwcs

Generalized World Coordinate System: provides tools for managing WCS in a general way
https://gwcs.readthedocs.io/en/latest/
45 stars 46 forks source link

Splitting the celestial axes of a transform causes errors #455

Open Cadair opened 1 year ago

Cadair commented 1 year ago

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..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..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..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..__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 commented 1 year ago

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.