Cubeviz parser to load all sorts of s3d cubes

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -3,10 +3,11 @@
 
 import numpy as np
 from astropy.io import fits
-from spectral_cube import SpectralCube
-from spectral_cube.io.fits import FITSReadError
 from specutils import Spectrum1D
 
+from astropy.wcs import WCS
+from astropy import units as u
+
 from jdaviz.core.registries import data_parser_registry
 
 __all__ = ['parse_data']
@@ -49,13 +50,27 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
         file_name = os.path.basename(file_obj)
 
         with fits.open(file_obj) as hdulist:
-            _parse_hdu(app, hdulist, file_name=data_label or file_name)
+            prihdr = hdulist[0].header
+            telescop = prihdr.get('TELESCOP', '').lower()
+            filetype = prihdr.get('FILETYPE', '').lower()
+            if telescop == 'jwst' and filetype == '3d ifu cube':
+                from glue.core import Data
+                unit = u.Unit(hdulist[1].header.get('BUNIT', 'count'))
+                flux = hdulist[1].data << unit
+                wcs = WCS(hdulist[1].header, hdulist)
+                data = Data(flux=flux, coords=wcs)  # Spectrum1D too slow
+                data_label = f'{file_obj}[SCI]'
+                app.add_data(data, data_label)
+                app.add_data_to_viewer('flux-viewer', data_label)
+                app.add_data_to_viewer('spectrum-viewer', data_label)
+            else:
+                _parse_hdu(app, hdulist, file_name=data_label or file_name)
 
     # If the data types are custom data objects, use explicit parsers. Note
     #  that this relies on the glue-astronomy machinery to turn the data object
     #  into something glue can understand.
-    elif isinstance(file_obj, SpectralCube):
-        _parse_spectral_cube(app, file_obj, data_type or 'flux', data_label)
+    elif isinstance(file_obj, Spectrum1D) and len(file_obj.shape) == 3:
+        _parse_spectrum1d_3d(app, file_obj)
     elif isinstance(file_obj, Spectrum1D):
         _parse_spectrum1d(app, file_obj)
 
@@ -67,49 +82,34 @@ def _parse_hdu(app, hdulist, file_name=None):
 
     file_name = file_name or "Unknown HDU object"
 
-    # WCS may only exist in a single extension (in this case, just the flux
-    #  flux extension), so first find and store then wcs information.
-    wcs = None
-
-    for hdu in hdulist:
-        if hdu.data is None or not hdu.is_image:
-            continue
-
-        try:
-            sc = SpectralCube.read(hdu, format='fits')
-        except (ValueError, FITSReadError):
-            continue
-        else:
-            wcs = sc.wcs
-
     # Now loop through and attempt to parse the fits extensions as spectral
     #  cube object. If the wcs fails to parse in any case, use the wcs
     #  information we scraped above.
     for hdu in hdulist:
         data_label = f"{file_name}[{hdu.name}]"
 
-        if hdu.data is None or not hdu.is_image:
+        if hdu.data is None or not hdu.is_image or len(hdu.data.shape) != 3:
             continue
 
-        # This is supposed to fail on attempting to load anything that
-        # isn't cube-shaped. But it's not terribly reliable
+        wcs = WCS(hdu.header, hdulist)
+
+        try:
+            flux_unit = u.Unit(hdu.header['BUNIT'])
+        except KeyError:
+            logging.warning("No flux units found in hdu, using u.count as a stand-in")
+            flux_unit = u.count
+        finally:
+            flux = hdu.data * flux_unit
+
+        # flux = np.moveaxis(flux, 1, 2)
+
         try:
-            sc = SpectralCube.read(hdu, format='fits')
-        except (ValueError, OSError):
-            # This will fail if the parsing of the wcs does not provide
-            # proper celestial axes
-            try:
-                hdu.header.update(wcs.to_header())
-                sc = SpectralCube.read(hdu)
-            except (ValueError, AttributeError) as e:
-                logging.warning(e)
-                continue
-        except FITSReadError as e:
+            sc = Spectrum1D(flux=flux, wcs=wcs)
+            app.data_collection[data_label] = sc
+        except Exception as e:
             logging.warning(e)
             continue
 
-        app.add_data(sc, data_label)
-
         # If the data type is some kind of integer, assume it's the mask/dq
         if hdu.data.dtype in (int, np.uint, np.uint32) or \
                 any(x in hdu.name.lower() for x in EXT_TYPES['mask']):
@@ -123,22 +123,52 @@ def _parse_hdu(app, hdulist, file_name=None):
             app.add_data_to_viewer('flux-viewer', data_label)
             app.add_data_to_viewer('spectrum-viewer', data_label)
 
