Fix bug in translating 3D Spectrum1D with only spectral axis defined

[rosteen]

Description

If a 3D flux is input to a Spectrum1D along with a spectral_axis (rather than a full 3D WCS), the resulting object will have a spectral GWCS of only one dimension. Personally I think this behavior is not ideal and might need rethinking at some point, but in either case this fix will stop the translator from trying to swap axes on the WCS in this case, thus fixing (I hope) the bug @pllim ran into in https://github.com/spacetelescope/jdaviz/pull/1040#issuecomment-1018006578.

Fixes #60

[codecov[bot]]

Codecov Report

Merging #61 (312a840) into main (7f3229d) will decrease coverage by 1.23%. The diff coverage is 84.09%.

@@            Coverage Diff             @@
##             main      #61      +/-   ##
==========================================
- Coverage   95.80%   94.56%   -1.24%     
==========================================
  Files          15       15              
  Lines        1145     1123      -22     
==========================================
- Hits         1097     1062      -35     
- Misses         48       61      +13     

Impacted Files	Coverage Δ
glue_astronomy/translators/spectrum1d.py	89.91% <80.00%> (-2.66%)	:arrow_down:
...lue_astronomy/translators/tests/test_spectrum1d.py	92.85% <92.85%> (-7.15%)	:arrow_down:

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 7f3229d...312a840. Read the comment docs.

[pllim]

Update from 2022-01-26: Using this PR and spacetelescope/jdaviz#1040 gives me the following traceback.

import numpy as np
from jdaviz import Cubeviz

sci_arr = np.random.random((5, 10, 10))
err_arr = np.sqrt(sci_arr)

viz = Cubeviz()
viz.load_data(sci_arr, data_type='flux', data_label='myarray[SCI]')
viz.load_data(err_arr, data_type='uncert', data_label='myarray[ERR]')
viz.app

      1 viz = Cubeviz()
----> 2 viz.load_data(sci_arr, data_type='flux', data_label='myarray[SCI]')
      3 viz.load_data(err_arr, data_type='uncert', data_label='myarray[ERR]')
      4 viz.app

.../jdaviz/core/helpers.py in load_data(self, data, parser_reference, **kwargs)
     59 
     60     def load_data(self, data, parser_reference=None, **kwargs):
---> 61         self.app.load_data(data, parser_reference=parser_reference, **kwargs)
     62 
     63     @property

.../jdaviz/app.py in load_data(self, file_obj, parser_reference, **kwargs)
    360                 # If the parser returns something other than known, assume it's
    361                 #  a message we want to make the user aware of.
--> 362                 msg = parser(self, file_obj, **kwargs)
    363 
    364                 if msg is not None:

.../jdaviz/configs/cubeviz/plugins/parsers.py in parse_data(app, file_obj, data_type, data_label)
     96 
     97         if file_obj.ndim == 3:
---> 98             _parse_ndarray_3d(app, file_obj, data_label, data_type)
     99         else:
    100             raise NotImplementedError(f'Unsupported data format: {file_obj}')

.../jdaviz/configs/cubeviz/plugins/parsers.py in _parse_ndarray_3d(app, file_obj, data_label, data_type)
    333     idx = np.arange(flux.shape[-1]) * u.pix
    334     sc = Spectrum1D(spectral_axis=idx, flux=flux)
--> 335     app.add_data(sc, data_label)
    336     _show_data_in_cubeviz_viewer(app, data_label, data_type)

.../jdaviz/app.py in add_data(self, data, data_label, notify_done)
    718             data_label = data.label
    719         data_label = data_label or "New Data"
--> 720         self.data_collection[data_label] = data
    721 
    722         # Send out a toast message

.../glue/glue/core/data_collection.py in __setitem__(self, key, data)
    408                 self.remove(existing_data)
    409 
--> 410         self.append(data)
    411 
    412     def __iter__(self):

.../glue/glue/core/data_collection.py in append(self, data)
     81                 s.register()
     82             msg = DataCollectionAddMessage(self, data)
---> 83             self.hub.broadcast(msg)
     84 
     85         self._sync_link_manager()

