JP-3502: Extend snowball flags to next integration (#238)

* Update jump.py * add fits * Update jump.py * Update grps flagging * fixes * Update jump.py * Update jump.py * Update jump.py * Update jump.py * cleanup * remove fits * Update test_jump.py * Update jump.py * Update CHANGES.rst * Update CHANGES.rst took suggestion Co-authored-by: Howard Bushouse <[email protected]> * Update src/stcal/jump/jump.py fix spacing Co-authored-by: Howard Bushouse <[email protected]> * Update src/stcal/jump/jump.py fuxed Co-authored-by: Howard Bushouse <[email protected]> * Update jump.py resolving comments * Update jump.py --------- Co-authored-by: Howard Bushouse <[email protected]>
spacetelescope · Feb 15, 2024 · 44f50f2 · 44f50f2
1 parent 5cdca94
commit 44f50f2
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 19 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -4,6 +4,13 @@
 Changes to API
 --------------
 
+jump
+~~~~
+
+- Add in the flagging of groups in the integration after a snowball
+  occurs. The saturated core of the snowball gets flagged as jump
+  for a number of groups passed in as a parameter [#238]
+
 Bug Fixes
 ---------
 

diff --git a/src/stcal/jump/jump.py b/src/stcal/jump/jump.py
@@ -56,6 +56,8 @@ def detect_jumps(
     minimum_groups=3,
     minimum_sigclip_groups=100,
     only_use_ints=True,
+    mask_persist_grps_next_int=True,
+    persist_grps_flagged=25,
 ):
     """
     This is the high-level controlling routine for the jump detection process.
@@ -299,6 +301,8 @@ def detect_jumps(
                 edge_size=edge_size,
                 sat_expand=sat_expand,
                 max_extended_radius=max_extended_radius,
+                mask_persist_grps_next_int=mask_persist_grps_next_int,
+                persist_grps_flagged=persist_grps_flagged,
             )
             log.info("Total snowballs = %i", total_snowballs)
             number_extended_events = total_snowballs
@@ -457,6 +461,8 @@ def detect_jumps(
                 edge_size=edge_size,
                 sat_expand=sat_expand,
                 max_extended_radius=max_extended_radius,
+                mask_persist_grps_next_int=mask_persist_grps_next_int,
+                persist_grps_flagged=persist_grps_flagged,
             )
             log.info("Total snowballs = %i", total_snowballs)
             number_extended_events = total_snowballs
@@ -503,6 +509,8 @@ def flag_large_events(
     sat_expand=2,
     edge_size=25,
     max_extended_radius=200,
+    mask_persist_grps_next_int=True,
+    persist_grps_flagged=5,
 ):
     """
     This routine controls the creation of expanded regions that are flagged as
@@ -540,7 +548,12 @@ def flag_large_events(
         required for a snowball to be created
     max_extended_radius : int
         The largest radius that a snowball or shower can be extended
-
+    mask_persist_grps_next_int : bool
+        The flag to turn on the extension of the flagging of the saturated cores of
+        snowballs.
+    persist_grps_flagged : int
+        How many groups to be flagged when the saturated cores are extended into
+        subsequent integrations
     Returns
     -------
     Nothing, gdq array is modified.
@@ -550,8 +563,8 @@ def flag_large_events(
 
     n_showers_grp = []
     total_snowballs = 0
-    nints = gdq.shape[0]
-    ngrps = gdq.shape[1]
+    nints, ngrps, nrows, ncols = gdq.shape
+    persist_jumps = np.zeros(shape=(nints, gdq.shape[2], gdq.shape[3]), dtype=np.uint8)
     for integration in range(nints):
         for group in range(1, ngrps):
             current_gdq = gdq[integration, group, :, :]
@@ -565,10 +578,8 @@ def flag_large_events(
             jump_ellipses = find_ellipses(gdq[integration, group, :, :], jump_flag, min_jump_area)
             if sat_required_snowball:
                 low_threshold = edge_size
-                nrows = gdq.shape[2]
                 high_threshold = max(0, nrows - edge_size)
-
-                gdq, snowballs = make_snowballs(
+                gdq, snowballs, persist_jumps = make_snowballs(
                     gdq,
                     integration,
                     group,
@@ -579,7 +590,9 @@ def flag_large_events(
                     min_sat_radius_extend,
                     sat_expand,
                     sat_flag,
+                    jump_flag,
                     max_extended_radius,
+                    persist_jumps,
                 )
             else:
                 snowballs = jump_ellipses
@@ -596,15 +609,26 @@ def flag_large_events(
                 num_grps_masked=0,
                 max_extended_radius=max_extended_radius,
             )
-    return gdq, total_snowballs
 
+    #  Test to see if the flagging of the saturated cores will be extended into the
+    #  subsequent integrations. Persist_jumps contains all the pixels that were saturated
+    #  in the cores of snowballs.
+    if mask_persist_grps_next_int:
+        for intg in range(1, nints):
+            if persist_grps_flagged >= 1:
+                last_grp_flagged = min(persist_grps_flagged, ngrps)
+                gdq[intg, 1:last_grp_flagged, :, :] = np.bitwise_or(gdq[intg, 1:last_grp_flagged, :, :],
+                                                                    np.repeat(persist_jumps[intg - 1, np.newaxis, :, :],
+                                                                    last_grp_flagged - 1, axis=0))
+    return gdq, total_snowballs
 
 def extend_saturation(
-    cube, grp, sat_ellipses, sat_flag, min_sat_radius_extend, expansion=2, max_extended_radius=200
+    cube, grp, sat_ellipses, sat_flag, jump_flag, min_sat_radius_extend, persist_jumps,
+    expansion=2, max_extended_radius=200
 ):
-    ncols = cube.shape[2]
-    nrows = cube.shape[1]
+    ngroups, nrows, ncols = cube.shape
     image = np.zeros(shape=(nrows, ncols, 3), dtype=np.uint8)
+    persist_image = np.zeros(shape=(nrows, ncols, 3), dtype=np.uint8)
     outcube = cube.copy()
     for ellipse in sat_ellipses:
         ceny = ellipse[0][0]
@@ -623,13 +647,29 @@ def extend_saturation(
                 alpha,
                 0,
                 360,
-                (0, 0, 22),
+                (0, 0, 22),  # in the RGB cube, set blue plane pixels of the ellipse to 22
                 -1,
             )
-            sat_ellipse = image[:, :, 2]
-            saty, satx = np.where(sat_ellipse == 22)
+            #  Create another non-extended ellipse that is used to create the
+            #  persist_jumps for this integration. This will be used to mask groups
+            #  in subsequent integrations.
+            sat_ellipse = image[:, :, 2]  # extract the Blue plane of the image
+            saty, satx = np.where(sat_ellipse == 22)  # find all the ellipse pixels in the ellipse
             outcube[grp:, saty, satx] = sat_flag
-    return outcube
+            persist_image = cv.ellipse(
+                persist_image,
+                (round(ceny), round(cenx)),
+                (round(ellipse[1][0] / 2), round(ellipse[1][1] / 2)),
+                alpha,
+                0,
+                360,
+                (0, 0, 22),
+                -1,
+            )
+            persist_ellipse = persist_image[:, :, 2]
+            persist_saty, persist_satx = np.where(persist_ellipse == 22)
+            persist_jumps[persist_saty, persist_satx] = jump_flag
+    return outcube, persist_jumps
 
 
 def extend_ellipses(
@@ -754,7 +794,9 @@ def make_snowballs(
     min_sat_radius,
     expansion,
     sat_flag,
+    jump_flag,
     max_extended_radius,
+    persist_jumps,
 ):
     # This routine will create a list of snowballs (ellipses) that have the
     # center
@@ -786,17 +828,19 @@ def make_snowballs(
                 ):
                     snowballs.append(jump)
     # extend the saturated ellipses that are larger than the min_sat_radius
-    gdq[integration, :, :, :] = extend_saturation(
+    gdq[integration, :, :, :], persist_jumps[integration, :, :] = extend_saturation(
         gdq[integration, :, :, :],
         group,
         sat_ellipses,
         sat_flag,
+        jump_flag,
         min_sat_radius,
+        persist_jumps[integration, :, :],
         expansion=expansion,
         max_extended_radius=max_extended_radius,
     )
 
-    return gdq, snowballs
+    return gdq, snowballs, persist_jumps
 
 
 def point_inside_ellipse(point, ellipse):

diff --git a/tests/test_jump.py b/tests/test_jump.py
@@ -156,6 +156,7 @@ def test_find_ellipse2():
 
 def test_extend_saturation_simple():
     cube = np.zeros(shape=(5, 7, 7), dtype=np.uint8)
+    persist_jumps = np.zeros(shape=(7, 7), dtype=np.uint8)
     grp = 1
     min_sat_radius_extend = 1
     cube[1, 3, 3] = DQFLAGS["SATURATED"]
@@ -165,8 +166,9 @@ def test_extend_saturation_simple():
     cube[1, 3, 2] = DQFLAGS["SATURATED"]
     cube[1, 2, 2] = DQFLAGS["JUMP_DET"]
     sat_circles = find_ellipses(cube[grp, :, :], DQFLAGS["SATURATED"], 1)
-    new_cube = extend_saturation(
-        cube, grp, sat_circles, DQFLAGS["SATURATED"], min_sat_radius_extend, expansion=1.1
+    new_cube, persist_jumps = extend_saturation(
+        cube, grp, sat_circles, DQFLAGS["SATURATED"], DQFLAGS["JUMP_DET"],
+        1.1, persist_jumps,
     )
 
     assert new_cube[grp, 2, 2] == DQFLAGS["SATURATED"]
@@ -430,7 +432,26 @@ def test_inside_ellipes5():
     result = point_inside_ellipse(point, ellipse)
     assert result
 
-
+@pytest.mark.skip(" used for local testing")
+def test_flag_persist_groups():
+#   gdq = fits.getdata("persistgdq.fits")
+    gdq = np.zeros(shape=(2,2,2,2))
+    print(gdq.shape[0])
+    gdq = gdq[:, 0:10, :, :]
+    total_snowballs = flag_large_events(
+        gdq,
+        DQFLAGS["JUMP_DET"],
+        DQFLAGS["SATURATED"],
+        min_sat_area=1,
+        min_jump_area=6,
+        expand_factor=1.9,
+        edge_size=0,
+        sat_required_snowball=True,
+        min_sat_radius_extend=2.5,
+        sat_expand=1.1,
+        mask_persist_grps_next_int=True,
+        persist_grps_flagged=0)
+#   fits.writeto("persitflaggedgdq.fits", gdq, overwrite=True)
 def test_calc_num_slices():
     n_rows = 20
     max_available_cores = 10