+    # _fix_axes(app)
+
+
+def _fix_axes(app):
+    # Calibrate viewers to work with the updated Spectrum1D that can handle
+    # cubes. This needs to be done because the shape is different from
+    # SpectralCube, which was used previously.
+    # Setting the y and x axes needs to happen in this order to avoid an
+    # error for setting both axes to the same value.
+    app.get_viewer("mask-viewer").viewer_options.y_att_world_selected = 1
+    app.get_viewer("mask-viewer").viewer_options.x_att_world_selected = 0
+    app.get_viewer("uncert-viewer").viewer_options.y_att_world_selected = 1
+    app.get_viewer("uncert-viewer").viewer_options.x_att_world_selected = 0
+    app.get_viewer("flux-viewer").viewer_options.y_att_world_selected = 1
+    app.get_viewer("flux-viewer").viewer_options.x_att_world_selected = 0
+    app.get_viewer("spectrum-viewer").viewer_options.x_att_selected = 5
+
+
+def _parse_spectrum1d_3d(app, file_obj):
+    # Load spectrum1d as a cube
+
+    for attr in ["flux", "mask", "uncertainty"]:
+        if attr == "mask":
+            flux = getattr(file_obj, attr) * file_obj.flux.unit
+        elif attr == "uncertainty":
+            flux = getattr(file_obj, attr)
+            flux = u.Quantity(flux.array) * file_obj.flux.unit
+        else:
+            flux = getattr(file_obj, attr)
+
+        flux = np.moveaxis(flux, 1, 0)
 
-def _parse_spectral_cube(app, file_obj, data_type='flux', data_label=None):
-    data_label = data_label or f"Unknown spectral cube[{data_type.upper()}]"
+        s1d = Spectrum1D(flux=flux, wcs=file_obj.wcs)
 
-    app.add_data(file_obj, data_label)
+        data_label = f"Unknown spectrum object[{attr}]"
+        app.data_collection[data_label] = s1d
 
-    if data_type == 'flux':
-        app.add_data_to_viewer('flux-viewer', f"{data_label}")
-        app.add_data_to_viewer('spectrum-viewer', f"{data_label}")
-    elif data_type == 'mask':
-        app.add_data_to_viewer('mask-viewer', f"{data_label}")
-    elif data_type == 'uncert':
-        app.add_data_to_viewer('uncert-viewer', f"{data_label}")
+        if attr == 'flux':
+            app.add_data_to_viewer('flux-viewer', f"{data_label}")
+            app.add_data_to_viewer('spectrum-viewer', f"{data_label}")
+        elif attr == 'mask':
+            app.add_data_to_viewer('mask-viewer', f"{data_label}")
+        elif attr == 'uncertainty':
+            app.add_data_to_viewer('uncert-viewer', f"{data_label}")
 
-    # TODO: SpectralCube does not store mask information
-    # TODO: SpectralCube does not store data quality information
+    # _fix_axes(app)
 
 
 def _parse_spectrum1d(app, file_obj):

jdaviz/configs/cubeviz/plugins/tests/test_parsers.py

@@ -1,12 +1,9 @@
-import os
-
 import astropy.units as u
 import numpy as np
 import pytest
 from astropy.io import fits
 from astropy.nddata import StdDevUncertainty
 from astropy.wcs import WCS
-from spectral_cube import SpectralCube
 from specutils import Spectrum1D
 
 from jdaviz.app import Application
@@ -49,28 +46,14 @@ def image_hdu_obj():
     return fits.HDUList([fits.PrimaryHDU(), flux_hdu, mask_hdu, uncert_hdu])
 
 
-@pytest.mark.filterwarnings('ignore:.* contains multiple slashes')
+@pytest.mark.filterwarnings('ignore')
 def test_fits_image_hdu_parse(image_hdu_obj, cubeviz_app):
     cubeviz_app.load_data(image_hdu_obj)
 
     assert len(cubeviz_app.data_collection) == 3
     assert cubeviz_app.data_collection[0].label.endswith('[FLUX]')
 
 
-@pytest.mark.filterwarnings('ignore:.* contains multiple slashes')
-def test_spectral_cube_parse(tmpdir, image_hdu_obj, cubeviz_app):
-    f = tmpdir.join("test_fits_image.fits")
-    path = os.path.join(f.dirname, f.basename)
-    image_hdu_obj.writeto(path)
-
-    sc = SpectralCube.read(path, hdu=1)
-
-    cubeviz_app.load_data(sc)
-
-    assert len(cubeviz_app.data_collection) == 1
-    assert cubeviz_app.data_collection[0].label.endswith('[FLUX]')
-
-
 def test_spectrum1d_parse(image_hdu_obj, cubeviz_app):
     spec = Spectrum1D(flux=np.random.sample(10) * u.Jy,
                       spectral_axis=np.arange(10) * u.nm,

jdaviz/configs/default/plugins/line_lists/line_lists.py

@@ -95,7 +95,7 @@ def _on_viewer_data_changed(self, msg=None):
 
         try:
             viewer_data = self.app.get_viewer('spectrum-viewer').data()
-        except TypeError:
+        except Exception:
             warn_message = SnackbarMessage("Line list plugin could not retrieve data from viewer",
                                            sender=self, color="error")
             self.hub.broadcast(warn_message)

jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py

@@ -74,8 +74,11 @@ def _on_viewer_data_changed(self, msg=None):
         if msg is not None and msg.viewer_id != self._viewer_id:
             return
 
-        self._viewer_data = self.app.get_data_from_viewer('spectrum-viewer',
-                                                          include_subsets=False)
+        try:
+            self._viewer_data = self.app.get_data_from_viewer('spectrum-viewer',
+                                                              include_subsets=False)
+        except Exception:  # Some data cannot be parsed as Spectrum1D
+            return
 
         self.dc_items = [data.label
                          for data in self.app.data_collection