Track loaded mask cube in cubeviz

spacetelescope · Nov 25, 2024 · 778c1cb · 778c1cb
1 parent 88d82f9
commit 778c1cb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/jdaviz/configs/cubeviz/helper.py b/jdaviz/configs/cubeviz/helper.py
@@ -22,6 +22,7 @@ class Cubeviz(CubeConfigHelper, LineListMixin):
 
     _loaded_flux_cube = None
     _loaded_uncert_cube = None
+    _loaded_mask_cube = None
     _cube_viewer_cls = CubevizImageView
 
     def __init__(self, *args, **kwargs):

diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -304,8 +304,9 @@ def _parse_hdulist(app, hdulist, file_name=None,
         app.add_data(sc, data_label)
 
         if data_type == 'mask':
-            # We no longer auto-populate the mask cube into a viewer
-            pass
+            # We no longer auto-populate the mask cube into a viewer, but we still want
+            # to keep track of this cube for use in, e.g., spectral extraction.
+            app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
         elif data_type == 'uncert':
             app.add_data_to_viewer(uncert_viewer_reference_name, data_label)
@@ -381,6 +382,8 @@ def _parse_jwst_s3d(app, hdulist, data_label, ext='SCI',
         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]
+    elif data_type == 'mask':
+        app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
 
 def _parse_esa_s3d(app, hdulist, data_label, ext='DATA', flux_viewer_reference_name=None):
@@ -429,8 +432,10 @@ def _parse_esa_s3d(app, hdulist, data_label, ext='DATA', flux_viewer_reference_n
 
     if data_type == 'flux':
         app._jdaviz_helper._loaded_flux_cube = app.data_collection[data_label]
-    if data_type == 'uncert':
+    elif data_type == 'uncert':
         app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
+    elif data_type == 'mask':
+        app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
 
 def _parse_spectrum1d_3d(app, file_obj, data_label=None,
@@ -480,7 +485,8 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
         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
+        elif attry == 'mask':
+            app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
 
 def _parse_spectrum1d(app, file_obj, data_label=None, spectrum_viewer_reference_name=None):
@@ -539,6 +545,8 @@ def _parse_ndarray(app, file_obj, data_label=None, data_type=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]
+    elif data_type == 'mask':
+        app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
 
 def _parse_gif(app, file_obj, data_label=None, flux_viewer_reference_name=None):  # pragma: no cover