JP-3448: Always flag at least one group when showers are detected

CHANGES.rst

@@ -1,7 +1,10 @@
 1.5.3 (unreleased)
 ==================
 
-- 
+jump
+----
+- Fix the code to at least always flag the group with the shower and the requested
+  groups after the primary shower. [#237]
 
 1.5.2 (2023-12-13)
 ==================

src/stcal/jump/jump.py

@@ -1,11 +1,13 @@
 import logging
 import multiprocessing
 import time
-
-import cv2 as cv
 import numpy as np
-from astropy import stats
-from astropy.convolution import Ring2DKernel, convolve
+import cv2 as cv
+import astropy.stats as stats
+
+from astropy.convolution import Ring2DKernel
+from astropy.convolution import convolve
+
 
 from . import constants
 from . import twopoint_difference as twopt
@@ -471,7 +473,6 @@ def detect_jumps(
     data /= gain_2d
     err /= gain_2d
     readnoise_2d /= gain_2d
-
     # Return the updated data quality arrays
     return gdq, pdq, total_primary_crs, number_extended_events, stddev
 
@@ -631,7 +632,8 @@ def extend_ellipses(
     # For a given DQ plane it will use the list of ellipses to create
     #  expanded ellipses of pixels with
     # the jump flag set.
-    plane = gdq_cube[intg, grp, :, :]
+    out_gdq_cube = gdq_cube.copy()
+    plane = gdq_cube[intg, grp, :, :].copy()
     ncols = plane.shape[1]
     nrows = plane.shape[0]
     image = np.zeros(shape=(nrows, ncols, 3), dtype=np.uint8)
@@ -671,15 +673,38 @@ def extend_ellipses(
         )
         jump_ellipse = image[:, :, 2]
         ngrps = gdq_cube.shape[1]
-        last_grp = min(grp + num_grps_masked, ngrps)
+        last_grp = find_last_grp(grp, ngrps, num_grps_masked)
         #  This loop will flag the number of groups
         for flg_grp in range(grp, last_grp):
             sat_pix = np.bitwise_and(gdq_cube[intg, flg_grp, :, :], sat_flag)
             saty, satx = np.where(sat_pix == sat_flag)
             jump_ellipse[saty, satx] = 0
-            gdq_cube[intg, flg_grp, :, :] = np.bitwise_or(gdq_cube[intg, flg_grp, :, :], jump_ellipse)
-    return gdq_cube, num_ellipses
+            out_gdq_cube[intg, flg_grp, :, :] = np.bitwise_or(gdq_cube[intg, flg_grp, :, :], jump_ellipse)
+    diff_cube = out_gdq_cube - gdq_cube
+    return out_gdq_cube, num_ellipses
+
+def find_last_grp(grp, ngrps, num_grps_masked_after):
+    """
+    Parameters
+    ----------
+    grp : int
+        The location of the shower
+
+    ngrps : int
+        The number of groups in the integration
+
+    num_grps_masked_after : int
+        The requested number of groups to be flagged after the shower
 
+    Returns
+    -------
+    last_grp : int
+        The index of the last group to flag for the shower
+
+    """
+    num_grps_masked_after += 1
+    last_grp = min(grp + num_grps_masked_after, ngrps)
+    return last_grp
 
 def find_circles(dqplane, bitmask, min_area):
     # Using an input DQ plane this routine will find the groups of pixels with at least the minimum
@@ -834,7 +859,7 @@ def find_faint_extended(
     Total number of showers detected.
 
     """
-    read_noise_2 = readnoise_2d**2
+    read_noise_2d_sqr = readnoise_2d**2
     data = indata.copy()
     data[gdq == sat_flag] = np.nan
     data[gdq == 1] = np.nan
@@ -849,12 +874,11 @@ def find_faint_extended(
         # calculate sigma for each pixel
         if nints <= minimum_sigclip_groups:
             median_diffs = np.nanmedian(first_diffs_masked[intg], axis=0)
-            sigma = np.sqrt(np.abs(median_diffs) + read_noise_2 / nframes)
+            sigma = np.sqrt(np.abs(median_diffs) + read_noise_2d_sqr / nframes)
             # The difference from the median difference for each group
             e_jump = first_diffs_masked[intg] - median_diffs[np.newaxis, :, :]
             # SNR ratio of each diff.
             ratio = np.abs(e_jump) / sigma[np.newaxis, :, :]
-
         #  The convolution kernel creation
         ring_2D_kernel = Ring2DKernel(inner, outer)
         ngrps = data.shape[1]
@@ -881,14 +905,14 @@ def find_faint_extended(
             extended_emission = np.zeros(shape=(nrows, ncols), dtype=np.uint8)
             exty, extx = np.where(masked_smoothed_ratio > snr_threshold)
             extended_emission[exty, extx] = 1
+
             #  find the contours of the extended emission
             contours, hierarchy = cv.findContours(extended_emission, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
             #  get the contours that are above the minimum size
             bigcontours = [con for con in contours if cv.contourArea(con) > min_shower_area]
             #  get the minimum enclosing rectangle which is the same as the
             # minimum enclosing ellipse
             ellipses = [cv.minAreaRect(con) for con in bigcontours]
-
             expand_by_ratio = True
             expansion = 1.0
             plane = gdq[intg, grp, :, :]
@@ -938,10 +962,12 @@ def find_faint_extended(
         # This is deferred until all showers are detected. because the showers
         # can flag future groups and would confuse the detection algorithm if
         # we worked on groups that already had some flagged showers.
+        total_showers = 0
         for showers in all_ellipses:
             intg = showers[0]
             grp = showers[1]
             ellipses = showers[2]
+            total_showers += len(ellipses)
             gdq, num = extend_ellipses(
                 gdq,
                 intg,
@@ -952,9 +978,9 @@ def find_faint_extended(
                 expansion=ellipse_expand,
                 expand_by_ratio=True,
                 num_grps_masked=num_grps_masked,
-                max_extended_radius=max_extended_radius,
+                max_extended_radius=max_extended_radius
             )
-    return gdq, len(all_ellipses)
+    return gdq, total_showers
 
 
 def calc_num_slices(n_rows, max_cores, max_available):

tests/test_jump.py

@@ -8,6 +8,7 @@
     find_faint_extended,
     flag_large_events,
     point_inside_ellipse,
+    find_last_grp
 )
 
 DQFLAGS = {"JUMP_DET": 4, "SATURATED": 2, "DO_NOT_USE": 1, "GOOD": 0, "NO_GAIN_VALUE": 8}
@@ -352,3 +353,15 @@
     assert calc_num_slices(n_rows, "3/4", max_available_cores) == 1
     n_rows = 9
     assert calc_num_slices(n_rows, "21", max_available_cores) == 9
+
+
+def test_find_last_grp():
+    assert (find_last_grp(grp=5, ngrps=7, num_grps_masked=0) == 6)
+    assert (find_last_grp(grp=5, ngrps=7, num_grps_masked=2) == 7)
+    assert (find_last_grp(grp=5, ngrps=7, num_grps_masked=3) == 7)
+    assert (find_last_grp(grp=5, ngrps=6, num_grps_masked=1) == 6)
+    assert (find_last_grp(grp=5, ngrps=6, num_grps_masked=0) == 6)
+    assert (find_last_grp(grp=5, ngrps=6, num_grps_masked=2) == 6)
+    assert (find_last_grp(grp=5, ngrps=8, num_grps_masked=0) == 6)
+    assert (find_last_grp(grp=5, ngrps=8, num_grps_masked=1) == 7)
+    assert (find_last_grp(grp=5, ngrps=8, num_grps_masked=2) == 8)