JP-2256: Apply median filter to IVM weight

CHANGES.rst

@@ -1,7 +1,12 @@
 1.14.1 (unreleased)
 ===================
 
-- 
+resample
+--------
+
+- Apply a median filter to IVM weight array to better handle saturated
+  pixels. [#7563]
+
 
 1.14.0 (2024-03-29)
 ===================
@@ -17,7 +22,7 @@ ami
 - Additional optional input arguments for greater user processing flexibility.
   See documentation for details. [#7862]
 
-- Bad pixel correction applied to data using new NRM reference file to calculate 
+- Bad pixel correction applied to data using new NRM reference file to calculate
   complex visibility support (M. Ireland method implemented by J. Kammerer). [#7862]
 
 - Make ``AmiAnalyze`` and ``AmiNormalize`` output conform to the OIFITS standard. [#7862]
@@ -53,7 +58,7 @@ charge_migration
   as DO_NOT_USE.  This group, and all subsequent groups, are then flagged as
   CHARGELOSS and DO_NOT_USE.  The four nearest pixel neighbor are then flagged
   in the same group. [#8336]
-  
+
 - Added warning handler for expected NaN and inf clipping in the
   ``sigma_clip`` function. [#8320]
 
@@ -340,7 +345,7 @@ tweakreg
 
 - Fixed a bug that caused failures instead of warnings when no GAIA sources
   were found within the bounding box of the input image. [#8334]
-  
+
 - Suppress AstropyUserWarnings regarding NaNs in the input data. [#8320]
 
 wfs_combine
@@ -903,6 +908,7 @@ resample
 - Update the following exposure time keywords: XPOSURE (EFFEXPTM),
   DURATION and TELAPSE. [#7793]
 
+
 residual_fringe
 ---------------
 
@@ -1202,6 +1208,10 @@ tweakreg
 - Added the 'GAIADR3' catalog to the available options for alignment;
   this has been enabled as the default option [#7611].
 
+- Apply median filtering to IVM weight array to better handle saturated
+  pixels. Added option of 'ivm-med', 'ivm-smed' (recommended default) to the
+  ``weight_type`` argument of ``ResampleStep``. [#7563]
+
 
 1.10.2 (2023-04-14)
 ===================

jwst/resample/resample_step.py

@@ -46,7 +46,7 @@ class ResampleStep(Step):
         pixfrac = float(default=1.0) # change back to None when drizpar reference files are updated
         kernel = string(default='square') # change back to None when drizpar reference files are updated
         fillval = string(default='INDEF' ) # change back to None when drizpar reference files are updated
-        weight_type = option('ivm', 'exptime', None, default='ivm')  # change back to None when drizpar ref update
+        weight_type = option('ivm', 'ivm-med', 'ivm-smed', 'exptime', None, default='ivm-smed')  # change back to None when drizpar ref update
         output_shape = int_list(min=2, max=2, default=None)  # [x, y] order
         crpix = float_list(min=2, max=2, default=None)
         crval = float_list(min=2, max=2, default=None)
@@ -58,7 +58,7 @@ class ResampleStep(Step):
         blendheaders = boolean(default=True)
         allowed_memory = float(default=None)  # Fraction of memory to use for the combined image.
         in_memory = boolean(default=True)
-    """
+    """  # noqa: E501
 
     reference_file_types = ['drizpars']
 

jwst/resample/resample_utils.py

@@ -2,9 +2,10 @@
 import logging
 import warnings
 
+import gwcs
 import numpy as np
+from scipy.ndimage import median_filter
 from astropy import units as u
-import gwcs
 
 from stdatamodels.dqflags import interpret_bit_flags
 from stdatamodels.jwst.datamodels.dqflags import pixel
@@ -170,23 +171,86 @@
     """
     dqmask = build_mask(model.dq, good_bits)
 
-    if weight_type == 'ivm':
+    if weight_type and weight_type.startswith('ivm'):
+        weight_type = weight_type.strip()
+        selective_median = weight_type.startswith('ivm-smed')
+
+        bitvalue = interpret_bit_flags(good_bits, mnemonic_map=pixel)
+        if bitvalue is None:
+            bitvalue = 0
+        saturation = pixel['SATURATED']
+
+        if selective_median and not (bitvalue & saturation):
+            selective_median = False
+            weight_type = 'ivm'
+
         if (model.hasattr("var_rnoise") and model.var_rnoise is not None and
                 model.var_rnoise.shape == model.data.shape):
             with np.errstate(divide="ignore", invalid="ignore"):
                 inv_variance = model.var_rnoise**-1
+
             inv_variance[~np.isfinite(inv_variance)] = 1
+
+            if weight_type != 'ivm':
+                ny, nx = inv_variance.shape
+
+                # apply a median filter to smooth the weight at saturated
+                # (or high read-out noise) single pixels. keep kernel size
+                # small to still give lower weight to extended CRs, etc.
+                ksz = weight_type[8 if selective_median else 7 :]
+                if ksz:
+                    kernel_size = int(ksz)
+                    if not (kernel_size % 2):
+                        raise ValueError(
+                            'Kernel size of the median filter in IVM weighting'
+                            ' must be an odd integer.'
+                        )
+                else:
+                    kernel_size = 3
+
+                ivm_copy = inv_variance.copy()
+
+                if selective_median:
+                    # apply median filter selectively only at
+                    # points of partially saturated sources:
+                    jumps = np.where(model.dq & saturation)
+                    w2 = kernel_size // 2
+                    for r, c in zip(*jumps):
+                        x1 = max(0, c - w2)
+                        x2 = min(nx, c + w2 + 1)
+                        y1 = max(0, r - w2)
+                        y2 = min(ny, r + w2 + 1)
+                        data = ivm_copy[y1:y2, x1:x2][dqmask[y1:y2, x1:x2]]
+                        if data.size:
+                            inv_variance[r, c] = np.median(data)
+                        # else: leave it as is
+
+                else:
+                    # apply median to the entire inv-var array:
+                    inv_variance = median_filter(
+                        inv_variance,
+                        size=kernel_size
+                    )
+                bad_dqmask = np.logical_not(dqmask)
+                inv_variance[bad_dqmask] = ivm_copy[bad_dqmask]
+
         else:
-            warnings.warn("var_rnoise array not available. Setting drizzle weight map to 1",
-                          RuntimeWarning)
+            warnings.warn(
+                "var_rnoise array not available. "
+                "Setting drizzle weight map to 1",
+                RuntimeWarning
+            )
             inv_variance = 1.0
+
         result = inv_variance * dqmask
+
     elif weight_type == 'exptime':
         if _check_for_tmeasure(model):
             exptime = model.meta.exposure.measurement_time
         else:
             exptime = model.meta.exposure.exposure_time
         result = exptime * dqmask
+
     else:
         result = np.ones(model.data.shape, dtype=model.data.dtype) * dqmask
 

jwst/resample/tests/test_utils.py

@@ -20,6 +20,7 @@
 
 
 DO_NOT_USE = dqflags.pixel["DO_NOT_USE"]
+SATURATED = dqflags.pixel["SATURATED"]
 GOOD = dqflags.pixel["GOOD"]
 
 
@@ -127,6 +128,46 @@ def test_build_driz_weight(weight_type):
     assert weight_map.dtype == np.float32
 
 
+def test_build_driz_weight_median():
+    """Check that correct weight map is returned of different weight types"""
+    model = ImageModel((10, 10))
+    model.dq[0] = DO_NOT_USE
+    model.meta.exposure.exposure_time = 10.0
+    model.var_rnoise += 0.1
+    model.var_rnoise[5, 5] = 1000
+
+    weight_map = build_driz_weight(
+        model,
+        weight_type="ivm-med3",
+        good_bits="SATURATED"
+    )
+    assert_array_equal(weight_map[0], 0)
+    assert_array_equal(weight_map[1:], 10.0)
+    assert weight_map.dtype == np.float32
+
+
+def test_build_driz_weight_selective_median():
+    """Check that correct weight map is returned of different weight types"""
+    model = ImageModel((10, 10))
+    model.dq[0] = DO_NOT_USE
+    model.meta.exposure.exposure_time = 10.0
+    model.var_rnoise += 0.1
+    model.var_rnoise[0, 5] = 1000
+    model.var_rnoise[5, 5] = 1000
+    model.dq[5, 5] = SATURATED
+    model.dq[0, 5] |= SATURATED
+
+    weight_map = build_driz_weight(
+        model,
+        weight_type="ivm-smed",
+        good_bits="SATURATED"
+    )
+
+    assert_array_equal(weight_map[0], 0)
+    assert_array_equal(weight_map[1:], 10.0)
+    assert weight_map.dtype == np.float32
+
+
 @pytest.mark.parametrize("weight_type", ["ivm", None])
 def test_build_driz_weight_zeros(weight_type):
     """Check that zero or not finite weight maps get set to 1"""