.../glue/glue/core/hub.py in broadcast(self, message)
    213             logging.getLogger(__name__).info("Broadcasting %s", message)
    214             for subscriber, handler in self._find_handlers(message):
--> 215                 handler(message)
    216 
    217     def __getstate__(self):

.../jdaviz/configs/default/plugins/collapse/collapse.py in _on_data_updated(self, msg)
     69             for i in range(len(self.data_items)):
     70                 try:
---> 71                     self.selected_data_item = self.data_items[i]
     72                 except (ValueError, TypeError):
     73                     continue

.../traitlets/traitlets.py in __set__(self, obj, value)
    604             raise TraitError('The "%s" trait is read-only.' % self.name)
    605         else:
--> 606             self.set(obj, value)
    607 
    608     def _validate(self, obj, value):

.../traitlets/traitlets.py in set(self, obj, value)
    593             # we explicitly compare silent to True just in case the equality
    594             # comparison above returns something other than True/False
--> 595             obj._notify_trait(self.name, old_value, new_value)
    596 
    597     def __set__(self, obj, value):

.../traitlets/traitlets.py in _notify_trait(self, name, old_value, new_value)
   1217 
   1218     def _notify_trait(self, name, old_value, new_value):
-> 1219         self.notify_change(Bunch(
   1220             name=name,
   1221             old=old_value,

.../ipywidgets/widgets/widget.py in notify_change(self, change)
    604                 # Send new state to front-end
    605                 self.send_state(key=name)
--> 606         super(Widget, self).notify_change(change)
    607 
    608     def __repr__(self):

.../traitlets/traitlets.py in notify_change(self, change)
   1227     def notify_change(self, change):
   1228         """Notify observers of a change event"""
-> 1229         return self._notify_observers(change)
   1230 
   1231     def _notify_observers(self, event):

.../traitlets/traitlets.py in _notify_observers(self, event)
   1264                 c = getattr(self, c.name)
   1265 
-> 1266             c(event)
   1267 
   1268     def _add_notifiers(self, handler, name, type):

.../jdaviz/configs/default/plugins/collapse/collapse.py in _on_data_item_selected(self, event)
     79 
     80         # Also set the spectral min and max to default to the full range
---> 81         cube = self._selected_data.get_object(cls=SpectralCube)
     82         self.selected_subset = "None"
     83         self.spectral_min = cube.spectral_axis[0].value

.../glue/glue/core/data.py in get_object(self, cls, **kwargs)
    287         handler, _ = data_translator.get_handler_for(cls)
    288 
--> 289         return handler.to_object(self, **kwargs)
    290 
    291     @property

.../glue-astronomy/glue_astronomy/translators/spectral_cube.py in to_object(self, data_or_subset, attribute)
     83             wcs = wcs.sub(subkeep)
     84 
---> 85         return SpectralCube(values, mask=mask, wcs=wcs, meta=data.meta)

.../spectral_cube/spectral_cube.py in __init__(self, data, wcs, mask, meta, fill_value, header, allow_huge_operations, beam, wcs_tolerance, use_dask, **kwargs)
   3532                  wcs_tolerance=0.0, use_dask=False, **kwargs):
   3533 
-> 3534         super(SpectralCube, self).__init__(data=data, wcs=wcs, mask=mask,
   3535                                            meta=meta, fill_value=fill_value,
   3536                                            header=header,

.../spectral_cube/spectral_cube.py in __init__(self, data, wcs, mask, meta, fill_value, header, allow_huge_operations, wcs_tolerance)
    191 
    192         # TODO: mask should be oriented? Or should we assume correctly oriented here?
--> 193         self._data, self._wcs = cube_utils._orient(data, wcs)
    194         self._wcs_tolerance = wcs_tolerance
    195         self._spectral_axis = None

.../spectral_cube/cube_utils.py in _orient(array, wcs)
    145         raise ValueError("Input array must be 3-dimensional")
    146 
--> 147     if wcs.wcs.naxis != 3:
    148         raise ValueError("Input WCS must be 3-dimensional")
    149 

AttributeError: 'PaddedSpectrumWCS' object has no attribute 'wcs'

[pllim]

Wait a minute... Why are we still using SpectralCube in Collapse plugin again? Oh no, does this mean I cannot do spacetelescope/jdaviz#1040 without getting spacetelescope/jdaviz#1006 in first? 😱

[rosteen]

@astrofrog I was working on the route you suggested of using NDCube's CompoundLowLevelWCS for a general solution of padding the WCS needed for additional axes, but I'm unsure about how many of the properties you had coded into your 2D solution need to be kept in my version. Seems like we do need a high level WCS object here, but just using HighLevelWCSWrapper around the CompoundLowLevelWCS isn't sufficient. If you have any advice/comments based on where I'm at so far I'd appreciate it!

[pllim]

After spectral-cube is out of the picture, I get new traceback when I try to give spectral_axis in pixels:

      1 viz = Cubeviz()
----> 2 viz.load_data(sci_arr, data_type='flux', data_label='myarray[SCI]')
      3 viz.load_data(err_arr, data_type='uncert', data_label='myarray[ERR]')
      4 viz.app

.../jdaviz/core/helpers.py in load_data(self, data, parser_reference, **kwargs)
     67 
     68     def load_data(self, data, parser_reference=None, **kwargs):
---> 69         self.app.load_data(data, parser_reference=parser_reference, **kwargs)
     70 
     71     @property

.../jdaviz/app.py in load_data(self, file_obj, parser_reference, **kwargs)
    368                 # If the parser returns something other than known, assume it's
    369                 #  a message we want to make the user aware of.
--> 370                 msg = parser(self, file_obj, **kwargs)
    371 
    372                 if msg is not None:

.../jdaviz/configs/cubeviz/plugins/parsers.py in parse_data(app, file_obj, data_type, data_label)
     96 
     97         if file_obj.ndim == 3:
---> 98             _parse_ndarray_3d(app, file_obj, data_label, data_type)
     99         else:
    100             raise NotImplementedError(f'Unsupported data format: {file_obj}')

.../jdaviz/configs/cubeviz/plugins/parsers.py in _parse_ndarray_3d(app, file_obj, data_label, data_type)
    334     sc = Spectrum1D(spectral_axis=idx, flux=flux)
    335     app.add_data(sc, data_label)
--> 336     _show_data_in_cubeviz_viewer(app, data_label, data_type)

.../jdaviz/configs/cubeviz/plugins/parsers.py in _show_data_in_cubeviz_viewer(app, data_label, data_type)
    259     if data_type == 'flux':
    260         app.add_data_to_viewer(f'{data_type}-viewer', data_label)
--> 261         app.add_data_to_viewer('spectrum-viewer', data_label)
    262     elif data_type in ('uncert', 'mask'):
    263         app.add_data_to_viewer(f'{data_type}-viewer', data_label)

.../jdaviz/app.py in add_data_to_viewer(self, viewer_reference, data_path, clear_other_data, ext)
    808         if data_id is not None:
    809             data_ids.append(data_id)
--> 810             self._update_selected_data_items(viewer_item['id'], data_ids)
    811         else:
    812             raise ValueError(

.../jdaviz/app.py in _update_selected_data_items(self, viewer_id, selected_items)
   1142                                               viewer_id=viewer_id,
   1143                                               sender=self)
-> 1144             self.hub.broadcast(add_data_message)
   1145 
   1146         # Remove any deselected data objects from viewer

.../glue/core/hub.py in broadcast(self, message)
    213             logging.getLogger(__name__).info("Broadcasting %s", message)
    214             for subscriber, handler in self._find_handlers(message):
--> 215                 handler(message)
    216 
    217     def __getstate__(self):

.../jdaviz/configs/cubeviz/plugins/slice/slice.py in _on_data_added(self, msg)
     69                 self._indicator_viewers.append(msg.viewer)
     70                 # cache wavelengths so that wavelength <> slice conversion can be done efficiently
---> 71                 x_all = msg.viewer.data()[0].spectral_axis
     72                 self._x_all = x_all.value
     73                 self.wavelength_unit = str(x_all.unit)

.../jdaviz/configs/specviz/plugins/viewers.py in data(self, cls)
     78                         # If spectrum, collapse via the defined statistic
     79                         if _class == Spectrum1D:
---> 80                             layer_data = layer_state.layer.get_object(cls=_class,
     81                                                                       statistic=statistic)
     82                         else:

.../glue/core/data.py in get_object(self, cls, **kwargs)
    295         handler, _ = data_translator.get_handler_for(cls)
    296 
--> 297         return handler.to_object(self, **kwargs)
    298 
    299     @property

.../glue_astronomy/translators/spectrum1d.py in to_object(self, data_or_subset, attribute, statistic)
    233             [attribute] if not hasattr(attribute, '__len__') else attribute)
    234 
--> 235         return Spectrum1D(**data_kwargs, **kwargs)

.../specutils/spectra/spectrum1d.py in __init__(self, flux, spectral_axis, wcs, velocity_convention, rest_value, redshift, radial_velocity, bin_specification, **kwargs)
    286                     spec_axis = self.wcs.spectral.pixel_to_world(np.arange(self.flux.shape[-1]))
    287             else:
--> 288                 spec_axis = self.wcs.pixel_to_world(np.arange(self.flux.shape[-1]))
    289 
    290             try:

.../specutils/utils/wcs_utils.py in pixel_to_world(self, *args, **kwargs)
    217             if orig_array.unit == '':
    218                 return u.Quantity(super().pixel_to_world_values(*args, **kwargs))
--> 219             return super().pixel_to_world(*args, **kwargs).to(
    220                 orig_array.unit, equivalencies=u.spectral())
    221 

.../gwcs/api.py in pixel_to_world(self, *pixel_arrays)
    297         """
    298         pixels = self._sanitize_pixel_inputs(*pixel_arrays)
--> 299         return self(*pixels, with_units=True)
    300 
    301     def array_index_to_world(self, *index_arrays):

.../gwcs/wcs.py in __call__(self, *args, **kwargs)
    378         if with_units:
    379             if self.output_frame.naxes == 1:
--> 380                 result = self.output_frame.coordinates(result)
    381             else:
    382                 result = self.output_frame.coordinates(*result)

.../gwcs/coordinate_frames.py in coordinates(self, *args)
    461         else:
    462             if hasattr(args[0], 'unit'):
--> 463                 return coord.SpectralCoord(*args).to(self.unit[0])
    464             else:
    465                 return coord.SpectralCoord(*args, self.unit[0])

.../astropy/units/decorators.py in wrapper(*func_args, **func_kwargs)
    302             # Call the original function with any equivalencies in force.
    303             with add_enabled_equivalencies(self.equivalencies):
--> 304                 return_ = wrapped_function(*func_args, **func_kwargs)
    305 
    306             # Return

.../astropy/coordinates/spectral_coordinate.py in __new__(cls, value, unit, observer, target, radial_velocity, redshift, **kwargs)
    191                 **kwargs):
    192 
--> 193         obj = super().__new__(cls, value, unit=unit, **kwargs)
    194 
    195         # There are two main modes of operation in this class. Either the

.../astropy/coordinates/spectral_quantity.py in __new__(cls, value, unit, doppler_rest, doppler_convention, **kwargs)
     55                 **kwargs):
     56 
---> 57         obj = super().__new__(cls, value, unit=unit, **kwargs)
     58 
     59         # If we're initializing from an existing SpectralQuantity, keep any

.../astropy/units/quantity.py in __new__(cls, value, unit, dtype, copy, order, subok, ndmin)
    422             if type(value) is not cls and not (subok and
    423                                                isinstance(value, cls)):
--> 424                 value = value.view(cls)
    425 
    426             if dtype is None and value.dtype.kind in 'iu':

.../astropy/coordinates/spectral_coordinate.py in __array_finalize__(self, obj)
    240 
    241     def __array_finalize__(self, obj):
--> 242         super().__array_finalize__(obj)
    243         self._radial_velocity = getattr(obj, '_radial_velocity', None)
    244         self._observer = getattr(obj, '_observer', None)

.../astropy/coordinates/spectral_quantity.py in __array_finalize__(self, obj)
     70 
     71     def __array_finalize__(self, obj):
---> 72         super().__array_finalize__(obj)
     73         self._doppler_rest = getattr(obj, '_doppler_rest', None)
     74         self._doppler_convention = getattr(obj, '_doppler_convention', None)

.../astropy/units/quantity.py in __array_finalize__(self, obj)
    547             unit = getattr(obj, '_unit', None)
    548             if unit is not None:
--> 549                 self._set_unit(unit)
    550 
    551         # Copy info if the original had `info` defined.  Because of the way the

.../astropy/units/quantity.py in _set_unit(self, unit)
   1929     def _set_unit(self, unit):
   1930         if unit is None or not unit.is_equivalent(self._equivalent_unit):
-> 1931             raise UnitTypeError(
   1932                 "{} instances require units equivalent to '{}'"
   1933                 .format(type(self).__name__, self._equivalent_unit) +

UnitTypeError: SpectralCoord instances require units equivalent to '(Unit("Hz"), Unit("m"), Unit("J"), Unit("1 / m"), Unit("km / s"))', so cannot set it to 'pix'.

[pllim]

Okay. I got it. With this PR, it works if I pass in only the flux like this Spectrum1D(flux=np.random.random((10, 10, 5))), where 5 is the length of the spectral axis. If I pass in spectral_axis explicitly, it crashes. Does the shape makes sense to you? Is spectral axis supposed to be in arr.shape[2]? I still get confused by all the swapping that happens (astropy/specutils#923).

https://github.com/astropy/specutils/blob/7f7b55272b7e36d827c130bb2bfa0ec6078f41a3/specutils/spectra/spectrum1d.py#L265-L269

[dhomeier]

https://github.com/glue-viz/glue-astronomy/blob/7f3229d01bdc2bff64fdb1a83738447cf696c59d/glue_astronomy/translators/tests/test_spectrum1d.py#L204-#234 will need to be changed to flux.T/spec.flux.T if that is the intended behaviour now. Still feels a bit weird...

[pllim]

Is there a way to push this forward? I need this for spacetelescope/jdaviz#1040 .

[dhomeier]

I could go ahead with adapting the tests, but @astrofrog and I are quite uncertain at this point what the logic of reordering axes here again is. specutils as far as I see it is consistent in using the last dimension as spectral axis, so it's rather confusing to have the glue loader swap this again, and (only?) for [G]WCS spectra.

[pllim]

In our FITS data, the spectral axis seems to be in NAXIS3, which in Python-speak, is axis=0, so maybe it is specutils who is putting it in the wrong end. Though perhaps that is necessary for performance reason? 🤷

[dhomeier]

NAXIS3 -> axis=0 is just the conversion from Fortran- to C-ordering, right? https://github.com/astropy/specutils/blob/af729120269e776e40d8a5e6c2e270617d7a8d6e/specutils/spectra/spectrum1d.py#L199-L205 notes that for the WCS the order is turned around again, but I think that is simply because by convention the WCS is kept in Fortran order. While the two have to match, I so not see a technical reason why the spectral axis has to be last, and performance-wise it would certainly be better to avoid those np.swapaxes. I suppose the spectral axis was simply found more frequently in NAXIS1, and to avoid even more confusion that was set as a rule. According to the comment in the data_translator the convention for glue is exactly the opposite, so I guess swapping the axes again (and fixing the tests accordingly) is unavoidable at this point. If the "glue-order" is common for large data cube, I think it is worth considering to allow that in specutils as well (e.g. MaNGA also uses NAXIS3, and the axes swapping seems to introduce a quite noticeable delay).

[pllim]

Yes, I think @rosteen mentioned that it would be nice if specutils is flexible enough to allow specutils in whatever axis but that's unlikely to happen any time soon. :(

[dhomeier]

@rosteen I have changed the expected labels in the test for now, so it can progress further, but does not necessarily mean they should remain like this. The current failure however raises already when just dereferencing data['World 1'] as an array-like object ( in the call to _return_single_array), which seems to indicate something is still more fundamentally broken in the implementation (or is exposing a new bug in astropy.wcs?).

[pllim]

What does this error mean?

        assert_equal(data['World 0'], [[10, 10, 10], [10, 10, 10]])
>       assert data['World 1'].shape == (2, 3)

glue_astronomy/translators/tests/test_spectrum1d.py:255: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.tox/py38-test/lib/python3.8/site-packages/glue/core/data.py:571: in __getitem__
    return self.get_data(key, view=view)
.tox/py38-test/lib/python3.8/site-packages/glue/core/data.py:1366: in get_data
    result = comp.data
.tox/py38-test/lib/python3.8/site-packages/glue/core/component.py:215: in data
    return self._calculate()
.tox/py38-test/lib/python3.8/site-packages/glue/core/component.py:313: in _calculate
    world_coords = pixel2world_single_axis(self._data.coords,
.tox/py38-test/lib/python3.8/site-packages/glue/core/coordinate_helpers.py:56: in pixel2world_single_axis
    result = wcs.pixel_to_world_values(*pixel)
.tox/py38-test/lib/python3.8/site-packages/ndcube/wcs/wrappers/compound_wcs.py:112: in pixel_to_world_values
    world_arrays_sub = w.pixel_to_world_values(*pixel_arrays_sub)
.tox/py38-test/lib/python3.8/site-packages/astropy/wcs/wcsapi/fitswcs.py:322: in pixel_to_world_values
    world = self.all_pix2world(*pixel_arrays, 0)

[dhomeier]

Per my speculation above, it's two levels further down ;-)

../../opt/lib/python3.9/site-packages/astropy/wcs/wcs.py:1351: in all_pix2world
    return self._array_converter(
../../opt/lib/python3.9/site-packages/astropy/wcs/wcs.py:1328: in _array_converter
    return _return_single_array(xy, origin)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

xy = array([[0, 0, 0]]), origin = 0

    def _return_single_array(xy, origin):
        if xy.shape[-1] != self.naxis:
>           raise ValueError(
                "When providing two arguments, the array must be "
                "of shape (N, {})".format(self.naxis))
E           ValueError: When providing two arguments, the array must be of shape (N, 1)

This apparently already when trying to access the shape attribute. Looks wicked. Or is it just that the shape is inconsistent with the no. of dimensions somewhere else?

[pllim]

Definitely something is still wonky here, unfortunately. Even if you ignore the data['World 1'] parts of the test (it seems to go into ndcube at some point and I got lost there), it crashes later at:

>>> s, o = data.coords.pixel_to_world(1, 2)
AttributeError: 'PaddedSpectrumWCS' object has no attribute 'pixel_to_world'

and

>>> px, py = data.coords.world_to_pixel(s, o)
AttributeError: 'PaddedSpectrumWCS' object has no attribute 'world_to_pixel'

The spec_new stuff works.

[pllim]

Maybe subclassing from CompoundLowLevelWCS is too low of a level?

[dhomeier]

Maybe subclassing from CompoundLowLevelWCS is too low of a level?

Yea, the pixel_to_world etc. are only provided by the HighLevel version, but that in turn is not supported by gwcs afair. Dunno if it can all be moved to the pixel_to_world_values level...

[rosteen]

As per PO feedback at sprint review today, they want the spectral axis to be dimensionless rather than pixels. Is it possible to do that upstream here? If not, I'll have to hack something up downstream in Jdaviz.

I seem to remember a decision specifically not to handle dimensionless "spectral" axes in specutils, with the compromise that we'd handle pixel units. I could be remembering wrong though, I'll look back through some notes/git discussion.

Btw, thanks @dhomeier and @pllim for looking at this, I left it in a state where I'd run into the limit of my WCS knowledge. I'm trying to find time this sprint to get back to it.

[pllim]

dimensionless "spectral" axes

Okay, just keep it as pixel here. I dealt with that part downstream. Thanks!

[pllim]

This can be closed without merged. Superseded by #68