removing previous code and adding modification to read_pattern code

src/stcal/saturation/saturation.py

@@ -10,7 +10,7 @@
 
 def flag_saturated_pixels(
     data, gdq, pdq, sat_thresh, sat_dq, atod_limit, dqflags,
-    n_pix_grow_sat=1, zframe=None, read_pattern=None, nframes=None):
+    n_pix_grow_sat=1, zframe=None, read_pattern=None):
     """
     Short Summary
     -------------
@@ -53,9 +53,6 @@ def flag_saturated_pixels(
     read_pattern : List[List[float or int]] or None
         The times or indices of the frames composing each group.
 
-    nframes : int
-        Number of frames per group.
-
     Returns
     -------
     gdq : int, 4D array
@@ -87,7 +84,20 @@ def flag_saturated_pixels(
             plane = data[ints, group, :, :]
 
             if read_pattern is not None:
-                nframes = None   # make sure the code does not go through this AND the nframes check
+                if group == 2:
+                    # Identify groups which we wouldn't expect to saturate by the third group,
+                    # on the basis of the first group
+                    mask = data[ints, 0, ...] / np.mean(read_pattern[0]) * read_pattern[2][-1] < sat_thresh
+
+                    # Identify groups with suspiciously large values in the second group
+                    mask &= data[ints, 1, ...] > sat_thresh / len(read_pattern[1])
+
+                    # Identify groups that are saturated in the third group
+                    mask &= np.where(gdq[ints, 2, :, :] & saturated, True, False)
+
+                    # Flag the 2nd group for the pixels passing that gauntlet
+                    gdq[ints, 1][mask] |= saturated
+
                 dilution_factor = np.mean(read_pattern[group]) / read_pattern[group][-1]
                 dilution_factor = np.where(no_sat_check_mask, 1, dilution_factor)
             else:
@@ -99,21 +109,6 @@ def flag_saturated_pixels(
             # and all following planes.
             np.bitwise_or(gdq[ints, group:, :, :], flagarray, gdq[ints, group:, :, :])
 
-            if nframes is not None:
-                if group == 2 and nframes > 1:
-                    # Look for pixels for which SATURATION is set in group 3. If the count
-                    # value in group 1 is less than 1/10 of the value that would be required
-                    # to trigger SATURATION (i.e., it isn't obviously a very bright source),
-                    # then also set the SATURATION flag for this pixel in group 2.
-                    saturation_in_goup3 = np.where(gdq[:, 2, ...] == saturated, True, False)
-                    if np.any(saturation_in_goup3):
-                        sat_nints, _, _ = np.shape(saturation_in_goup3)
-                        for ni in range(sat_nints):
-                            satgp3mask = saturation_in_goup3[ni]
-                            gp1mask = np.where(data[ni, 0][satgp3mask] < sat_thresh/10., True, False)
-                            if np.any(gp1mask):
-                                gdq[ni, 1, ...] = np.where(gp1mask, saturated, gdq[ni, 1, ...])
-
             # for A/D floor, the flag is only set of the current plane
             np.bitwise_or(gdq[ints, group, :, :], flaglowarray, gdq[ints, group, :, :])
 

tests/test_saturation.py

@@ -20,7 +20,6 @@ def test_basic_saturation_flagging():
     pdq = np.zeros((20, 20)).astype("uint32")
     sat_thresh = np.ones((20, 20)) * 100000.0
     sat_dq = np.zeros((20, 20)).astype("uint32")
-    nframes = 3
 
     # Add ramp values up to the saturation limit
     data[0, 0, 5, 5] = 0
@@ -40,33 +39,6 @@ def test_basic_saturation_flagging():
     assert np.all(gdq[0, satindex:, 5, 5] == DQFLAGS["SATURATED"])
 
 
-def test_nframes_saturation_flagging():
-    # Create inputs, data, and saturation maps
-    data = np.zeros((1, 5, 20, 20)).astype("float32")
-    gdq = np.zeros((1, 5, 20, 20)).astype("uint32")
-    pdq = np.zeros((20, 20)).astype("uint32")
-    sat_thresh = np.ones((20, 20)) * 100000.0
-    sat_dq = np.zeros((20, 20)).astype("uint32")
-    nframes = 3
-
-    # Add ramp values up to the saturation limit
-    data[0, 0, 5, 5] = 0
-    data[0, 1, 5, 5] = 20000
-    data[0, 2, 5, 5] = 60000
-    data[0, 3, 5, 5] = 60000  # Signal reaches saturation limit
-    data[0, 4, 5, 5] = 62000
-
-    # Set saturation value in the saturation model
-    satvalue = 60000
-    sat_thresh[5, 5] = satvalue
-
-    gdq, pdq, _ = flag_saturated_pixels(data, gdq, pdq, sat_thresh, sat_dq, ATOD_LIMIT, DQFLAGS, nframes=nframes)
-
-    # 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"])
-
-
 def test_read_pattern_saturation_flagging():
     """Check that the saturation threshold varies depending on how the reads
     are allocated into resultants."""
@@ -81,7 +53,7 @@ def test_read_pattern_saturation_flagging():
     # Add ramp values up to the saturation limit
     data[0, 0, 5, 5] = 0
     data[0, 1, 5, 5] = 20000
-    data[0, 2, 5, 5] = 40000
+    data[0, 2, 5, 5] = 60000
     data[0, 3, 5, 5] = 60000  # Signal reaches saturation limit
     data[0, 4, 5, 5] = 62000