Add unit test for the double flagging of pixels

spacetelescope · Jan 26, 2024 · 8f30516 · 8f30516
1 parent c5c0248
commit 8f30516
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 4 deletions.
diff --git a/src/stcal/jump/jump.py b/src/stcal/jump/jump.py
@@ -427,11 +427,11 @@ def detect_jumps(
             # save the neighbors to be flagged that will be in the next slice
             previous_row_above_gdq = row_above_gdq.copy()
             k += 1
-            # remove redundant bits in pixels that have jump flagged but were
-            # already flagged as do_not_use or saturated.
-            gdq[gdq == np.bitwise_or(dqflags['DO_NOT_USE'], dqflags['JUMP_DET'])] = \
+        # remove redundant bits in pixels that have jump flagged but were
+        # already flagged as do_not_use or saturated.
+        gdq[gdq == np.bitwise_or(dqflags['DO_NOT_USE'], dqflags['JUMP_DET'])] = \
                 dqflags['DO_NOT_USE']
-            gdq[gdq == np.bitwise_or(dqflags['SATURATED'], dqflags['JUMP_DET'])] = \
+        gdq[gdq == np.bitwise_or(dqflags['SATURATED'], dqflags['JUMP_DET'])] = \
                 dqflags['SATURATED']
 
         #  This is the flag that controls the flagging of snowballs.

diff --git a/tests/test_jump.py b/tests/test_jump.py
@@ -8,6 +8,7 @@
     find_faint_extended,
     flag_large_events,
     point_inside_ellipse,
+    detect_jumps,
 )
 
 DQFLAGS = {"JUMP_DET": 4, "SATURATED": 2, "DO_NOT_USE": 1, "GOOD": 0, "NO_GAIN_VALUE": 8}
@@ -30,6 +31,53 @@ def _cube(ngroups, readnoise=10):
     return _cube
 
 
+def test_multiprocessing():
+    nints = 1
+    nrows = 13
+    ncols = 2
+    ngroups = 13
+    readnoise = 10
+    frames_per_group = 1
+
+    data = np.zeros(shape=(nints, ngroups, nrows, ncols), dtype=np.float32)
+    readnoise_2d = np.ones((nrows, ncols), dtype=np.float32) * readnoise
+    gain_2d = np.ones((nrows, ncols), dtype=np.float32) * 4
+    gdq = np.zeros(shape=(nints, ngroups, nrows, ncols), dtype=np.uint32)
+    pdq = np.zeros(shape=(nrows, ncols), dtype=np.uint32)
+    err = np.zeros(shape=(nrows, ncols), dtype=np.float32)
+    num_cores = "1"
+    data[0, 4:, 5, 1] = 2000
+    gdq[0, 4:, 6, 1] = DQFLAGS['DO_NOT_USE']
+    gdq, pdq, total_primary_crs, number_extended_events, stddev = detect_jumps(
+        frames_per_group, data, gdq, pdq, err, gain_2d, readnoise_2d, rejection_thresh=5, three_grp_thresh=6,
+        four_grp_thresh=7, max_cores=num_cores, max_jump_to_flag_neighbors=10000, min_jump_to_flag_neighbors=100,
+        flag_4_neighbors=True, dqflags=DQFLAGS)
+    print(data[0, 4, :, :])
+    print(gdq[0, 4, :, :])
+    assert gdq[0, 4, 5, 1] == DQFLAGS['JUMP_DET']
+    assert gdq[0, 4, 6, 1] == DQFLAGS['DO_NOT_USE']
+
+    # This section of code will fail without the fixes for PR #239 that prevent
+    # the double flagging pixels with jump which already have do_not_use or saturation set.
+    num_cores = "5"
+    data = np.zeros(shape=(nints, ngroups, nrows, ncols), dtype=np.float32)
+    gdq = np.zeros(shape=(nints, ngroups, nrows, ncols), dtype=np.uint32)
+    pdq = np.zeros(shape=(nrows, ncols), dtype=np.uint32)
+    readnoise_2d = np.ones((nrows, ncols), dtype=np.float32) * readnoise
+    gain_2d = np.ones((nrows, ncols), dtype=np.float32) * 3
+    err = np.zeros(shape=(nrows, ncols), dtype=np.float32)
+    data[0, 4:, 5, 1] = 2000
+    gdq[0, 4:, 6, 1] = DQFLAGS['DO_NOT_USE']
+    gdq, pdq, total_primary_crs, number_extended_events, stddev = detect_jumps(
+        frames_per_group, data, gdq, pdq, err, gain_2d, readnoise_2d, rejection_thresh=5, three_grp_thresh=6,
+        four_grp_thresh=7, max_cores=num_cores, max_jump_to_flag_neighbors=10000, min_jump_to_flag_neighbors=100,
+        flag_4_neighbors=True, dqflags=DQFLAGS)
+    assert gdq[0, 4, 5, 1] == DQFLAGS['JUMP_DET']
+    assert gdq[0, 4, 6, 1] == DQFLAGS['DO_NOT_USE'] #This value would have been 5 without the fix.
+
+
+
+
 def test_find_simple_ellipse():
     plane = np.zeros(shape=(5, 5), dtype=np.uint8)
     plane[2, 2] = DQFLAGS["JUMP_DET"]