JP-2556: Handle charge spilling out of saturated pixels into neighbors

CHANGES.rst

@@ -1,6 +1,11 @@
 0.6.4 (unreleased)
 ==================
 
+saturation
+---------- 
+
+- Added in functionality to deal with charge spilling from saturated pixels onto neighboring pixels [#83]
+
 0.6.3 (2022-04-27)
 ==================
 

setup.cfg

@@ -27,6 +27,7 @@ setup_requires =
 install_requires =
     numpy>=1.17
     astropy>=5.0.4
+    scipy>=1.6.0
 
 package_dir =
     =src

src/stcal/saturation/saturation.py

@@ -1,12 +1,15 @@
 import numpy as np
 import logging
 
+import copy
+from scipy import ndimage
+
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
 
 
 def flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq, atod_limit,
-                          dqflags):
+                          dqflags, n_pix_grow_sat=1):
     """
     Short Summary
     -------------
@@ -39,6 +42,10 @@ def flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq, atod_limit,
         A dictionary with at least the following keywords:
         DO_NOT_USE, SATURATED, AD_FLOOR, NO_SAT_CHECK
 
+    n_pix_grow_sat : int
+        Number of pixels that each flagged saturated pixel should be 'grown',
+        to account for charge spilling. Default is 1.
+
 
     Returns
     -------
@@ -85,6 +92,15 @@ def flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq, atod_limit,
             np.bitwise_or(gdq[ints, group, :, :], flaglowarray,
                           gdq[ints, group, :, :])
 
+            # now, flag any pixels that border saturated pixels (not A/D floor pix)
+            if n_pix_grow_sat > 0:
+                gdq_slice = copy.copy(gdq[ints, group, :, :]).astype(int)
+                only_sat = np.bitwise_and(gdq_slice, saturated).astype(np.uint8)
+                box_dim = (n_pix_grow_sat * 2) + 1
+                struct = np.ones((box_dim, box_dim)).astype(bool)
+                dialated = ndimage.binary_dilation(only_sat, structure=struct).astype(only_sat.dtype)
+                gdq[ints, group, :, :] = np.bitwise_or(gdq[ints, group, :, :], (dialated * saturated))
+
     n_sat = np.any(np.any(np.bitwise_and(gdq, saturated), axis=0), axis=0).sum()
     log.info(f'Detected {n_sat} saturated pixels')
     n_floor = np.any(np.any(np.bitwise_and(gdq, ad_floor), axis=0), axis=0).sum()

tests/test_saturation.py

@@ -8,6 +8,11 @@
 
 from stcal.saturation.saturation import flag_saturated_pixels
 
+# dictionary with required DQ flags
+DQFLAGS = {'DO_NOT_USE': 1, 'SATURATED': 2, 'AD_FLOOR': 64,
+           'NO_SAT_CHECK': 2097152}
+ATOD_LIMIT = 65535.  # Hard DN limit of 16-bit A-to-D converter
+
 
 def test_basic_saturation_flagging():
 
@@ -29,15 +34,46 @@ def test_basic_saturation_flagging():
     satvalue = 60000
     sat_thresh[5, 5] = satvalue
 
-    # dictionary with required DQ flags
-    dqflags = {'DO_NOT_USE': 1, 'SATURATED': 2, 'AD_FLOOR': 64,
-               'NO_SAT_CHECK': 2097152}
-
-    atod_limit = 65535.  # Hard DN limit of 16-bit A-to-D converter
-
     gdq, pdq = flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq,
-                                     atod_limit, dqflags)
+                                     ATOD_LIMIT, DQFLAGS)
+    print(gdq[:,:, 2:6, 2:6])
 
     # Make sure that groups with signal > saturation limit get flagged
     satindex = np.argmax(data[0, :, 5, 5] == satvalue)
-    assert np.all(gdq[0, satindex:, 5, 5] == dqflags['SATURATED'])
+    assert np.all(gdq[0, satindex:, 5, 5] == DQFLAGS['SATURATED'])
+
+
+def test_adjacent_pixel_flagging():
+    """ Test to see if specified number of adjacent pixels next to a saturated
+        pixel are also flagged, and that the edges of the dq array are treated
+        correctly when this is done. """
+
+    # Create inputs, data, and saturation maps
+    data = np.ones((1, 2, 5, 5)).astype('float32')
+    gdq = np.zeros((1, 2, 5, 5)).astype('uint32')
+    pdq = np.zeros((5, 5)).astype('uint32')
+    sat_thresh = np.ones((5, 5)) * 60000   # sat. thresh is 60000
+    sat_dq = np.zeros((5, 5)).astype('uint32')
+
+    # saturate a few pixels just in the first group
+    # (0, 0) and (1, 1) to test adjacent pixels
+    data[0, 0, 0, 0] = 62000
+    data[0, 0, 0, 1] = 62000
+    data[0, 0, 3, 3] = 62000
+
+    gdq, pdq = flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq,
+                                     ATOD_LIMIT, DQFLAGS)
+
+    sat_locs = np.where(np.bitwise_and(gdq, DQFLAGS['SATURATED']) ==
+                        DQFLAGS['SATURATED'])
+
+    assert sat_locs[0].all() == 0
+    assert np.all(sat_locs[1] == np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                                           0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                                           1, 1, 1, 1, 1, 1]))
+    assert np.all(sat_locs[2] == np.array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3,
+                                           4, 4, 4, 0, 0, 0, 1, 1, 1, 2, 2, 2,
+                                           3, 3, 3, 4, 4, 4]))
+    assert np.all(sat_locs[3] == np.array([0, 1, 2, 0, 1, 2, 2, 3, 4, 2, 3, 4,
+                                           2, 3, 4, 0, 1, 2, 0, 1, 2, 2, 3, 4,
+                                           2, 3, 4, 2, 3, 4]))