Level 1 and Level 2 data product support in Rampviz

CHANGES.rst

@@ -21,7 +21,8 @@ New Features
 
 - The standalone version of jdaviz now uses solara instead of voila, resulting in faster load times. [#2909]
 
-- New configuration for ramp/Level 1 data products from Roman WFI and JWST [#3120, #3148, #3167, #3171]
+- New configuration for ramp/Level 1 and rate image/Level 2 data products from Roman WFI and
+  JWST [#3120, #3148, #3167, #3171, #3194]
 
 - Unit columns are now visible by default in the results table in model fitting. [#3196]
 

jdaviz/configs/default/plugins/data_quality/data_quality.py

@@ -106,8 +106,9 @@
         def is_image_viewer(viewer):
             from jdaviz.configs.imviz.plugins.viewers import ImvizImageView
             from jdaviz.configs.cubeviz.plugins.viewers import CubevizImageView
+            from jdaviz.configs.rampviz.plugins.viewers import RampvizImageView
 
-            return isinstance(viewer, (ImvizImageView, CubevizImageView))
+            return isinstance(viewer, (ImvizImageView, CubevizImageView, RampvizImageView))
 
         viewer_filter_names = [filt.__name__ for filt in self.viewer.filters]
         if 'is_image_viewer' not in viewer_filter_names:
@@ -151,6 +152,9 @@
 
     @property
     def dq_layer_selected_flattened(self):
+        if not hasattr(self, 'dq_layer'):
+            return []
+
         selected_dq = self.dq_layer.selected_obj
         if not len(selected_dq):
             return []
@@ -170,6 +174,7 @@
             return []
 
         dq = selected_dq[0].get_image_data()
+
         return np.unique(dq[~np.isnan(dq)])
 
     @property
@@ -210,11 +215,15 @@
 
             # for cubeviz, also change uncert-viewer defaults to
             # map the out-of-bounds regions to the cmap's `bad` color:
-            if self.app.config == 'cubeviz':
-                uncert_viewer = self.app.get_viewer(
-                    self.app._jdaviz_helper._default_uncert_viewer_reference_name
+            if self.app.config in ('cubeviz', 'rampviz'):
+                viewer = self.app.get_viewer(
+                    getattr(
+                        self.app._jdaviz_helper,
+                        '_default_uncert_viewer_reference_name', 'level-2'
+                    )
                 )
-                for layer in uncert_viewer.layers:
+
+                for layer in viewer.layers:
                     # allow bad alpha for image layers, not subsets:
                     if not hasattr(layer, 'subset_array'):
                         layer.composite._allow_bad_alpha = True

jdaviz/configs/default/plugins/plot_options/plot_options.py

@@ -1164,10 +1164,12 @@ def _viewer_is_image_viewer(self):
         from jdaviz.configs.imviz.plugins.viewers import ImvizImageView
         from jdaviz.configs.cubeviz.plugins.viewers import CubevizImageView
         from jdaviz.configs.mosviz.plugins.viewers import MosvizImageView, MosvizProfile2DView
+        from jdaviz.configs.rampviz.plugins.viewers import RampvizImageView
 
         def _is_image_viewer(viewer):
             return isinstance(viewer, (ImvizImageView, CubevizImageView,
-                                       MosvizImageView, MosvizProfile2DView))
+                                       MosvizImageView, MosvizProfile2DView,
+                                       RampvizImageView))
 
         viewers = self.viewer.selected_obj
         if not isinstance(viewers, list):

jdaviz/configs/imviz/plugins/coords_info/coords_info.py

@@ -262,6 +262,8 @@
             # cubeviz case:
             return arr[int(round(x)), int(round(y)), viewer.state.slices[-1]]
         elif image.ndim == 2:
+            if isinstance(viewer, RampvizImageView):
+                x, y = y, x
             return arr[int(round(y)), int(round(x))]
         else:  # pragma: no cover
             raise ValueError(f'does not support ndim={image.ndim}')
@@ -380,6 +382,11 @@
         elif isinstance(viewer, RampvizImageView):
             coords_status = False
 
+            slice_plugin = self.app._jdaviz_helper.plugins.get('Slice', None)
+            if slice_plugin is not None and len(image.shape) == 3:
+                # float to be compatible with default value of nan
+                self._dict['slice'] = float(viewer.slice)
+
         elif isinstance(viewer, MosvizImageView):
 
             if data_has_valid_wcs(image, ndim=2):

jdaviz/configs/imviz/plugins/parsers.py

@@ -232,7 +232,7 @@ def _parse_image(app, file_obj, data_label, ext=None, parent=None):
     for data, data_label in data_iter:
 
         # if the science extension hasn't been identified yet, do so here:
-        if sci_ext is None and ('[DATA' in data_label or '[SCI' in data_label):
+        if sci_ext is None and data.ndim == 2 and ('[DATA' in data_label or '[SCI' in data_label):
             sci_ext = data_label
 
         if isinstance(data.coords, GWCS) and (data.coords.bounding_box is not None):

jdaviz/configs/rampviz/helper.py

@@ -1,4 +1,4 @@
-from jdaviz.core.events import SliceSelectSliceMessage
+from jdaviz.core.events import SliceSelectSliceMessage, NewViewerMessage
 from jdaviz.core.helpers import CubeConfigHelper
 from jdaviz.configs.rampviz.plugins.viewers import RampvizImageView
 
@@ -49,8 +49,9 @@ def load_data(self, data, data_label=None, **kwargs):
     def _set_x_axis(self, msg={}):
         group_viewer = self.app.get_viewer(self._default_group_viewer_reference_name)
         ref_data = group_viewer.state.reference_data
-        group_viewer.state.x_att = ref_data.id["Pixel Axis 0 [z]"]
-        group_viewer.state.y_att = ref_data.id["Pixel Axis 1 [y]"]
+        if ref_data:
+            group_viewer.state.x_att = ref_data.id["Pixel Axis 0 [z]"]
+            group_viewer.state.y_att = ref_data.id["Pixel Axis 1 [y]"]
 
     def select_group(self, group_index):
         """
@@ -97,3 +98,27 @@ def get_data(self, data_label=None, spatial_subset=None,
         return self._get_data(data_label=data_label, spatial_subset=spatial_subset,
                               temporal_subset=temporal_subset,
                               cls=cls, use_display_units=use_display_units)
+
+    def create_image_viewer(self, viewer_name=None, data=None):
+        """
+        Create a new image viewer.
+
+        Parameters
+        ----------
+        viewer_name : str or `None`
+            Viewer name/ID to use. If `None`, it is auto-generated.
+
+        Returns
+        -------
+        viewer : `~jdaviz.configs.imviz.plugins.viewers.ImvizImageView`
+            Image viewer instance.
+
+        """
+        from jdaviz.configs.rampviz.plugins.viewers import RampvizImageView
+
+        # Cannot assign data to real Data because it loads but it will
+        # not update checkbox in Data menu.
+
+        return self.app._on_new_viewer(
+            NewViewerMessage(RampvizImageView, data=None, sender=self.app),
+            vid=viewer_name, name=viewer_name)

jdaviz/configs/rampviz/plugins/parsers.py

@@ -11,6 +11,7 @@
     standardize_metadata, download_uri_to_path,
     PRIHDR_KEY, standardize_roman_metadata
 )
+from jdaviz.configs.imviz.plugins.parsers import _parse_image as _parse_image_imviz
 
 try:
     from roman_datamodels import datamodels as rdd
@@ -24,7 +25,7 @@
 
 @data_parser_registry("ramp-data-parser")
 def parse_data(app, file_obj, data_type=None, data_label=None,
-               parent=None, cache=None, local_path=None, timeout=None,
+               ext=None, parent=None, cache=None, local_path=None, timeout=None,
                integration=0):
     """
     Attempts to parse a data file and auto-populate available viewers in
@@ -104,7 +105,9 @@
 
         with fits.open(file_obj) as hdulist:
             _parse_hdulist(
-                app, hdulist, file_name=data_label or file_name,
+                app, hdulist,
+                ext=ext, parent=parent,
+                file_name=data_label or file_name,
                 integration=integration,
                 group_viewer_reference_name=group_viewer_reference_name,
                 diff_viewer_reference_name=diff_viewer_reference_name,
@@ -158,7 +161,7 @@
     # swap axes per the conventions of ramp cubes
     # (group axis comes first) and the default in
     # rampviz (group axis expected last)
-    return np.swapaxes(x, 0, -1)
+    return np.transpose(x, (1, 2, 0))
 
 
 def _roman_3d_to_glue_data(
@@ -190,12 +193,15 @@
     ramp_cube_data_label = f"{data_label}[DATA]"
     ramp_diff_data_label = f"{data_label}[DIFF]"
 
+    data_reshaped = move_group_axis_last(data)
+    diff_data_reshaped = move_group_axis_last(diff_data)
+
     # load these cubes into the cache:
     app._jdaviz_helper.cube_cache[ramp_cube_data_label] = NDDataArray(
-        move_group_axis_last(data)
+        data_reshaped
     )
     app._jdaviz_helper.cube_cache[ramp_diff_data_label] = NDDataArray(
-        move_group_axis_last(diff_data)
+        diff_data_reshaped
     )
 
     if meta is not None:
@@ -204,14 +210,14 @@
     # load these cubes into the app:
     _parse_ndarray(
         app,
-        file_obj=move_group_axis_last(data),
+        file_obj=data_reshaped,
         data_label=ramp_cube_data_label,
         viewer_reference_name=group_viewer_reference_name,
         meta=meta
     )
     _parse_ndarray(
         app,
-        file_obj=move_group_axis_last(diff_data),
+        file_obj=diff_data_reshaped,
         data_label=ramp_diff_data_label,
         viewer_reference_name=diff_viewer_reference_name,
         meta=meta
@@ -224,7 +230,9 @@
 
 
 def _parse_hdulist(
-    app, hdulist, file_name=None,
+    app, hdulist,
+    ext=None, parent=None,
+    file_name=None,
     integration=None,
     group_viewer_reference_name=None,
     diff_viewer_reference_name=None,
@@ -235,8 +243,27 @@
         file_name = file_name or "Unknown HDU object"
 
     hdu = hdulist[1]  # extension containing the ramp
-    if hdu.header['NAXIS'] != 4:
-        raise ValueError(f"Expected a ramp with NAXIS=4 (with axes:"
+
+    if hdu.header['NAXIS'] == 2:
+        # this may be a calibrated Level 2 image:
+        _parse_image_imviz(app, hdulist, data_label=file_name, ext=ext, parent=parent)
+        new_viewer_name = 'level-2'
+
+        # create a level-2 viewer if none exists:
+        if new_viewer_name not in app.get_viewer_reference_names():
+            app._jdaviz_helper.create_image_viewer(viewer_name=new_viewer_name)
+
+        # add the SCI extension to the level-2 viewer:
+        if not ext:
+            idx = 1
+        elif ext and ('SCI' in ext or ext == '*'):
+            idx = len(ext) - ext.index('SCI')
+
+        app.add_data_to_viewer(new_viewer_name, app.data_collection[-idx].label)
+        return
+
+    elif hdu.header['NAXIS'] != 4:
+        raise ValueError(f"Expected a ramp with NAXIS=4 (with dimensions: "
                          f"integrations, groups, x, y), but got "
                          f"NAXIS={hdu.header['NAXIS']}.")
 
@@ -269,27 +296,6 @@
 def _parse_ramp_cube(app, ramp_cube_data, flux_unit, file_name,
                      group_viewer_reference_name, diff_viewer_reference_name,
                      meta=None):
-
-    # Identify NIRSpec IRS2 detector mode, which needs special treatment.
-    # jdox: https://jwst-docs.stsci.edu/jwst-near-infrared-spectrograph/nirspec-instrumentation/
-    # nirspec-detectors/nirspec-detector-readout-modes-and-patterns/nirspec-irs2-detector-readout-mode
-    if 'meta.model_type' in meta:
-        # this is a Level1bModel, which has metadata in a Node rather
-        # than a dictionary:
-        from_jwst_nirspec_irs2 = (
-            meta.get('meta._primary_header.TELESCOP') == 'JWST' and
-            meta.get('meta._primary_header.INSTRUME') == 'NIRSPEC' and
-            'IRS2' in meta.get('meta._primary_header.READPATT', '')
-        )
-    else:
-        # assume this was parsed from FITS:
-        header = meta.get('_primary_header', {})
-        from_jwst_nirspec_irs2 = (
-            header.get('TELESCOP') == 'JWST' and
-            header.get('INSTRUME') == 'NIRSPEC' and
-            'IRS2' in header.get('READPATT', '')
-        )
-
     # last axis is the group axis, first two are spatial axes:
     diff_data = np.vstack([
         # begin with a group of zeros, so
@@ -298,15 +304,8 @@
         np.diff(ramp_cube_data, axis=0)
     ])
 
-    if from_jwst_nirspec_irs2:
-        # JWST/NIRSpec in IRS2 readout needs an additional axis swap for x and y:
-        def move_axes(x):
-            return np.swapaxes(move_group_axis_last(x), 0, 1)
-    else:
-        move_axes = move_group_axis_last
-
-    ramp_cube = NDDataArray(move_axes(ramp_cube_data), unit=flux_unit, meta=meta)
-    diff_cube = NDDataArray(move_axes(diff_data), unit=flux_unit, meta=meta)
+    ramp_cube = NDDataArray(move_group_axis_last(ramp_cube_data), unit=flux_unit, meta=meta)
+    diff_cube = NDDataArray(move_group_axis_last(diff_data), unit=flux_unit, meta=meta)
 
     group_data_label = app.return_data_label(file_name, ext="DATA")
     diff_data_label = app.return_data_label(file_name, ext="DIFF")
@@ -322,6 +321,10 @@
         # load these cubes into the cache:
         app._jdaviz_helper.cube_cache[data_label] = data_entry
 
+        # set as reference data in this viewer
+        viewer = app.get_viewer(viewer_ref)
+        viewer.state.reference_data = app.data_collection[data_label]
+
     app._jdaviz_helper._loaded_flux_cube = app.data_collection[group_data_label]
 
 

jdaviz/configs/rampviz/plugins/ramp_extraction/ramp_extraction.py

@@ -138,7 +138,7 @@ def integration_viewer(self):
 
     def _on_data_added(self, msg={}):
         # only perform the default collapse after the first data load:
-        if len(self.app.data_collection) == 2:
+        if len(self.app._jdaviz_helper.cube_cache) and not self.extraction_available:
             self.extract(add_data=True)
             self.integration_viewer._initialize_x_axis()
 

jdaviz/configs/rampviz/plugins/viewers.py

@@ -129,8 +129,12 @@ def _default_integration_viewer_reference_name(self):
     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:
-            self.state.x_att = ref_data.id["Pixel Axis 0 [z]"]
+
+        if ref_data:
+            if "Pixel Axis 0 [z]" in [comp.label for comp in ref_data.components]:
+                self.state.x_att = ref_data.id["Pixel Axis 0 [z]"]
+            else:
+                self.state.x_att = ref_data.id["Pixel Axis 0 [y]"]
 
     def set_plot_axes(self):
         self.figure.axes[1].tick_format = None

jdaviz/configs/rampviz/tests/test_parser.py

@@ -12,34 +12,23 @@ def test_load_rectangular_ramp(rampviz_helper, jwst_level_1b_rectangular_ramp):
 
     parsed_cube_shape = rampviz_helper.app.data_collection[0].shape
     assert parsed_cube_shape == (
-        original_cube_shape[2], original_cube_shape[1], original_cube_shape[0]
+        original_cube_shape[1], original_cube_shape[2], original_cube_shape[0]
     )
 
 
-def test_load_nirspec_irs2(rampviz_helper, jwst_level_1b_rectangular_ramp):
-    # update the Level1bModel to have the header cards that are
-    # expected for an exposure from NIRSpec in IRS2 readout mode
-    jwst_level_1b_rectangular_ramp.update(
-        {
-            'meta': {
-                '_primary_header': {
-                    "TELESCOP": "JWST",
-                    "INSTRUME": "NIRSPEC",
-                    "READPATT": "NRSIRS2"
-                }
-            }
-        }
-    )
+def test_load_level_1_and_2(
+        rampviz_helper,
+        jwst_level_1b_rectangular_ramp,
+        jwst_level_2c_rate_image
+):
+    # load level 1 ramp and level 2 rate image
     rampviz_helper.load_data(jwst_level_1b_rectangular_ramp)
+    rampviz_helper.load_data(jwst_level_2c_rate_image)
 
-    # drop the integration axis
-    original_cube_shape = jwst_level_1b_rectangular_ramp.shape[1:]
-
-    # on ramp cube load (1), the parser loads a diff cube (2) and
-    # the ramp extraction plugin produces a default extraction (3):
-    assert len(rampviz_helper.app.data_collection) == 3
+    # confirm that a "level-2" viewer is created, and that
+    # the rate image is loaded into it
+    assert len(rampviz_helper.viewers) == 4
+    assert 'level-2' in rampviz_helper.viewers
 
-    parsed_cube_shape = rampviz_helper.app.data_collection[0].shape
-    assert parsed_cube_shape == (
-        original_cube_shape[1], original_cube_shape[2], original_cube_shape[0]
-    )
+    layers = rampviz_helper.app.get_viewer('level-2').layers
+    assert len(layers) == 1