JP-3581: badpix_selfcal step (#8500)

CHANGES.rst

@@ -49,6 +49,12 @@ associations
 - Exclude NIRISS SOSS data taken with uncalibrated filter F277W from spec2 and
   tso3 associations. [#8549]
 
+badpix_selfcal
+--------------
+
+- Added new optional step ``badpix_selfcal`` to the spec2 pipeline to self-calibrate
+  bad pixels in IFU data. [#8500]
+
 combine_1d
 ----------
 
@@ -204,6 +210,9 @@ pipeline
   virtual ("v") slits and the construction of output file names for each
   type. [#8442]
 
+- Added new optional step ``badpix_selfcal`` to the ``calwebb_spec2`` to self-calibrate
+  bad pixels in IFU data. [#8500]
+
 pixel_replace
 -------------
 

docs/jwst/badpix_selfcal/arguments.rst

@@ -0,0 +1,18 @@
+Step Arguments
+==============
+The ``badpix_selfcal`` step has the following optional arguments.
+
+``--flagfrac_lower`` (float, default=0.001)
+  The fraction of pixels to flag as outliers on the low-flux 
+  side of the smoothed-subtracted image.
+
+``--flagfrac_upper`` (float, default=0.001)
+  The fraction of pixels to flag as outliers on the high-flux
+  side of the smoothed-subtracted image.
+
+``--kernel_size`` (integer, default=15)
+  The size of the kernel to use for the median filter, which is applied 
+  in the spectral direction to make the smoothed image.
+
+``--save_flagged_bkg`` (boolean, default=False)
+  Whether to save the flagged background exposures to fits files.

docs/jwst/badpix_selfcal/description.rst

@@ -0,0 +1,60 @@
+Description
+===========
+
+:Class: `jwst.badpix_selfcal.BadpixSelfcalStep`
+:Alias: badpix_selfcal
+
+Overview
+--------
+The ``badpix_selfcal`` step flags bad pixels in the input data using a self-calibration 
+technique based on median filtering along the spectral axis. 
+When additional exposures are available, those are used in combination with the science
+exposure to identify bad pixels; when unavailable, the step will be skipped with a warning
+unless the ``force_single`` parameter is set True. In that case, the science data alone is
+used as its own "background".
+This correction is applied to `~jwst.datamodels.IFUImageModel` data
+directly after the :ref:`assign_wcs <assign_wcs_step>` correction has been applied
+in the :ref:`calwebb_spec2 <calwebb_spec2>` pipeline.
+
+Input details
+-------------
+The input data must be in the form of a `~jwst.datamodels.IFUImageModel` or
+a `~jwst.datamodels.ModelContainer` containing exactly one
+science exposure and any number of additional exposures.
+A fits or association file 
+that can be read into one of these data models is also acceptable.
+Any exposure with the metadata attribute ``asn.exptype`` set to 
+``background`` or ``selfcal`` will be used in conjunction with the science
+exposure to construct the combined background image. 
+
+Algorithm
+---------
+The algorithm relies on the assumption that bad pixels are outliers in the data along
+the spectral axis. The algorithm proceeds as follows:
+
+* A combined background image is created. If additional (``selfcal`` or ``background``)
+  exposures are available, 
+  the pixelwise minimum of all background, selfcal, and science exposures is taken. 
+  If no additional exposures are available, the science data itself is passed in 
+  without modification, serving as the "background image" for the rest of the procedure, 
+  i.e., true self-calibration.
+* The combined background image is median-filtered, ignoring NaNs, along the spectral axis 
+  with a user-specified kernel size. The default kernel size is 15 pixels.
+* The difference between the original background image and the median-filtered background image
+  is taken. The highest- and lowest-flux pixels in this difference image are
+  flagged as bad pixels. The default fraction of pixels to flag is 0.1% of the total number of pixels
+  on each of the high-flux and low-flux ends of the distribution. These fractions can be adjusted
+  using the ``flagfrac_lower`` and ``flagfrac_upper`` parameters for the low- and high-flux ends
+  of the distribution, respectively. The total fraction of flagged pixels is thus 
+  ``flagfrac_lower + flagfrac_upper``.
+* The bad pixels are flagged in the input data by setting the DQ flag to
+  "OTHER_BAD_PIXEL" and "DO_NOT_USE".
+* The bad pixels are also flagged in each exposure with ``asn.exptype`` equal to ``background``,
+  if available.
+
+Output product
+--------------
+The output is a new copy of the input `~jwst.datamodels.IFUImageModel`, with the
+bad pixels flagged.  If the entire ``calwebb_spec2`` pipeline is run, the background
+exposures passed into the ``background`` step will include the flags from the
+``badpix_selfcal`` step.

docs/jwst/badpix_selfcal/index.rst

@@ -0,0 +1,14 @@
+.. _badpix_selfcal_step:
+
+==========================
+Bad Pixel Self-Calibration
+==========================
+
+.. toctree::
+   :maxdepth: 2
+
+   description.rst
+   arguments.rst
+   reference_files.rst
+
+.. automodapi:: jwst.badpix_selfcal

docs/jwst/badpix_selfcal/reference_files.rst

@@ -0,0 +1,4 @@
+Reference Files
+===============
+The ``badpix_selfcal`` step does not use any reference files.
+

docs/jwst/package_index.rst

@@ -13,6 +13,7 @@ Package Index
    associations/index.rst
    background_step/index.rst
    background_subtraction/index.rst
+   badpix_selfcal/index.rst
    barshadow/index.rst
    charge_migration/index.rst
    combine_1d/index.rst

docs/jwst/pipeline/calwebb_spec2.rst

@@ -47,6 +47,8 @@ TSO exposures. The instrument mode abbreviations used in the table are as follow
 +==========================================================+=====+=====+=====+=====+=====+=====+=================+======+========+=====+
 | :ref:`assign_wcs <assign_wcs_step>`                      | |c| | |c| | |c| | |c| | |c| | |c| |       |c|       | |c|  |  |c|   | |c| |
 +----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+
+| :ref:`badpix_selfcal <badpix_selfcal_step>`\ :sup:`2`    |     |     | |c| |     |     | |c| |                 |      |        |     |
++----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+
 | :ref:`msaflagopen <msaflagopen_step>`                    |     | |c| | |c| |     |     |     |                 |      |        |     |
 +----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+
 | :ref:`nsclean <nsclean_step>`                            | |c| | |c| | |c| |     |     |     |                 |      |        |     |

jwst/badpix_selfcal/__init__.py

@@ -0,0 +1,4 @@
+from .badpix_selfcal_step import BadpixSelfcalStep
+
+
+__all__ = ['BadpixSelfcalStep']

jwst/badpix_selfcal/badpix_selfcal.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+import numpy as np
+import jwst.datamodels as dm
+from jwst.outlier_detection.outlier_detection_ifu import medfilt
+from stdatamodels.jwst.datamodels.dqflags import pixel
+import warnings
+
+
+def badpix_selfcal(minimg: np.ndarray,
+                   flagfrac_lower: float = 0.001,
+                   flagfrac_upper: float = 0.001,
+                   kernel_size: int = 15,
+                   dispaxis=None) -> np.ndarray:
+    """
+    Flag residual artifacts as bad pixels in the DQ array of a JWST exposure
+
+    Parameters
+    ----------
+    minimg : np.ndarray
+        Selfcal data of shape (x, y), i.e., after some operation has
+        already been taken to combine multiple exposures,
+        typically a MIN operation.
+    flagfrac_lower : float
+        Fraction of pixels to flag on the low end
+    flagfrac_upper : float
+        Fraction of pixels to flag on the high end
+    kernel_size : int
+        Size of kernel for median filter
+    dispaxis : int
+        Dispersion axis, either 1 or 2. If None, a two-dimensional
+        median filter is applied.
+
+    Returns
+    -------
+    flagged_indices : np.ndarray
+        Indices of the flagged pixels,
+        shaped like output from np.where
+    """
+    if dispaxis not in [1, 2, None]:
+        raise ValueError("dispaxis must be either 1 or 2, or None.")
+
+    # Determine outliers using median filter
+    if dispaxis is None:
+        kern_size = (kernel_size, kernel_size)
+    elif dispaxis == 2:
+        # check that this is the right way around!
+        kern_size = (kernel_size, 1)
+    elif dispaxis == 1:
+        kern_size = (1, kernel_size)
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=RuntimeWarning, message="All-NaN")
+        smoothed = medfilt(minimg, kern_size)
+    minimg_hpf = minimg - smoothed
+
+    # Flag outliers using percentile cutoff
+    flag_low, flag_high = np.nanpercentile(minimg_hpf, [flagfrac_lower * 100, (1 - flagfrac_upper) * 100])
+    bad = (minimg_hpf > flag_high) | (minimg_hpf < flag_low)
+    flagged_indices = np.where(bad)
+
+    return flagged_indices
+
+
+def apply_flags(input_model: dm.IFUImageModel, flagged_indices: np.ndarray) -> dm.IFUImageModel:
+    """
+    Apply the flagged indices to the input model. Sets the flagged pixels to NaN
+    and the DQ flag to DO_NOT_USE + OTHER_BAD_PIXEL
+
+    Parameters
+    ----------
+    input_model : dm.IFUImageModel
+        Input science data to be corrected
+    flagged_indices : np.ndarray
+        Indices of the flagged pixels,
+        shaped like output from np.where
+
+    Returns
+    -------
+    output_model : dm.IFUImageModel
+        Flagged data model
+    """
+
+    input_model.dq[flagged_indices] |= pixel["DO_NOT_USE"] + pixel["OTHER_BAD_PIXEL"]
+
+    input_model.data[flagged_indices] = np.nan
+    input_model.err[flagged_indices] = np.nan
+    input_model.var_poisson[flagged_indices] = np.nan
+    input_model.var_rnoise[flagged_indices] = np.nan
+    input_model.var_flat[flagged_indices] = np.nan
+
+    return input_model