PoC: Cubeviz parser for Roman L1 (ramp) products

CHANGES.rst

@@ -680,6 +680,8 @@ Cubeviz
 - Add Spectral Extraction plugin for Cubeviz, which converts spectral cubes
   to 1D spectra with propagated uncertainties [#2039]
 
+- Support for loading Roman WFI Level 1 (ramp) files [#2376]
+
 Imviz
 ^^^^^
 

jdaviz/app.py

@@ -104,10 +104,13 @@
         # should return the converted values. Note that original_units
         # gives the units of the values array, which might not be the same
         # as the original native units of the component in the data.
-        if cid.label == "flux":
+        if cid.label == 'Pixel Axis 0 [z]' and target_units == '':
+            # handle ramps loaded into Cubeviz by avoiding conversion
+            # of the groups axis:
+            return values
+        elif cid.label == "flux":
             try:
                 spec = data.get_object(cls=Spectrum1D)
-
             except RuntimeError:
                 data = data.get_object(cls=NDDataArray)
                 spec = Spectrum1D(flux=data.data * u.Unit(original_units))
@@ -1255,6 +1258,9 @@
             elif axis == 'flux':
                 sv = self.get_viewer(self._jdaviz_helper._default_spectrum_viewer_reference_name)
                 return sv.data()[0].flux.unit
+            elif axis == 'data':
+                sv = self.get_viewer(self._jdaviz_helper._default_spectrum_viewer_reference_name)
+                return sv.data()[0].unit
             else:
                 raise ValueError(f"could not find units for axis='{axis}'")
         try:

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -15,12 +15,21 @@
 from jdaviz.utils import standardize_metadata, PRIHDR_KEY, download_uri_to_path
 
 
+try:
+    from roman_datamodels import datamodels as rdd
+except ImportError:
+    HAS_ROMAN_DATAMODELS = False
+else:
+    HAS_ROMAN_DATAMODELS = True
+
 __all__ = ['parse_data']
 
 EXT_TYPES = dict(flux=['flux', 'sci', 'data'],
                  uncert=['ivar', 'err', 'error', 'var', 'uncert'],
                  mask=['mask', 'dq', 'quality', 'data_quality'])
 
+cubeviz_ramp_meta_flag = '_roman_ramp'
+
 
 @data_parser_registry("cubeviz-data-parser")
 def parse_data(app, file_obj, data_type=None, data_label=None,
@@ -79,6 +88,19 @@
             _parse_gif(app, file_obj, data_label,
                        flux_viewer_reference_name=flux_viewer_reference_name)
             return
+        elif file_obj.lower().endswith('.asdf'):
+            if not HAS_ROMAN_DATAMODELS:
+                raise ImportError(
+                    "ASDF detected but roman-datamodels is not installed."
+                )
+            with rdd.open(file_obj) as pf:
+                _roman_3d_to_glue_data(
+                    app, pf, data_label,
+                    flux_viewer_reference_name=flux_viewer_reference_name,
+                    spectrum_viewer_reference_name=spectrum_viewer_reference_name,
+                    uncert_viewer_reference_name=uncert_viewer_reference_name
+                )
+            return
 
         # try parsing file_obj as a URI/URL:
         file_obj = download_uri_to_path(
@@ -158,7 +180,18 @@
         _parse_ndarray(app, file_obj, data_label=data_label, data_type=data_type,
                        flux_viewer_reference_name=flux_viewer_reference_name,
                        uncert_viewer_reference_name=uncert_viewer_reference_name)
+
         app.get_tray_item_from_name("Spectral Extraction").disabled_msg = ""
+
+    elif HAS_ROMAN_DATAMODELS and isinstance(file_obj, rdd.DataModel):
+        with rdd.open(file_obj) as pf:
+            _roman_3d_to_glue_data(
+                app, pf, data_label,
+                flux_viewer_reference_name=flux_viewer_reference_name,
+                spectrum_viewer_reference_name=spectrum_viewer_reference_name,
+                uncert_viewer_reference_name=uncert_viewer_reference_name
+            )
+
     else:
         raise NotImplementedError(f'Unsupported data format: {file_obj}')
 
@@ -472,7 +505,8 @@
 
 def _parse_ndarray(app, file_obj, data_label=None, data_type=None,
                    flux_viewer_reference_name=None,
-                   uncert_viewer_reference_name=None):
+                   uncert_viewer_reference_name=None,
+                   meta={}):
     if data_label is None:
         data_label = app.return_data_label(file_obj)
 
@@ -486,7 +520,9 @@
     if not hasattr(flux, 'unit'):
         flux = flux << u.count
 
-    meta = standardize_metadata({'_orig_spatial_wcs': None})
+    meta.update({'_orig_spatial_wcs': None})
+    meta = standardize_metadata(meta)
+
     s3d = Spectrum1D(flux=flux, meta=meta)
     app.add_data(s3d, data_label)
 
@@ -531,3 +567,88 @@
     else:
         data_type = ''
     return data_type
+
+
+def _roman_3d_to_glue_data(
+    app, file_obj, data_label,
+    flux_viewer_reference_name=None,
+    spectrum_viewer_reference_name=None,
+    uncert_viewer_reference_name=None,
+):
+    """
+    Parse a Roman 3D ramp cube file (Level 1),
+    usually with suffix '_uncal.asdf'.
+    """
+    def _swap_axes(x):
+        # swap axes per the conventions of Roman cubes
+        # (group axis comes first) and the default in
+        # Cubeviz (wavelength axis expected last)
+        return np.swapaxes(x, 0, -1)
+
+    # update viewer reference names for Roman ramp cubes:
+    # app._update_viewer_reference_name()
+
+    data = file_obj.data
+
+    if data_label is None:
+        data_label = app.return_data_label(file_obj)
+
+    # last axis is the group axis, first two are spatial axes:
+    diff_data = np.vstack([
+        # begin with a group of zeros, so
+        # that `diff_data.ndim == data.ndim`
+        np.zeros((1, *data[0].shape)),
+        np.diff(data, axis=0)
+    ])
+
+    meta = {cubeviz_ramp_meta_flag: True}
+
+    # load the `data` cube into what's usually the "flux-viewer"
+    _parse_ndarray(
+        app,
+        file_obj=_swap_axes(data),
+        data_label=f"{data_label}[DATA]",
+        data_type="flux",
+        flux_viewer_reference_name=flux_viewer_reference_name,
+        meta=meta
+    )
+
+    # load the diff of the data cube
+    # into what's usually the "uncert-viewer"
+    _parse_ndarray(
+        app,
+        file_obj=_swap_axes(diff_data),
+        data_type="uncert",
+        data_label=f"{data_label}[DIFF]",
+        uncert_viewer_reference_name=uncert_viewer_reference_name,
+        meta=meta
+    )
+
+    # If the default Cubeviz viewers are still their defaults, rename them to
+    # names that are appropriate for the Roman ramp files that we just parsed:
+    if 'flux-viewer' in app.get_viewer_reference_names():
+        app._update_viewer_reference_name('flux-viewer', 'group-viewer')
+        app._update_viewer_reference_name('uncert-viewer', 'group-diff-viewer')
+        app._update_viewer_reference_name('spectrum-viewer', 'integration-viewer')
+
+    # the default collapse function in the profile viewer is "sum",
+    # but for ramp files, "median" is more useful:
+    viewer = app.get_viewer('integration-viewer')
+    viewer.state.function = 'median'
+
+    # some Cubeviz plugins aren't relevant for ramps, so remove them:
+    remove_tray_items = [
+        'g-line-list',
+        'specviz-line-analysis',
+        'cubeviz-moment-maps',
+        'g-gaussian-smooth'
+    ]
+
+    for item_name in remove_tray_items:
+        item_names = [
+            tray_item['name'] for tray_item in app.state.tray_items
+        ]
+
+        app.state.tray_items.pop(
+            item_names.index(item_name)
+        )

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

@@ -7,6 +7,9 @@
 from astropy import units as u
 from astropy.units import UnitsWarning
 from traitlets import Bool, Int, Unicode, observe
+from specutils import Spectrum1D
+from glue_jupyter.bqplot.image import BqplotImageView
+from glue_jupyter.bqplot.profile import BqplotProfileView
 
 from jdaviz.configs.cubeviz.plugins.viewers import (WithSliceIndicator, WithSliceSelection,
                                                     CubevizImageView)
@@ -19,6 +22,7 @@
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import PluginTemplateMixin
 from jdaviz.core.user_api import PluginUserApi
+from jdaviz.configs.cubeviz.plugins.parsers import cubeviz_ramp_meta_flag
 
 
 __all__ = ['Slice']
@@ -147,6 +151,73 @@
     def slice_selection_viewers(self):
         return [v for v in self.app._viewer_store.values() if isinstance(v, WithSliceSelection)]
 
+    def slice_indicator(self):
+        return self.spectrum_viewer.slice_indicator
+
+    def _watch_viewer(self, viewer, watch=True):
+        if isinstance(viewer, BqplotImageView):
+            if watch and viewer not in self._watched_viewers:
+                self._watched_viewers.append(viewer)
+                viewer.state.add_callback('slices',
+                                          self._viewer_slices_changed)
+            elif not watch and viewer in self._watched_viewers:
+                viewer.state.remove_callback('slices',
+                                             self._viewer_slices_changed)
+                self._watched_viewers.remove(viewer)
+        elif isinstance(viewer, BqplotProfileView) and watch:
+            viewer_data = viewer.data()
+            if self._x_all is None and len(viewer.data()):
+                if hasattr(viewer_data, 'spectral_axis'):
+                    # cache wavelengths so that wavelength <> slice
+                    # conversion can be done efficiently
+                    self._update_data(viewer_data[0].spectral_axis)
+                else:
+                    sample_index = np.arange(viewer_data[0].shape[-1]) * u.one
+                    self._update_data(sample_index)
+
+            if viewer not in self._indicator_viewers:
+                self._indicator_viewers.append(viewer)
+                # if the units (or data) change, we need to update internally
+                viewer.state.add_callback("reference_data",
+                                          self._update_reference_data)
+
+    def _on_data_added(self, msg):
+        if isinstance(msg.viewer, BqplotImageView):
+            if len(msg.data.shape) == 3:
+                self.max_value = msg.data.shape[-1] - 1  # Same as i_end in Export Plot plugin
+                self._watch_viewer(msg.viewer, True)
+                msg.viewer.state.slices = (0, 0, int(self.slice))
+
+        elif isinstance(msg.viewer, BqplotProfileView):
+            self._watch_viewer(msg.viewer, True)
+
+    def _update_reference_data(self, reference_data):
+        if reference_data is None:
+            return  # pragma: no cover
+
+        if reference_data.meta.get(cubeviz_ramp_meta_flag, False):
+            sample_index = np.arange(reference_data['flux'].shape[-1]) * u.one
+            self._update_data(sample_index)
+        else:
+            self._update_data(reference_data.get_object(cls=Spectrum1D).spectral_axis)
+
+    def _update_data(self, x_all):
+        self._x_all = x_all.value
+
+        if self.wavelength == -1:
+            if len(x_all):
+                # initialize at middle of cube
+                self.slice = int(len(x_all)/2)
+            else:
+                # leave in the pre-init state and don't update the wavelength/slice
+                return
+
+        # Also update unit when data is updated
+        self.wavelength_unit = x_all.unit.to_string()
+
+        # force wavelength to update from the current slider value
+        self._on_slider_updated({'new': self.slice})
+
     @property
     def slice_indicator_viewers(self):
         return [v for v in self.app._viewer_store.values() if isinstance(v, WithSliceIndicator)]

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

@@ -7,7 +7,9 @@
 from specutils.utils import QuantityModel
 from traitlets import Bool, List, Unicode, observe
 
-from jdaviz.core.events import SnackbarMessage, GlobalDisplayUnitChanged
+from jdaviz.core.events import (
+    SnackbarMessage, GlobalDisplayUnitChanged, ViewerRenamedMessage
+)
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import (PluginTemplateMixin,
                                         SelectPluginComponent,
@@ -25,6 +27,7 @@
 from jdaviz.configs.default.plugins.model_fitting.initializers import (MODELS,
                                                                        initialize,
                                                                        get_model_parameters)
+from jdaviz.configs.cubeviz.plugins.parsers import cubeviz_ramp_meta_flag
 
 __all__ = ['ModelFitting']
 
@@ -158,6 +161,9 @@
         self.hub.subscribe(self, GlobalDisplayUnitChanged,
                            handler=self._on_global_display_unit_changed)
 
+        self.hub.subscribe(self, ViewerRenamedMessage,
+                           handler=self._on_viewer_renamed)
+
     @property
     def _default_spectrum_viewer_reference_name(self):
         return getattr(
@@ -328,7 +334,12 @@
             # during initial init, this can trigger before the component is initialized
             return
 
+        if self.dataset.selected_obj:
+            if self.dataset.selected_obj.meta.get(cubeviz_ramp_meta_flag, False):
+                return
+
         selected_spec = self.dataset.selected_spectrum
+
         if selected_spec is None:
             return
 
@@ -980,3 +991,8 @@
 
             spectrum.mask = subset_mask
         return spectrum
+
+    def _on_viewer_renamed(self, msg):
+        if self.dataset._viewers:
+            if msg.old_viewer_ref in self.dataset._viewers:
+                self.dataset._viewers = [self._default_spectrum_viewer_reference_name]

jdaviz/configs/specviz/plugins/line_analysis/line_analysis.py

@@ -167,7 +167,9 @@
             self.disabled_msg = 'Line Analysis unavailable without spectral data'
             return
 
-        if viewer_data.spectral_axis.unit == u.pix:
+        if not hasattr(viewer_data, 'spectral_axis'):
+            self.disabled_msg = 'Line Analysis plugin unavailable when viewing ramps'
+        elif viewer_data.spectral_axis.unit == u.pix:
             # disable the plugin until we can address this properly (either using the wavelength
             # solution to support pixels in line-lists, or properly displaying the extracted
             # 1d spectrum in wavelength-space)