Resolving traceback bug in cubeviz spectral extraction (spacetelesco…

…pe#2646)

* update cubeviz parsers to save metadata for flux/uncert data, spec extract no longer assumes 1 set of data loaded, spec extract tests failing

* cube_data_type saved as single metadata, string used instead of bool

* _parse_spectrum1d_3d needed uncertainty instead of uncert, cleaned in spec_extract and conditional added so subsets are ignored

* updated change log

* changed from release 3.9 to 3.8.2

* removed conditionals in parser, and added one in spec_extract to counteract mask data

* updated methodology to complexity

* add test that runs gaussian smooth before spectral extraction

* added line to test to ensure gs_cube added to viewer

* disable spec extract until data is loaded into cubeviz, updated test for this case and after data is loaded

* typo in disable message fixed

---------

Co-authored-by: Gilbert Green <[email protected]>
Co-authored-by: Gilbert Green <[email protected]>

CHANGES.rst

@@ -84,6 +84,8 @@ Bug Fixes
 
 Cubeviz
 ^^^^^^^
+- Fixes Spectral Extraction's assumptions of one data per viewer, and flux data only in
+  flux-viewer/uncertainty data only in uncert-viewer. [#2646]
 
 Imviz
 ^^^^^

jdaviz/configs/cubeviz/helper.py

@@ -22,6 +22,9 @@ class Cubeviz(ImageConfigHelper, LineListMixin):
     _default_flux_viewer_reference_name = "flux-viewer"
     _default_image_viewer_reference_name = "image-viewer"
 
+    _loaded_flux_cube = None
+    _loaded_uncert_cube = None
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.app.hub.subscribe(self, AddDataMessage,

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -13,6 +13,7 @@
 from jdaviz.core.registries import data_parser_registry
 from jdaviz.utils import standardize_metadata, PRIHDR_KEY
 
+
 __all__ = ['parse_data']
 
 EXT_TYPES = dict(flux=['flux', 'sci', 'data'],
@@ -58,6 +59,7 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
             spectrum_viewer_reference_name=spectrum_viewer_reference_name,
             uncert_viewer_reference_name=uncert_viewer_reference_name
         )
+        app.get_tray_item_from_name("Spectral Extraction").disabled_msg = ""
     elif isinstance(file_obj, str):
         if file_obj.lower().endswith('.gif'):  # pragma: no cover
             _parse_gif(app, file_obj, data_label,
@@ -104,6 +106,7 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
                     spectrum_viewer_reference_name=spectrum_viewer_reference_name,
                     uncert_viewer_reference_name=uncert_viewer_reference_name
                 )
+        app.get_tray_item_from_name("Spectral Extraction").disabled_msg = ""
 
     # 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
@@ -121,11 +124,14 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
                 app, file_obj, data_label=data_label,
                 spectrum_viewer_reference_name=spectrum_viewer_reference_name
             )
+        app.get_tray_item_from_name("Spectral Extraction").disabled_msg = ""
+
     elif isinstance(file_obj, np.ndarray) and file_obj.ndim == 3:
         _parse_ndarray(app, file_obj, data_label=data_label, data_type=data_type,
                        flux_viewer_reference_name=flux_viewer_reference_name,
                        spectrum_viewer_reference_name=spectrum_viewer_reference_name,
                        uncert_viewer_reference_name=uncert_viewer_reference_name)
+        app.get_tray_item_from_name("Spectral Extraction").disabled_msg = ""
     else:
         raise NotImplementedError(f'Unsupported data format: {file_obj}')
 
@@ -252,12 +258,14 @@ def _parse_hdulist(app, hdulist, file_name=None,
 
         elif data_type == 'uncert':
             app.add_data_to_viewer(uncert_viewer_reference_name, data_label)
+            app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
 
         else:  # flux
             # Add flux to top left image viewer
             app.add_data_to_viewer(flux_viewer_reference_name, data_label)
             # Add flux to spectrum viewer
             app.add_data_to_viewer(spectrum_viewer_reference_name, data_label)
+            app._jdaviz_helper._loaded_flux_cube = app.data_collection[data_label]
 
 
 def _parse_jwst_s3d(app, hdulist, data_label, ext='SCI',
@@ -307,6 +315,11 @@ def _parse_jwst_s3d(app, hdulist, data_label, ext='SCI',
     if viewer_name == flux_viewer_reference_name:
         app.add_data_to_viewer(spectrum_viewer_reference_name, data_label)
 
+    if data_type == 'flux':
+        app._jdaviz_helper._loaded_flux_cube = app.data_collection[data_label]
+    elif data_type == 'uncert':
+        app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
+
 
 def _parse_esa_s3d(app, hdulist, data_label, ext='DATA', flux_viewer_reference_name=None,
                    spectrum_viewer_reference_name=None):
@@ -353,6 +366,11 @@ def _parse_esa_s3d(app, hdulist, data_label, ext='DATA', flux_viewer_reference_n
     app.add_data_to_viewer(flux_viewer_reference_name, data_label)
     app.add_data_to_viewer(spectrum_viewer_reference_name, data_label)
 
+    if data_type == 'flux':
+        app._jdaviz_helper._loaded_flux_cube = app.data_collection[data_label]
+    if data_type == 'uncert':
+        app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
+
 
 def _parse_spectrum1d_3d(app, file_obj, data_label=None,
                          flux_viewer_reference_name=None, spectrum_viewer_reference_name=None,
@@ -400,8 +418,10 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
         if attr == 'flux':
             app.add_data_to_viewer(flux_viewer_reference_name, cur_data_label)
             app.add_data_to_viewer(spectrum_viewer_reference_name, cur_data_label)
+            app._jdaviz_helper._loaded_flux_cube = app.data_collection[cur_data_label]
         elif attr == 'uncertainty':
             app.add_data_to_viewer(uncert_viewer_reference_name, cur_data_label)
+            app._jdaviz_helper._loaded_uncert_cube = app.data_collection[cur_data_label]
         # We no longer auto-populate the mask cube into a viewer
 
 
@@ -446,8 +466,10 @@ def _parse_ndarray(app, file_obj, data_label=None, data_type=None,
     if data_type == 'flux':
         app.add_data_to_viewer(flux_viewer_reference_name, data_label)
         app.add_data_to_viewer(spectrum_viewer_reference_name, data_label)
+        app._jdaviz_helper._loaded_flux_cube = app.data_collection[data_label]
     elif data_type == 'uncert':
         app.add_data_to_viewer(uncert_viewer_reference_name, data_label)
+        app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
 
 
 def _parse_gif(app, file_obj, data_label=None, flux_viewer_reference_name=None,

jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py

@@ -21,6 +21,7 @@
 from jdaviz.core.user_api import PluginUserApi
 from jdaviz.configs.cubeviz.plugins.parsers import _return_spectrum_with_correct_units
 
+
 __all__ = ['SpectralExtraction']
 
 ASTROPY_LT_5_3_2 = Version(astropy.__version__) < Version('5.3.2')
@@ -51,6 +52,7 @@ class SpectralExtraction(PluginTemplateMixin, DatasetSelectMixin,
     filename = Unicode().tag(sync=True)
     extracted_spec_available = Bool(False).tag(sync=True)
     overwrite_warn = Bool(False).tag(sync=True)
+
     # export_enabled controls whether saving to a file is enabled via the UI.  This
     # is a temporary measure to allow server-installations to disable saving server-side until
     # saving client-side is supported
@@ -83,6 +85,11 @@ def __init__(self, *args, **kwargs):
             # on the user's machine, so export support in cubeviz should be disabled
             self.export_enabled = False
 
+        self.disabled_msg = (
+            "Spectral Extraction requires a single dataset to be loaded into Cubeviz, "
+            "please load data to enable this plugin."
+        )
+
     @property
     def user_api(self):
         return PluginUserApi(
@@ -107,15 +114,8 @@ def collapse_to_spectrum(self, add_data=True, **kwargs):
             Additional keyword arguments passed to the NDDataArray collapse operation.
             Examples include ``propagate_uncertainties`` and ``operation_ignores_mask``.
         """
-        # get glue Data objects for the spectral cube and uncertainties
-        flux_viewer = self._app.get_viewer(
-            self._app._jdaviz_helper._default_flux_viewer_reference_name
-        )
-        uncert_viewer = self._app.get_viewer(
-            self._app._jdaviz_helper._default_uncert_viewer_reference_name
-        )
-        [spectral_cube] = flux_viewer.data()
-        [uncert_cube] = uncert_viewer.data()
+        spectral_cube = self._app._jdaviz_helper._loaded_flux_cube
+        uncert_cube = self._app._jdaviz_helper._loaded_uncert_cube
 
         # This plugin collapses over the *spatial axes* (optionally over a spatial subset,
         # defaults to ``No Subset``). Since the Cubeviz parser puts the fluxes

jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py

@@ -6,6 +6,7 @@
 from astropy.nddata import NDDataArray, StdDevUncertainty
 from specutils import Spectrum1D
 from regions import CirclePixelRegion, PixCoord
+from astropy.utils.exceptions import AstropyUserWarning
 
 ASTROPY_LT_5_3_2 = Version(astropy.__version__) < Version('5.3.2')
 
@@ -21,7 +22,7 @@ def test_version_before_nddata_update(cubeviz_helper, spectrum1d_cube_with_uncer
 def test_version_after_nddata_update(cubeviz_helper, spectrum1d_cube_with_uncerts):
     # Also test that plugin is disabled before data is loaded.
     plg = cubeviz_helper.plugins['Spectral Extraction']
-    assert plg._obj.disabled_msg == ''
+    assert plg._obj.disabled_msg != ''
 
     cubeviz_helper.load_data(spectrum1d_cube_with_uncerts)
 
@@ -36,6 +37,7 @@ def test_version_after_nddata_update(cubeviz_helper, spectrum1d_cube_with_uncert
     # Collapse the spectral cube using the methods in jdaviz:
     collapsed_cube_s1d = plg.collapse_to_spectrum(add_data=False)  # returns Spectrum1D
 
+    assert plg._obj.disabled_msg == ''
     assert isinstance(spectral_cube, NDDataArray)
     assert isinstance(collapsed_cube_s1d, Spectrum1D)
 
@@ -45,6 +47,58 @@ def test_version_after_nddata_update(cubeviz_helper, spectrum1d_cube_with_uncert
     )
 
 
+@pytest.mark.skipif(ASTROPY_LT_5_3_2, reason='Needs astropy 5.3.2 or later')
+def test_gauss_smooth_before_spec_extract(cubeviz_helper, spectrum1d_cube_with_uncerts):
+    # Also test if gaussian smooth plugin is run before spec extract
+    # that spec extract yields results of correct cube data
+    gs_plugin = cubeviz_helper.plugins['Gaussian Smooth']._obj
+
+    # give uniform unit uncertainties for spec extract test:
+    spectrum1d_cube_with_uncerts.uncertainty = StdDevUncertainty(
+        np.ones_like(spectrum1d_cube_with_uncerts.data)
+    )
+
+    cubeviz_helper.load_data(spectrum1d_cube_with_uncerts)
+
+    gs_plugin.dataset_selected = f'{cubeviz_helper.app.data_collection[0].label}'
+    gs_plugin.mode_selected = 'Spatial'
+    gs_plugin.stddev = 3
+
+    with pytest.warns(
+            AstropyUserWarning,
+            match='The following attributes were set on the data object, but will be ignored'):
+        gs_plugin.vue_apply()
+
+    gs_data_label = cubeviz_helper.app.data_collection[2].label
+    cubeviz_helper.app.add_data_to_viewer('flux-viewer', gs_data_label)
+
+    # create a subset with a single pixel:
+    regions = [
+        # create a subset with a single pixel:
+        CirclePixelRegion(PixCoord(0, 1), radius=0.7),
+        # two-pixel region:
+        CirclePixelRegion(PixCoord(0.5, 0), radius=1.2)
+    ]
+    cubeviz_helper.load_regions(regions)
+
+    extract_plugin = cubeviz_helper.plugins['Spectral Extraction']
+    extract_plugin.function = "Sum"
+    expected_uncert = 2
+
+    extract_plugin.spatial_subset = 'Subset 1'
+    collapsed_spec = extract_plugin.collapse_to_spectrum()
+
+    # this single pixel has two wavelengths, and all uncertainties are unity
+    # irrespective of which collapse function is applied:
+    assert len(collapsed_spec.flux) == 2
+    assert np.all(np.equal(collapsed_spec.uncertainty.array, 1))
+
+    # this two-pixel region has four unmasked data points per wavelength:
+    extract_plugin.spatial_subset = 'Subset 2'
+    collapsed_spec_2 = extract_plugin.collapse_to_spectrum()
+    assert np.all(np.equal(collapsed_spec_2.uncertainty.array, expected_uncert))
+
+
 @pytest.mark.skipif(ASTROPY_LT_5_3_2, reason='Needs astropy 5.3.2 or later')
 @pytest.mark.parametrize(
     "function, expected_uncert",