Update Cubeviz to work on data with spectral axis last

CHANGES.rst

@@ -53,8 +53,9 @@ Other Changes and Additions
 
 - Change default collapse function to sum.
   This affects collapsed spectrum in Cubeviz and its Collapse plugin default. [#1229, #1237]
-
 - Data dropdowns in plugins are now filtered to only applicable entries. [#1221]
+- Cube data now has spectral axis last in the backend, to match specutils Spectrum1D
+  axis order and work with updated glue-astronomy translators. [#1174]
 
 2.4 (2022-03-29)
 ================

jdaviz/app.py

@@ -833,7 +833,10 @@ def add_data_to_viewer(self, viewer_reference, data_path,
         data_label = self._build_data_label(data_path, ext=ext)
         data_id = self._data_id_from_label(data_label)
 
-        data_ids = viewer_item['selected_data_items'] if not clear_other_data else []
+        if clear_other_data:
+            self._update_selected_data_items(viewer_item['id'], [])
+
+        data_ids = viewer_item['selected_data_items']
 
         if data_id is not None:
             data_ids.append(data_id)

jdaviz/configs/cubeviz/helper.py

@@ -3,7 +3,9 @@
 from jdaviz.core.helpers import ConfigHelper
 from jdaviz.configs.default.plugins.line_lists.line_list_mixin import LineListMixin
 from jdaviz.configs.specviz import Specviz
-from jdaviz.core.events import SliceSelectWavelengthMessage, SliceSelectSliceMessage
+from jdaviz.core.events import (AddDataMessage,
+                                SliceSelectWavelengthMessage,
+                                SliceSelectSliceMessage)
 
 __all__ = ['Cubeviz', 'CubeViz']
 
@@ -16,6 +18,22 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.app.hub.subscribe(self, SliceSelectWavelengthMessage,
                                handler=self.select_wavelength)
+        self.app.hub.subscribe(self, AddDataMessage,
+                               handler=self._set_spectrum_x_axis)
+
+    def _set_spectrum_x_axis(self, msg):
+        if msg.viewer_id != "cubeviz-3":
+            return
+        viewer = self.app.get_viewer("spectrum-viewer")
+        ref_data = viewer.state.reference_data
+        if ref_data and ref_data.ndim == 3:
+            for att_name in ["Wave", "Wavelength", "Freq", "Frequency",
+                             "Wavenumber", "Velocity", "Energy"]:
+                if att_name in ref_data.component_ids():
+                    viewer.state.x_att = ref_data.id[att_name]
+                    break
+            else:
+                viewer.state.x_att_pixel = ref_data.id["Pixel Axis 2 [x]"]
 
     def select_slice(self, slice):
         """

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -4,6 +4,7 @@
 import numpy as np
 from astropy import units as u
 from astropy.io import fits
+from astropy.time import Time
 from astropy.wcs import WCS
 from specutils import Spectrum1D
 
@@ -133,13 +134,13 @@ def _parse_hdu(app, hdulist, file_name=None):
             app.add_data_to_viewer('uncert-viewer', data_label)
 
         if any(x in hdu.name.lower() for x in EXT_TYPES['flux']):
+            # Add flux to top left image viewer
             app.add_data_to_viewer('flux-viewer', data_label)
+            # Add flux to spectrum viewer
             app.add_data_to_viewer('spectrum-viewer', data_label)
 
 
 def _parse_jwst_s3d(app, hdulist, data_label, ext='SCI', viewer_name='flux-viewer'):
-    from specutils import Spectrum1D
-
     # Manually inject MJD-OBS until we can support GWCS, see
     # https://github.com/spacetelescope/jdaviz/issues/690 and
     # https://github.com/glue-viz/glue-astronomy/issues/59
@@ -150,7 +151,6 @@ def _parse_jwst_s3d(app, hdulist, data_label, ext='SCI', viewer_name='flux-viewe
                     hdulist[ext].header['MJD-OBS'] = hdulist[ext].header[key]
                     break
                 else:
-                    from astropy.time import Time
                     t = Time(hdulist[ext].header[key])
                     hdulist[ext].header['MJD-OBS'] = t.mjd
                     break

jdaviz/configs/cubeviz/plugins/slice/slice.py

@@ -55,14 +55,19 @@ def _on_linked_changed(self, event):
                                           self._viewer_slices_changed)
 
     def _on_data_added(self, msg):
-        if len(msg.data.shape) == 3 and isinstance(msg.viewer, BqplotImageView):
-            self.max_value = msg.data.shape[0] - 1
+        if isinstance(msg.viewer, BqplotImageView):
+            if len(msg.data.shape) == 3:
+                self.max_value = msg.data.shape[-1] - 1
 
-            if msg.viewer not in self._watched_viewers:
-                self._watched_viewers.append(msg.viewer)
+                if msg.viewer not in self._watched_viewers:
+                    self._watched_viewers.append(msg.viewer)
 
-                msg.viewer.state.add_callback('slices',
-                                              self._viewer_slices_changed)
+                    msg.viewer.state.add_callback('slices',
+                                                  self._viewer_slices_changed)
+
+            else:
+                if msg.viewer in self._watched_viewers:
+                    self._watched_viewers.remove(msg.viewer)
 
         elif isinstance(msg.viewer, BqplotProfileView):
             if msg.viewer not in self._indicator_viewers:
@@ -98,8 +103,8 @@ def _viewer_slices_changed(self, value):
         # so we'll update the slider to match which will trigger
         # the slider observer (_on_slider_updated) and sync across
         # any other applicable viewers
-        if len(value):
-            self.slider = float(value[0])
+        if len(value) == 3:
+            self.slider = float(value[-1])
 
     def _on_select_slice_message(self, msg):
         # NOTE: by setting the slice index, the observer (_on_slider_updated)
@@ -135,6 +140,6 @@ def _on_slider_updated(self, event):
 
         if self.linked:
             for viewer in self._watched_viewers:
-                viewer.state.slices = (value, 0, 0)
+                viewer.state.slices = (0, 0, value)
             for viewer in self._indicator_viewers:
                 viewer._update_slice_indicator(value)

jdaviz/configs/cubeviz/plugins/viewers.py

@@ -35,6 +35,20 @@ class CubevizImageView(BqplotImageView, JdavizViewerMixin):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._initialize_toolbar_nested()
+        self.state.add_callback('reference_data', self._initial_x_axis)
+
+    def _initial_x_axis(self, *args):
+        # Make sure that the x_att is correct on data load
+        ref_data = self.state.reference_data
+        if ref_data and ref_data.ndim == 3:
+            for att_name in ["Right Ascension", "RA", "Galactic Longitude"]:
+                if att_name in ref_data.component_ids():
+                    x_att = att_name
+                    self.state.x_att_world = ref_data.id[x_att]
+                    break
+            else:
+                x_att = "Pixel Axis 0 [z]"
+                self.state.x_att = ref_data.id[x_att]
 
     def set_plot_axes(self):
         self.figure.axes[1].tick_format = None

jdaviz/configs/default/plugins/collapse/collapse.py

@@ -91,11 +91,9 @@ def _link_collapse_data(self):
                     self.app.data_collection[i].meta["Plugin"] == "Collapse"):
                 links = [LinkSame(pc_old[0], pc_new[0]),
                          LinkSame(pc_old[1], pc_new[1])]
-
-            # Else, link collapse data to cube (pc_old)
             else:
-                links = [[LinkSame(pc_new[1], pc_old[2]),
-                          LinkSame(pc_new[0], pc_old[1])]]
+                links = [LinkSame(pc_old[0], pc_new[1]),
+                         LinkSame(pc_old[1], pc_new[0])]
 
             self.app.data_collection.add_link(links)
 

jdaviz/configs/default/plugins/collapse/tests/test_collapse.py

@@ -23,5 +23,5 @@ def test_linking_after_collapse(cubeviz_helper, spectral_cube_wcs):
     assert dc[1].label == 'Collapsed 1 Unknown spectrum object[FLUX]'
     assert len(dc.external_links) == 2
 
-    assert dc.external_links[1].cids1[0] is dc[1].pixel_component_ids[0]
-    assert dc.external_links[1].cids2[0] is dc[0].pixel_component_ids[1]
+    assert dc.external_links[1].cids1[0] is dc[0].pixel_component_ids[1]
+    assert dc.external_links[1].cids2[0] is dc[1].pixel_component_ids[0]

jdaviz/configs/default/plugins/gaussian_smooth/gaussian_smooth.py

@@ -54,7 +54,8 @@ def _on_data_selected(self, event):
             return
 
         # NOTE: if this is ever used anywhere else, it should be moved into DatasetSelect
-        self.selected_data_is_1d = len(self.dataset.selected_dc_item.data.shape) == 1
+        if self.dataset.selected_dc_item is not None:
+            self.selected_data_is_1d = len(self.dataset.selected_dc_item.data.shape) == 1
 
     def vue_spectral_smooth(self, *args, **kwargs):
         # Testing inputs to make sure putting smoothed spectrum into

jdaviz/configs/default/plugins/model_fitting/fitting_backend.py

@@ -7,7 +7,7 @@
 
 import astropy.units as u
 
-from specutils.spectra import Spectrum1D
+from specutils import Spectrum1D
 from specutils.fitting import fit_lines
 
 __all__ = ['fit_model_to_spectrum']
@@ -189,7 +189,7 @@ def collect_result(results):
 
     # Build output 3D spectrum
     funit = spectrum.flux.unit
-    output_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis,
+    output_spectrum = Spectrum1D(wcs=spectrum.wcs,
                                  flux=output_flux_cube * funit)
 
     return fitted_models, output_spectrum
@@ -231,7 +231,7 @@ def __call__(self):
             # spectrum reference into the callable somehow prevents it
             # to execute. This behavior was seen also with other functions
             # passed to the callable.
-            flux = self.cube[x, y, :]  # transposed!
+            flux = self.cube[x, y, :]
             sp = Spectrum1D(spectral_axis=self.wave, flux=flux)
 
             fitted_model = fit_lines(sp, self.model, window=self.window)

jdaviz/configs/default/plugins/model_fitting/model_fitting.py

@@ -2,7 +2,6 @@
 import numpy as np
 
 import astropy.units as u
-from astropy.wcs import WCSSUB_SPECTRAL
 from specutils import Spectrum1D
 from specutils.utils import QuantityModel
 from traitlets import Bool, List, Unicode, observe
@@ -441,17 +440,7 @@ def vue_fit_model_to_cube(self, *args, **kwargs):
             return
 
         # Get the primary data component
-        attribute = data.main_components[0]
-        component = data.get_component(attribute)
-        temp_values = data.get_data(attribute)
-
-        # Transpose the axis order
-        values = np.moveaxis(temp_values, 0, -1) * u.Unit(component.units)
-
-        # We manually create a Spectrum1D object from the flux information
-        #  in the cube we select
-        wcs = data.coords.sub([WCSSUB_SPECTRAL])
-        spec = Spectrum1D(flux=values, wcs=wcs)
+        spec = data.get_object(Spectrum1D, statistic=None)
 
         # TODO: in vuetify >2.3, timeout should be set to -1 to keep open
         #  indefinitely
@@ -483,9 +472,6 @@ def vue_fit_model_to_cube(self, *args, **kwargs):
             temp_label = "{} ({}, {})".format(self.model_label, m["x"], m["y"])
             self.app.fitted_models[temp_label] = m["model"]
 
-        # Transpose the axis order back
-        values = np.moveaxis(fitted_spectrum.flux.value, -1, 0)
-
         count = max(map(lambda s: int(next(iter(re.findall(r"\d$", s)), 0)),
                         self.data_collection.labels)) + 1
 
@@ -494,9 +480,8 @@ def vue_fit_model_to_cube(self, *args, **kwargs):
         # Create new glue data object
         output_cube = Data(label=label,
                            coords=data.coords)
-        output_cube['flux'] = values
-        output_cube.get_component('flux').units = \
-            fitted_spectrum.flux.unit.to_string()
+        output_cube['flux'] = fitted_spectrum.flux.value
+        output_cube.get_component('flux').units = fitted_spectrum.flux.unit.to_string()
 
         # Add to data collection, tracking that it came from this plugin to be filtered out
         # from dataset-select dropdowns
@@ -547,7 +532,7 @@ def vue_register_spectrum(self, event):
         # Link the result spectrum to the reference data of the spectrum viewer
 
         ref_data = self.app.get_viewer('spectrum-viewer').state.reference_data
-        data_id = ref_data.world_component_ids[0]
+        data_id = ref_data.world_component_ids[-1]
         model_id = self.app.session.data_collection[label].world_component_ids[0]
         self.app.session.data_collection.add_link(LinkSame(data_id, model_id))
 

jdaviz/conftest.py

@@ -52,7 +52,7 @@ def image_2d_wcs(request):
 def spectral_cube_wcs(request):
     # A simple spectral cube WCS used by some tests
     wcs = WCS(naxis=3)
-    wcs.wcs.ctype = 'RA---TAN', 'DEC--TAN', 'FREQ'
+    wcs.wcs.ctype = 'FREQ', 'DEC--TAN', 'RA---TAN'
     wcs.wcs.set()
     return wcs
 
@@ -61,7 +61,7 @@ def spectral_cube_wcs(request):
 def spectrum1d_cube_wcs(request):
     # A simple spectrum1D WCS used by some tests
     wcs = WCS(naxis=3)
-    wcs.wcs.ctype = 'RA---TAN', 'DEC--TAN', 'WAVE-LOG'
+    wcs.wcs.ctype = 'WAVE-LOG', 'DEC--TAN', 'RA---TAN'
     wcs.wcs.set()
     return wcs