Repository: spacetelescope/stcal
Issue category: Performance Issue

Jump performance
Resolves [JP-3697](https://jira.stsci.edu/browse/JP-3697)

Closes #337

A number of changes to `twopoint_difference.py` and `jump.py` reduce memory usage by a factor of about 2 and improve runtime by a factor of anywhere from 3 to 20.  Tested output files (noted below) are unchanged, except for MIRI shower detection in `jw01701012001_02105_00001_mirifulong`.  This is due to a difference in replacing infilled saturated pixels after performing a convolution; the new version identifies fewer showers.  An example comment on this is here: 
https://jira.stsci.edu/browse/JP-3697?focusedId=760740&page=com.atlassian.jira.plugin.system.issuetabpanels%3Acomment-tabpanel#comment-760740
Overflagging of jumps should no longer cause the pipeline to grind to a halt due to `np.where` combined with `for` loops over pixels; it should now have virtually no impact on performance.

Files tested, along with memory usage and runtimes in the jump step (all files are public except for program 5263):

| File                                                                   |  Old Runtime | Old Memory  |   New Runtime | New Memory |
| -------------------------------------------| --------| ---------| ----| -- |
| jw01668002001_04101_00001_mirimage   | 68 s | 27 GB | 15 s | 11 GB |
| jw01283001001_03101_00001_mirimage    | 420 s | 70 GB | 120s  | 30 GB |
| jw05263001001_05101_00001_nrs1             | 500 s | 40 GB | 30 s | 23 GB |
| jw02079004001_03201_00001_nrca1          | 35 s | 2 GB | 6 s | 2 GB |
| jw01701012001_02105_00001_mirifulong     | 13 s | 0.4 GB | 4 s | 0.4 GB |
| jw01366001001_04101_00001-seg001_nis | 155 s | 40 GB | 35 s | 23 GB |

A significant fraction of the line changes are due to fixing a mixture of 8-space indents and 4-space indents in `twopoint_difference.py`.  Most of these lines were not actually changed.


<!-- if you can't perform these tasks due to permissions, please ask a maintainer to do them -->
## Tasks
- [ ] update or add relevant tests
- [ ] update relevant docstrings and / or `docs/` page
- [ ] Does this PR change any API used downstream? (if not, label with `no-changelog-entry-needed`)
  - [ ] write news fragment(s) in `changes/`: `echo "changed something" > changes/<PR#>.<changetype>.rst` (see below for change types)
  - [ ] run regression tests with this branch installed (`"git+https://github.com/<fork>/stcal@<branch>"`)
    - [ ] [`jwst` regression test](https://github.com/spacetelescope/RegressionTests/actions/workflows/jwst.yml) 
    - [ ] [`romancal` regression test](https://github.com/spacetelescope/RegressionTests/actions/workflows/romancal.yml)

<details><summary>news fragment change types...</summary>

- ``changes/<PR#>.apichange.rst``: change to public API
- ``changes/<PR#>.bugfix.rst``: fixes an issue
- ``changes/<PR#>.general.rst``: infrastructure or miscellaneous change
</details>



Functions to edit:
src/stcal/jump/jump.py:extend_saturation
src/stcal/jump/jump.py:extend_ellipses
src/stcal/jump/jump.py:find_last_grp
src/stcal/jump/jump.py:find_faint_extended
src/stcal/jump/jump.py:get_bigcontours
src/stcal/jump/jump.py:diff_meddiff_int
src/stcal/jump/jump.py:diff_meddiff_grp
src/stcal/jump/twopoint_difference.py:find_crs_old
Functions added: src/stcal/jump/jump.py:ellipse_subim
src/stcal/jump/jump.py:convolve_fast
src/stcal/jump/twopoint_difference.py:propagate_flags

Patch:
diff --git a/changes/337.general.rst b/changes/337.general.rst
new file mode 100644
index 00000000..f73f7e3a
--- /dev/null
+++ b/changes/337.general.rst
@@ -0,0 +1,1 @@
+Performance improvements for jump step targeting both runtime and memory consumption. Results are mostly identical, but there are some differences in the MIRI shower flagging (it now flags somewhat fewer showers). 
diff --git a/src/stcal/jump/jump.py b/src/stcal/jump/jump.py
index 47105e3c..6c89ccb0 100644
--- a/src/stcal/jump/jump.py
+++ b/src/stcal/jump/jump.py
@@ -5,6 +5,7 @@
 import multiprocessing
 import time
 import warnings
+from scipy import signal
 
 import numpy as np
 import cv2 as cv
@@ -402,9 +403,9 @@ def flag_large_events(gdq, jump_flag, sat_flag, jump_data):
                 expansion=jump_data.expand_factor, num_grps_masked=0,
             )
 
-    #  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.
+    #  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 jump_data.mask_persist_grps_next_int:
         for intg in range(1, nints):
             if jump_data.persist_grps_flagged >= 1:
@@ -445,9 +446,6 @@ def extend_saturation(cube, grp, sat_ellipses, jump_data, persist_jumps):
         3D (nints, nrows, ncols) uint8
     """
     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()
     satcolor = 22  # (0, 0, 22) is a dark blue in RGB
     for ellipse in sat_ellipses:
         ceny = ellipse[0][0]
@@ -460,31 +458,92 @@ def extend_saturation(cube, grp, sat_ellipses, jump_data, persist_jumps):
             axis2 = ellipse[1][1] + jump_data.sat_expand
             axis1 = min(axis1, jump_data.max_extended_radius)
             axis2 = min(axis2, jump_data.max_extended_radius)
-            axes = (round(axis1 / 2), round(axis2 / 2))
 
             alpha = ellipse[2]
-            color = (0, 0, satcolor)  # in the RGB cube, set blue plane pixels of the ellipse to 22
-            image = cv.ellipse(image, cen, axes, alpha, 0, 360, color, -1,)
 
-            #  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.
+            indx, sat_ellipse = ellipse_subim(
+                ceny, cenx, axis1, axis2, alpha, satcolor, (nrows, ncols))
+            (iy1, iy2, ix1, ix2) = indx
 
-            # extract the Blue plane of the image
-            sat_ellipse = image[:, :, 2]  
+            # 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.
 
-            # find all the ellipse pixels in the ellipse
-            saty, satx = np.where(sat_ellipse == 22)  
+            is_sat = sat_ellipse == satcolor
+            for i in range(grp, cube.shape[0]):
+                cube[i][iy1:iy2, ix1:ix2][is_sat] = jump_data.fl_sat
 
-            outcube[grp:, saty, satx] = jump_data.fl_sat
-            axes = (round(ellipse[1][0] / 2), round(ellipse[1][1] / 2))
-            persiste_image = cv.ellipse(persist_image, cen, axes, alpha, 0, 360, color, -1,)
+            ax1, ax2 = (ellipse[1][0], ellipse[1][1])
+            indx, persist_ellipse = ellipse_subim(
+                ceny, cenx, ax1, ax2, alpha, satcolor, (nrows, ncols))
+            (iy1, iy2, ix1, ix2) = indx
 
-            persist_ellipse = persist_image[:, :, 2]
-            persist_saty, persist_satx = np.where(persist_ellipse == 22)
-            persist_jumps[persist_saty, persist_satx] = jump_data.fl_jump
+            persist_mask = persist_ellipse == satcolor
+            persist_jumps[iy1:iy2, ix1:ix2][persist_mask] = jump_data.fl_jump
+
+    return cube, persist_jumps
+
+
+def ellipse_subim(ceny, cenx, axis1, axis2, alpha, value, shape):
+    """Draw a filled ellipse in a small array at a given (returned) location
+    Parameters
+    ----------
+    ceny : float
+        Center of the ellipse in y (second axis of an image)
+    cenx : float
+        Center of the ellipse in x (first axis of an image)
+    axis1 : float
+        One (full) axis of the ellipse
+    axis2 : float
+        The other (full) axis of the ellipse
+    alpha : float
+        Angle (in degrees) between axis1 and x
+    value : unsigned 8-bit integer
+        Value to fill the image with
+    shape : (int, int)
+        The shape of the full 2D array into which the returned
+        subimage should be placed.
+    Returns
+    -------
+    indx : (int, int, int, int)
+        Indices (iy1, iy2, ix1, ix2) such that
+        fullimage[iy1:iy2, ix1:ix2] = subimage (see below)
+    subimage : 2D 8-bit unsigned int array
+        Small image containing the ellipse, goes into fullimage
+        as described above.
+    """
+    yc, xc = round(ceny), round(cenx)
+
+    # How big of a subarray do we need for the subimage?
+
+    dn_over_2 = max(round(axis1/2), round(axis2/2)) + 2
+
+    # Note that the convention between which index is x and which
+    # is y is a little confusing here.  To cv.ellipse, the first
+    # coordinate corresponds to the second Python index.  That is
+    # why x and y are a bit mixed up below.
+
+    ix1 = max(yc - dn_over_2, 0)
+    ix2 = min(yc + dn_over_2 + 1, shape[1])
+    iy1 = max(xc - dn_over_2, 0)
+    iy2 = min(xc + dn_over_2 + 1, shape[0])
+
+    image = np.zeros(shape=(iy2 - iy1, ix2 - ix1, 3), dtype=np.uint8)
+    image = cv.ellipse(
+        image,
+        (yc - ix1, xc - iy1),
+        (round(axis1 / 2), round(axis2 / 2)),
+        alpha,
+        0,
+        360,
+        (0, 0, value),
+        -1,
+    )
+
+    # The last ("blue") part contains the filled ellipse that we want.
+    subimage = image[:, :, 2]
+    return (iy1, iy2, ix1, ix2), subimage
 
-    return outcube, persist_jumps
 
 
 def extend_ellipses(
@@ -497,7 +556,7 @@ def extend_ellipses(
     Parameters
     ----------
     gdq_cube : ndarray
-        Group DQ cube for an integration.
+        Group DQ cube for an integration.  Modified in-place.
 
     intg : int
         The current integration.
@@ -522,8 +581,8 @@ def extend_ellipses(
 
     Returns
     -------
-    out_gdq_cube : ndarray
-        Computed 3-D group DQ array.
+    gdq_cube : ndarray
+        Computed 3-D group DQ array, modified in-place
 
     num_ellipses : int
         The number of ellipses passed in as a parameter.
@@ -531,9 +590,7 @@ 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.
-    out_gdq_cube = gdq_cube.copy()
-    _, _, nrows, ncols = gdq_cube.shape
-    image = np.zeros(shape=(nrows, ncols, 3), dtype=np.uint8)
+    _, ngroups, nrows, ncols = gdq_cube.shape
     num_ellipses = len(ellipses)
     for ellipse in ellipses:
         ceny = ellipse[0][0]
@@ -541,48 +598,26 @@ def extend_ellipses(
         axes = compute_axes(expand_by_ratio, ellipse, expansion, jump_data)
 
         alpha = ellipse[2]
-        cen = (round(ceny), round(cenx))
-        color = (0, 0, jump_data.fl_jump)
-        image = cv.ellipse(image, cen, axes, alpha, 0, 360, color, -1)
-
-        jump_ellipse = image[:, :, 2]
-        ngrps = gdq_cube.shape[1]
-        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, :, :], jump_data.fl_sat)
-            jump_ellipse[sat_pix == jump_data.fl_sat] = 0
-            out_gdq_cube[intg, flg_grp, :, :] = np.bitwise_or(gdq_cube[intg, flg_grp, :, :], jump_ellipse)
-
-    return out_gdq_cube, num_ellipses
-
-
-def find_last_grp(grp, ngrps, num_grps_masked):
-    """
-    Find the last group based on current group and number of groups masked.
-
-    Parameters
-    ----------
-    grp : int
-        The location of the shower
-
-    ngrps : int
-        The number of groups in the integration
+        # Get the expanded ellipse in a subimage, along with the
+        # indices that place this subimage within the full array.
+        axis1 = axes[0]*2
+        axis2 = axes[1]*2
+        indx, jump_ellipse = ellipse_subim(
+            ceny, cenx, axis1, axis2, alpha, jump_data.fl_jump, (nrows, ncols))
+        (iy1, iy2, ix1, ix2) = indx
+        
+        # Propagate forward by num_grps_masked groups.
 
-    num_grps_masked : int
-        The requested number of groups to be flagged after the shower
+        for flg_grp in range(grp, min(grp + num_grps_masked + 1, ngroups)):
 
-    Returns
-    -------
-    last_grp : int
-        The index of the last group to flag for the shower
+            # Only propagate the snowball forward to unsaturated pixels.
 
-    """
-    num_grps_masked += 1
-    last_grp = min(grp + num_grps_masked, ngrps)
-    return last_grp
+            sat_pix = gdq_cube[intg, flg_grp, iy1:iy2, ix1:ix2] & jump_data.fl_sat
+            jump_ellipse[sat_pix == jump_data.fl_sat] = 0
+            gdq_cube[intg, flg_grp, iy1:iy2, ix1:ix2] |= jump_ellipse
 
+    return gdq_cube, num_ellipses
 
 def find_ellipses(dqplane, bitmask, min_area):
     """
@@ -796,21 +831,19 @@ def find_faint_extended(
 
     all_ellipses = []
 
-    first_diffs_masked = np.ma.masked_array(first_diffs, mask=np.isnan(first_diffs))
     warnings.filterwarnings("ignore")
 
     read_noise_2 = readnoise_2d**2
     if nints >= jump_data.minimum_sigclip_groups:
-        mean, median, stddev = stats.sigma_clipped_stats(first_diffs_masked, sigma=5, axis=0)
+        mean, median, stddev = stats.sigma_clipped_stats(first_diffs, sigma=5, axis=0)
     else:
-        median_diffs = np.nanmedian(first_diffs_masked, axis=(0, 1))
+        median_diffs = np.nanmedian(first_diffs, axis=(0, 1))
         sigma = np.sqrt(np.abs(median_diffs) + read_noise_2 / jump_data.nframes)
 
     for intg in range(nints):
         if nints < jump_data.minimum_sigclip_groups:
             # The difference from the median difference for each group
-            median_diffs, ratio = diff_meddiff_int(
-                    intg, median_diffs, sigma, first_diffs_masked)
+            ratio = diff_meddiff_int(intg, median_diffs, sigma, first_diffs)
 
         #  The convolution kernel creation
         ring_2D_kernel = Ring2DKernel(
@@ -818,8 +851,7 @@ def find_faint_extended(
         first_good_group = find_first_good_group(gdq[intg, :, :, :], jump_data.fl_dnu)
         for grp in range(first_good_group + 1, ngrps):
             if nints >= jump_data.minimum_sigclip_groups:
-                median_diffs, ratio = diff_meddiff_grp(
-                        intg, grp, median, stddev, first_diffs_masked)
+                ratio = diff_meddiff_grp(intg, grp, median, stddev, first_diffs)
 
             bigcontours = get_bigcontours(
                     ratio, intg, grp, gdq, pdq, jump_data, ring_2D_kernel)
@@ -1070,25 +1102,19 @@ def get_bigcontours(ratio, intg, grp, gdq, pdq, jump_data, ring_2D_kernel):
     sat_flag = jump_data.fl_sat
     dnu_flag = jump_data.fl_dnu
 
-    #  mask pixels that are already flagged as jump
+    #  mask pixels that are already flagged as jump, sat, or dnu
     combined_pixel_mask = np.bitwise_or(gdq[intg, grp, :, :], pdq[:, :])
-    jump_pixels_array = np.bitwise_and(combined_pixel_mask, jump_flag)
-    masked_ratio[jump_pixels_array == jump_flag] = np.nan
 
-    #  mask pixels that are already flagged as sat.
-    sat_pixels_array = np.bitwise_and(combined_pixel_mask, sat_flag)
-    masked_ratio[sat_pixels_array == sat_flag] = np.nan
-
-    #  mask pixels that are already flagged as do not use
-    dnu_pixels_array = np.bitwise_and(combined_pixel_mask, dnu_flag)
-    dnuy, dnux = np.where(dnu_pixels_array == dnu_flag)  # dnuy, dnux used twice
-    masked_ratio[dnuy, dnux] = np.nan
-
-    masked_smoothed_ratio = convolve(masked_ratio.filled(np.nan), ring_2D_kernel)
-    del masked_ratio
+    jump_sat_or_dnu = np.bitwise_and(combined_pixel_mask, jump_flag|sat_flag|dnu_flag) != 0
+    masked_ratio[jump_sat_or_dnu] = np.nan
+    
+    kernel = ring_2D_kernel.array
+    
+    # Equivalent to but faster than
+    # masked_smoothed_ratio = convolve(masked_ratio, ring_2D_kernel, preserve_nan=True)
+    
+    masked_smoothed_ratio = convolve_fast(masked_ratio, kernel)
 
-    #  mask out the pixels that got refilled by the convolution
-    masked_smoothed_ratio[dnuy, dnux] = np.nan
     extended_emission = (masked_smoothed_ratio > jump_data.extend_snr_threshold).astype(np.uint8)
 
     #  find the contours of the extended emission
@@ -1100,6 +1126,65 @@ def get_bigcontours(ratio, intg, grp, gdq, pdq, jump_data, ring_2D_kernel):
     return bigcontours 
 
 
+
+def convolve_fast(inarray, kernel, copy=False):
+    """Convolve an array with a kernel, interpolating over NaNs.
+    Faster version of astropy.convolution.convolve(preserve_nan=True)
+    Parameters
+    ----------
+    inarray : 2D array of floats
+        Array for convolution
+    kernel : 2D array of floats
+        Convolution kernel.  Both dimensions must be odd.
+    copy : bool
+        Make a copy of inarray to avoid modifying NaN values.  Default False.
+    Returns
+    -------
+    convolved_array : 2D array of floats
+        Convolution of inarray and kernel, interpolating over NaNs.
+    """
+
+    # We will mask nan pixels by setting them to zero.  We
+    # will convolve by our kernel, then divide by the weight
+    # given by the valid pixels convolved with the kernel in
+    # order to normalize.  Finally, we will reset the
+    # initially nan pixels to nan.
+    #
+    # This function is equivalent to
+    # convolved_array = astropy.convolution.convolve(inarray, kernel, preserve_nan=True)
+    # but runs in about half the time.
+
+    if copy:
+        array = inarray.copy()
+    else:
+        array = inarray
+
+    good = np.isfinite(array)
+    array[~good] = 0
+
+    convolved_array = signal.oaconvolve(array, kernel, mode='same')
+
+    # Embed the flag in a larger array to reproduce the behavior at
+    # the edge with a fill value of zero.
+
+    padded_good_arr = np.ones((good.shape[0] + kernel.shape[0] - 1,
+                               good.shape[1] + kernel.shape[1] - 1))
+    n = kernel.shape[0]//2
+    padded_good_arr[n:-n, n:-n] = good
+    norm = signal.oaconvolve(padded_good_arr, kernel, mode='valid')
+
+    # Avoid dividing by a tiny number due to roundoff error.
+
+    good &= norm > 1e-3*np.mean(kernel)
+    convolved_array /= norm
+
+    # Replace NaNs
+
+    convolved_array[~good] = np.nan
+
+    return convolved_array
+
+
 def diff_meddiff_int(intg, median_diffs, sigma, first_diffs_masked):
     """
     Compute the SNR ratio of each difference.
@@ -1120,26 +1205,16 @@ def diff_meddiff_int(intg, median_diffs, sigma, first_diffs_masked):
 
     Returns
     -------
-    median_diffs : ndarray
-        Median of first differences
-
     ratio : ndarray
         SNR ratio
     """
-    if intg > 0:
-        e_jump = first_diffs_masked[intg] - median_diffs[np.newaxis, :, :]
 
-        # SNR ratio of each diff.
-        ratio = np.abs(e_jump) / sigma[np.newaxis, :, :]
-    else:
-        # The difference from the median difference for each group
-        e_jump = first_diffs_masked[intg] - median_diffs[np.newaxis, :, :]
+    e_jump = first_diffs_masked[intg] - median_diffs[np.newaxis, :, :]
 
-        # SNR ratio of each diff.
-        ratio = np.abs(e_jump) / sigma[np.newaxis, :, :]
-        median_diffs = np.nanmedian(first_diffs_masked, axis=(0, 1))
+    # SNR ratio of each diff.
+    ratio = np.abs(e_jump) / sigma[np.newaxis, :, :]
 
-    return median_diffs, ratio
+    return ratio
 
 
 def diff_meddiff_grp(intg, grp, median, stddev, first_diffs_masked):
@@ -1165,9 +1240,6 @@ def diff_meddiff_grp(intg, grp, median, stddev, first_diffs_masked):
 
     Returns
     -------
-    median_diffs : ndarray
-        Median of first differences
-
     ratio : ndarray
         SNR ratio
     """
@@ -1180,7 +1252,7 @@ def diff_meddiff_grp(intg, grp, median, stddev, first_diffs_masked):
     # SNR ratio of each diff.
     ratio = np.abs(e_jump) / sigma[np.newaxis, :, :]
 
-    return median_diffs, ratio
+    return ratio
 
 
 def nan_invalid_data(data, gdq, jump_data):
diff --git a/src/stcal/jump/twopoint_difference.py b/src/stcal/jump/twopoint_difference.py
index ade8dc6c..511761bb 100644
--- a/src/stcal/jump/twopoint_difference.py
+++ b/src/stcal/jump/twopoint_difference.py
@@ -4,6 +4,7 @@
 import numpy as np
 import warnings
 from astropy import stats
+from astropy.utils.exceptions import AstropyUserWarning
 
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
@@ -244,29 +245,23 @@ def find_crs_old(
         return gdq, row_below_gdq, row_above_gdq, 0, dummy
     else:
         # set 'saturated' or 'do not use' pixels to nan in data
-        dat[np.where(np.bitwise_and(gdq, sat_flag))] = np.nan
-        dat[np.where(np.bitwise_and(gdq, dnu_flag))] = np.nan
-        dat[np.where(np.bitwise_and(gdq, dnu_flag + sat_flag))] = np.nan
-
+        dat[gdq & (dnu_flag | sat_flag) != 0] = np.nan
+        
         # calculate the differences between adjacent groups (first diffs)
-        # use mask on data, so the results will have sat/donotuse groups masked
+        # Bad data will be NaN; np.nanmedian will be used later.
         first_diffs = np.diff(dat, axis=1)
-
+        first_diffs_finite = np.isfinite(first_diffs)
+        
         # calc. the median of first_diffs for each pixel along the group axis
-        first_diffs_masked = np.ma.masked_array(first_diffs, mask=np.isnan(first_diffs))
-        median_diffs = np.ma.median(first_diffs_masked, axis=(0, 1))
+        warnings.filterwarnings("ignore", ".*All-NaN slice encountered.*", RuntimeWarning)
+        median_diffs = np.nanmedian(first_diffs, axis=(0, 1))
+        warnings.resetwarnings()
         # calculate sigma for each pixel
         sigma = np.sqrt(np.abs(median_diffs) + read_noise_2 / nframes)
 
         # reset sigma so pxels with 0 readnoise are not flagged as jumps
-        sigma[np.where(sigma == 0.)] = np.nan
-
-        # compute 'ratio' for each group. this is the value that will be
-        # compared to 'threshold' to classify jumps. subtract the median of
-        # first_diffs from first_diffs, take the absolute value and divide by sigma.
-        e_jump_4d = first_diffs - median_diffs[np.newaxis, :, :]
-        ratio_all = np.abs(first_diffs - median_diffs[np.newaxis, np.newaxis, :, :]) / \
-                    sigma[np.newaxis, np.newaxis, :, :]
+        sigma[sigma == 0.] = np.nan
+
         # Test to see if there are enough groups to use sigma clipping
         if (only_use_ints and nints >= minimum_sigclip_groups) or \
            (not only_use_ints and total_groups >= minimum_sigclip_groups):
@@ -275,241 +270,200 @@ def find_crs_old(
             warnings.filterwarnings("ignore", ".*All-NaN slice encountered.*", RuntimeWarning)
             warnings.filterwarnings("ignore", ".*Mean of empty slice.*", RuntimeWarning)
             warnings.filterwarnings("ignore", ".*Degrees of freedom <= 0.*", RuntimeWarning)
+            warnings.filterwarnings("ignore", ".*Input data contains invalid values*", AstropyUserWarning)
 
             if only_use_ints:
-                mean, median, stddev = stats.sigma_clipped_stats(first_diffs_masked, sigma=normal_rej_thresh,
-                                                                 axis=0)
-                clipped_diffs = stats.sigma_clip(first_diffs_masked, sigma=normal_rej_thresh,
-                                                 axis=0, masked=True)
+                clipped_diffs, alow, ahigh = stats.sigma_clip(
+                    first_diffs, sigma=normal_rej_thresh,
+                    axis=0, masked=True, return_bounds=True)
             else:
-                mean, median, stddev = stats.sigma_clipped_stats(first_diffs_masked, sigma=normal_rej_thresh,
-                                                                 axis=(0, 1))
-                clipped_diffs = stats.sigma_clip(first_diffs_masked, sigma=normal_rej_thresh,
-                                                 axis=(0, 1), masked=True)
-            jump_mask = np.logical_and(clipped_diffs.mask, np.logical_not(first_diffs_masked.mask))
-            jump_mask[np.bitwise_and(jump_mask, gdq[:, 1:, :, :] == sat_flag)] = False
-            jump_mask[np.bitwise_and(jump_mask, gdq[:, 1:, :, :] == dnu_flag)] = False
-            jump_mask[np.bitwise_and(jump_mask, gdq[:, 1:, :, :] == (dnu_flag + sat_flag))] = False
-            gdq[:, 1:, :, :] = np.bitwise_or(gdq[:, 1:, :, :], jump_mask *
-                                             np.uint8(dqflags["JUMP_DET"]))
+                clipped_diffs, alow, ahigh = stats.sigma_clip(
+                    first_diffs, sigma=normal_rej_thresh,
+                    axis=(0, 1), masked=True, return_bounds=True)
+            # get the standard deviation from the bounds of sigma clipping
+            stddev = 0.5*(ahigh - alow)/normal_rej_thresh
+            jump_candidates = clipped_diffs.mask
+            sat_or_dnu_not_set = gdq[:, 1:] & (sat_flag | dnu_flag) == 0
+            jump_mask = jump_candidates & first_diffs_finite & sat_or_dnu_not_set
+            del clipped_diffs
+            gdq[:, 1:] |= jump_mask * np.uint8(jump_flag)
+
             # if grp is all jump set to do not use
             for integ in range(nints):
                 for grp in range(ngrps):
-                    if np.all(np.bitwise_or(np.bitwise_and(gdq[integ, grp, :, :], jump_flag),
-                                            np.bitwise_and(gdq[integ, grp, :, :], dnu_flag))):
-                        jumpy, jumpx = np.where(gdq[integ, grp, :, :] == jump_flag)
-                        gdq[integ, grp, jumpy, jumpx] = 0
+                    if np.all(gdq[integ, grp] & (jump_flag | dnu_flag) != 0):
+                        # The line below matches the comment above, but not the
+                        # old logic.  Leaving it for now.
+                        #gdq[integ, grp] |= dnu_flag
+                        
+                        jump_only = gdq[integ, grp, :, :] == jump_flag
+                        gdq[integ, grp][jump_only] = 0
+                        
             warnings.resetwarnings()
         else:  # There are not enough groups for sigma clipping
 
-                # set 'saturated' or 'do not use' pixels to nan in data
-                dat[np.where(np.bitwise_and(gdq, sat_flag))] = np.nan
-                dat[np.where(np.bitwise_and(gdq, dnu_flag))] = np.nan
-
-                # calculate the differences between adjacent groups (first diffs)
-                # use mask on data, so the results will have sat/donotuse groups masked
-                first_diffs = np.diff(dat, axis=1)
-
-                if total_usable_diffs >= min_diffs_single_pass:
-                    warnings.filterwarnings("ignore", ".*All-NaN slice encountered.*", RuntimeWarning)
-                    median_diffs = np.nanmedian(first_diffs, axis=(0, 1))
-                    warnings.resetwarnings()
-                    # calculate sigma for each pixel
-                    sigma = np.sqrt(np.abs(median_diffs) + read_noise_2 / nframes)
-                    # reset sigma so pixels with 0 read noise are not flagged as jumps
-                    sigma[np.where(sigma == 0.)] = np.nan
-
-                    # compute 'ratio' for each group. this is the value that will be
-                    # compared to 'threshold' to classify jumps. subtract the median of
-                    # first_diffs from first_diffs, take the abs. value and divide by sigma.
-                    e_jump = first_diffs - median_diffs[np.newaxis, np.newaxis, :, :]
-
-                    ratio = np.abs(e_jump) / sigma[np.newaxis, np.newaxis, :, :]
-                    # XXX Increased memory consumption with np.ma.masked_greater
-                    masked_ratio = np.ma.masked_greater(ratio, normal_rej_thresh)
-                    #  The jump mask is the ratio greater than the threshold and the difference is usable
-                    jump_mask = np.logical_and(masked_ratio.mask, np.logical_not(first_diffs_masked.mask))
-                    gdq[:, 1:, :, :] = np.bitwise_or(gdq[:, 1:, :, :], jump_mask *
-                                                     np.uint8(dqflags["JUMP_DET"]))
-                else:  # low number of diffs requires iterative flagging
-                    # calculate the differences between adjacent groups (first diffs)
-                    # use mask on data, so the results will have sat/donotuse groups masked
-                    first_diffs = np.abs(np.diff(dat, axis=1))
-
-                    # calc. the median of first_diffs for each pixel along the group axis
-                    median_diffs = calc_med_first_diffs(first_diffs)
-
-                    # calculate sigma for each pixel
-                    sigma = np.sqrt(np.abs(median_diffs) + read_noise_2 / nframes)
-                    # reset sigma so pxels with 0 readnoise are not flagged as jumps
-                    sigma[np.where(sigma == 0.0)] = np.nan
-
-                    # compute 'ratio' for each group. this is the value that will be
-                    # compared to 'threshold' to classify jumps. subtract the median of
-                    # first_diffs from first_diffs, take the abs. value and divide by sigma.
-                    e_jump = first_diffs - median_diffs[np.newaxis, :, :]
-                    ratio = np.abs(e_jump) / sigma[np.newaxis, :, :]
-
-                    # create a 2d array containing the value of the largest 'ratio' for each pixel
-                    warnings.filterwarnings("ignore", ".*All-NaN slice encountered.*", RuntimeWarning)
-                    max_ratio = np.nanmax(ratio, axis=1)
-                    warnings.resetwarnings()
-                    # now see if the largest ratio of all groups for each pixel exceeds the threshold.
-                    # there are different threshold for 4+, 3, and 2 usable groups
-                    num_unusable_groups = np.sum(np.isnan(first_diffs), axis=(0, 1))
-                    int4cr, row4cr, col4cr = np.where(
-                        np.logical_and(ndiffs - num_unusable_groups >= 4, max_ratio > normal_rej_thresh)
-                    )
-                    int3cr, row3cr, col3cr = np.where(
-                        np.logical_and(ndiffs - num_unusable_groups == 3, max_ratio > three_diff_rej_thresh)
-                    )
-                    int2cr, row2cr, col2cr = np.where(
-                        np.logical_and(ndiffs - num_unusable_groups == 2, max_ratio > two_diff_rej_thresh)
-                    )
-                    # get the rows, col pairs for all pixels with at least one CR
+            if total_usable_diffs >= min_diffs_single_pass:
+
+                # compute 'ratio' for each group. this is the value that will be
+                # compared to 'threshold' to classify jumps. subtract the median of
+                # first_diffs from first_diffs, take the abs. value and divide by sigma.
+                # The jump mask is the ratio greater than the threshold and the
+                # difference is usable.  Loop over integrations to minimize the memory
+                # footprint.
+                jump_mask = np.zeros(first_diffs.shape, dtype=bool)
+                for i in range(nints):
+                    absdiff = np.abs(first_diffs[i] - median_diffs[np.newaxis, :])
+                    ratio = absdiff / sigma[np.newaxis, :]
+                    jump_candidates = ratio > normal_rej_thresh
+                    jump_mask = jump_candidates & first_diffs_finite[i]
+                    gdq[i, 1:] |= jump_mask * np.uint8(jump_flag)
+                    
+            else:  # low number of diffs requires iterative flagging
+                
+                # calc. the median of first_diffs for each pixel along the group axis
+                # Do not overwrite first_diffs, median_diffs, sigma.
+                first_diffs_abs = np.abs(first_diffs)
+                median_diffs_iter = calc_med_first_diffs(first_diffs_abs)
+
+                # calculate sigma for each pixel
+                sigma_iter = np.sqrt(np.abs(median_diffs_iter) + read_noise_2 / nframes)
+                # reset sigma so pxels with 0 readnoise are not flagged as jumps
+                sigma_iter[sigma_iter == 0.0] = np.nan
+
+                # compute 'ratio' for each group. this is the value that will be
+                # compared to 'threshold' to classify jumps. subtract the median of
+                # first_diffs from first_diffs, take the abs. value and divide by sigma.
+                e_jump = first_diffs_abs - median_diffs_iter[np.newaxis, :, :]
+                ratio = np.abs(e_jump) / sigma_iter[np.newaxis, :, :]
+                # create a 2d array containing the value of the largest 'ratio' for each pixel
+                warnings.filterwarnings("ignore", ".*All-NaN slice encountered.*", RuntimeWarning)
+                max_ratio = np.nanmax(ratio, axis=1)
+                warnings.resetwarnings()
+                # now see if the largest ratio of all groups for each pixel exceeds the threshold.
+                # there are different threshold for 4+, 3, and 2 usable groups
+                num_unusable_groups = np.sum(np.isnan(first_diffs_abs), axis=(0, 1))
+                int4cr, row4cr, col4cr = np.where(
+                    np.logical_and(ndiffs - num_unusable_groups >= 4, max_ratio > normal_rej_thresh)
+                )
+                int3cr, row3cr, col3cr = np.where(
+                    np.logical_and(ndiffs - num_unusable_groups == 3, max_ratio > three_diff_rej_thresh)
+                )
+                int2cr, row2cr, col2cr = np.where(
+                    np.logical_and(ndiffs - num_unusable_groups == 2, max_ratio > two_diff_rej_thresh)
+                )
+                # get the rows, col pairs for all pixels with at least one CR
 #                    all_crs_int = np.concatenate((int4cr, int3cr, int2cr))
-                    all_crs_row = np.concatenate((row4cr, row3cr, row2cr))
-                    all_crs_col = np.concatenate((col4cr, col3cr, col2cr))
-
-                    # iterate over all groups of the pix w/ an initial CR to look for subsequent CRs
-                    # flag and clip the first CR found. recompute median/sigma/ratio
-                    # and repeat the above steps of comparing the max 'ratio' for each pixel
-                    # to the threshold to determine if another CR can be flagged and clipped.
-                    # repeat this process until no more CRs are found.
-                    for j in range(len(all_crs_row)):
-                        # get arrays of abs(diffs), ratio, readnoise for this pixel
-                        pix_first_diffs = first_diffs[:, :, all_crs_row[j], all_crs_col[j]]
-                        pix_ratio = ratio[:, :, all_crs_row[j], all_crs_col[j]]
-                        pix_rn2 = read_noise_2[all_crs_row[j], all_crs_col[j]]
-    
-                        # Create a mask to flag CRs. pix_cr_mask = 0 denotes a CR
-                        pix_cr_mask = np.ones(pix_first_diffs.shape, dtype=bool)
-
-                        # set the largest ratio as a CR
-                        location = np.unravel_index(np.nanargmax(pix_ratio), pix_ratio.shape)
-                        pix_cr_mask[location] = 0
-                        new_CR_found = True
-
-                        # loop and check for more CRs, setting the mask as you go and
-                        # clipping the group with the CR. stop when no more CRs are found
-                        # or there is only one two diffs left (which means there is
-                        # actually one left, since the next CR will be masked after
-                        # checking that condition)
-                        while new_CR_found and (ndiffs - np.sum(np.isnan(pix_first_diffs)) > 2):
-                            new_CR_found = False
-
-                            # set CRs to nans in first diffs to clip them
-                            pix_first_diffs[~pix_cr_mask] = np.nan
-
-                            # recalculate median, sigma, and ratio
-                            new_pix_median_diffs = calc_med_first_diffs(pix_first_diffs)
-
-                            new_pix_sigma = np.sqrt(np.abs(new_pix_median_diffs) + pix_rn2 / nframes)
-                            new_pix_ratio = np.abs(pix_first_diffs - new_pix_median_diffs) / new_pix_sigma
-
-                            # check if largest ratio exceeds threshold appropriate for num remaining groups
-
-                            # select appropriate thresh. based on number of remaining groups
-                            rej_thresh = normal_rej_thresh
-                            if ndiffs - np.sum(np.isnan(pix_first_diffs)) == 3:
-                                rej_thresh = three_diff_rej_thresh
-                            if ndiffs - np.sum(np.isnan(pix_first_diffs)) == 2:
-                                rej_thresh = two_diff_rej_thresh
-                            max_idx = np.nanargmax(new_pix_ratio)
-                            location = np.unravel_index(max_idx, new_pix_ratio.shape)
-                            if new_pix_ratio[location] > rej_thresh:
-                                new_CR_found = True
-                                pix_cr_mask[location] = 0
-                            unusable_diffs = np.sum(np.isnan(pix_first_diffs))
-                        # Found all CRs for this pix - set flags in input DQ array
-                        gdq[:, 1:, all_crs_row[j], all_crs_col[j]] = np.bitwise_or(
-                             gdq[:, 1:, all_crs_row[j], all_crs_col[j]],
-                            dqflags["JUMP_DET"] * np.invert(pix_cr_mask),
-                        )
-    cr_integ, cr_group, cr_row, cr_col = np.where(np.bitwise_and(gdq, jump_flag))
-    num_primary_crs = len(cr_group)
-    if flag_4_neighbors:  # iterate over each 'jump' pixel
-        for j in range(len(cr_group)):
-            ratio_this_pix = ratio_all[cr_integ[j], cr_group[j] - 1, cr_row[j], cr_col[j]]
-
-            # Jumps must be in a certain range to have neighbors flagged
-            if (ratio_this_pix < max_jump_to_flag_neighbors) and (
-                ratio_this_pix > min_jump_to_flag_neighbors
-            ):
-                integ = cr_integ[j]
-                group = cr_group[j]
-                row = cr_row[j]
-                col = cr_col[j]
-
-                # This section saves flagged neighbors that are above or
-                # below the current range of row. If this method
-                # running in a single process, the row above and below are
-                # not used. If it is running in multiprocessing mode, then
-                # the rows above and below need to be returned to
-                # find_jumps to use when it reconstructs the full group dq
-                # array from the slices.
-
-                # Only flag adjacent pixels if they do not already have the
-                # 'SATURATION' or 'DONOTUSE' flag set
-                if row != 0:
-                    if (gdq[integ, group, row - 1, col] & sat_flag) == 0 and (
-                        gdq[integ, group, row - 1, col] & dnu_flag
-                    ) == 0:
-                        gdq[integ, group, row - 1, col] = np.bitwise_or(
-                            gdq[integ, group, row - 1, col], jump_flag
-                        )
-                else:
-                    row_below_gdq[integ, cr_group[j], cr_col[j]] = jump_flag
-
-                if row != nrows - 1:
-                    if (gdq[integ, group, row + 1, col] & sat_flag) == 0 and (
-                        gdq[integ, group, row + 1, col] & dnu_flag
-                    ) == 0:
-                        gdq[integ, group, row + 1, col] = np.bitwise_or(
-                            gdq[integ, group, row + 1, col], jump_flag
-                        )
-                else:
-                    row_above_gdq[integ, cr_group[j], cr_col[j]] = jump_flag
-
-                # Here we are just checking that we don't flag neighbors of
-                # jumps that are off the detector.
-                if (
-                    cr_col[j] != 0
-                    and (gdq[integ, group, row, col - 1] & sat_flag) == 0
-                    and (gdq[integ, group, row, col - 1] & dnu_flag) == 0
-                ):
-                    gdq[integ, group, row, col - 1] = np.bitwise_or(
-                        gdq[integ, group, row, col - 1], jump_flag
+                all_crs_row = np.concatenate((row4cr, row3cr, row2cr))
+                all_crs_col = np.concatenate((col4cr, col3cr, col2cr))
+
+                # iterate over all groups of the pix w/ an initial CR to look for subsequent CRs
+                # flag and clip the first CR found. recompute median/sigma/ratio
+                # and repeat the above steps of comparing the max 'ratio' for each pixel
+                # to the threshold to determine if another CR can be flagged and clipped.
+                # repeat this process until no more CRs are found.
+                for j in range(len(all_crs_row)):
+                    # get arrays of abs(diffs), ratio, readnoise for this pixel.
+                    pix_first_diffs = first_diffs_abs[:, :, all_crs_row[j], all_crs_col[j]]
+                    pix_ratio = ratio[:, :, all_crs_row[j], all_crs_col[j]]
+                    pix_rn2 = read_noise_2[all_crs_row[j], all_crs_col[j]]
+
+                    # Create a mask to flag CRs. pix_cr_mask = 0 denotes a CR
+                    pix_cr_mask = np.ones(pix_first_diffs.shape, dtype=bool)
+
+                    # set the largest ratio as a CR
+                    location = np.unravel_index(np.nanargmax(pix_ratio), pix_ratio.shape)
+                    pix_cr_mask[location] = 0
+                    new_CR_found = True
+
+                    # loop and check for more CRs, setting the mask as you go and
+                    # clipping the group with the CR. stop when no more CRs are found
+                    # or there is only one two diffs left (which means there is
+                    # actually one left, since the next CR will be masked after
+                    # checking that condition)
+                    while new_CR_found and (ndiffs - np.sum(np.isnan(pix_first_diffs)) > 2):
+                        new_CR_found = False
+
+                        # set CRs to nans in first diffs to clip them
+                        pix_first_diffs[~pix_cr_mask] = np.nan
+
+                        # recalculate median, sigma, and ratio
+                        new_pix_median_diffs = calc_med_first_diffs(pix_first_diffs)
+
+                        new_pix_sigma = np.sqrt(np.abs(new_pix_median_diffs) + pix_rn2 / nframes)
+                        new_pix_ratio = np.abs(pix_first_diffs - new_pix_median_diffs) / new_pix_sigma
+
+                        # check if largest ratio exceeds threshold appropriate for num remaining groups
+
+                        # select appropriate thresh. based on number of remaining groups
+                        rej_thresh = normal_rej_thresh
+                        if ndiffs - np.sum(np.isnan(pix_first_diffs)) == 3:
+                            rej_thresh = three_diff_rej_thresh
+                        if ndiffs - np.sum(np.isnan(pix_first_diffs)) == 2:
+                            rej_thresh = two_diff_rej_thresh
+                        max_idx = np.nanargmax(new_pix_ratio)
+                        location = np.unravel_index(max_idx, new_pix_ratio.shape)
+                        if new_pix_ratio[location] > rej_thresh:
+                            new_CR_found = True
+                            pix_cr_mask[location] = 0
+                        unusable_diffs = np.sum(np.isnan(pix_first_diffs))
+                    # Found all CRs for this pix - set flags in input DQ array
+                    gdq[:, 1:, all_crs_row[j], all_crs_col[j]] = np.bitwise_or(
+                         gdq[:, 1:, all_crs_row[j], all_crs_col[j]],
+                        dqflags["JUMP_DET"] * np.invert(pix_cr_mask),
                     )
-
-                if (
-                    cr_col[j] != ncols - 1
-                    and (gdq[integ, group, row, col + 1] & sat_flag) == 0
-                    and (gdq[integ, group, row, col + 1] & dnu_flag) == 0
-                ):
-                    gdq[integ, group, row, col + 1] = np.bitwise_or(
-                        gdq[integ, group, row, col + 1], jump_flag
-                    )
-
-    # flag n groups after jumps above the specified thresholds to account for
-    # the transient seen after ramp jumps
-    flag_e_threshold = [after_jump_flag_e1, after_jump_flag_e2]
-    flag_groups = [after_jump_flag_n1, after_jump_flag_n2]
-    for cthres, cgroup in zip(flag_e_threshold, flag_groups):
-        if cgroup > 0:
-            cr_intg, cr_group, cr_row, cr_col = np.where(np.bitwise_and(gdq, jump_flag))
-            for j in range(len(cr_group)):
-                intg = cr_intg[j]
-                group = cr_group[j]
-                row = cr_row[j]
-                col = cr_col[j]
-                if e_jump_4d[intg, group - 1, row, col] >= cthres:
-                    for kk in range(group, min(group + cgroup + 1, ngroups)):
-                        if (gdq[intg, kk, row, col] & sat_flag) == 0 and (
-                            gdq[intg, kk, row, col] & dnu_flag
-                        ) == 0:
-                            gdq[intg, kk, row, col] = np.bitwise_or(gdq[intg, kk, row, col], jump_flag)
-
+                    
+    num_primary_crs = np.sum(gdq & jump_flag == jump_flag)
+    
+    # Flag the four neighbors using bitwise or, shifting the reference
+    # boolean flag on pixel right, then left, then up, then down.
+    # Flag neighbors above the threshold for which neither saturation 
+    # nor donotuse is set.
+    
+    if flag_4_neighbors:
+        for i in range(nints):
+            for j in range(ngroups - 1):
+                ratio = np.abs(first_diffs[i, j] - median_diffs)/sigma
+                jump_set = gdq[i, j + 1] & jump_flag != 0
+                flag = (ratio < max_jump_to_flag_neighbors) & \
+                    (ratio > min_jump_to_flag_neighbors) & \
+                    (jump_set)
+
+                # Dilate the flag by one pixel in each direction.
+                flagsave = flag.copy()
+                flag[1:] |= flagsave[:-1]
+                flag[:-1] |= flagsave[1:]
+                flag[:, 1:] |= flagsave[:, :-1]
+                flag[:, :-1] |= flagsave[:, 1:]
+                sat_or_dnu_notset = gdq[i, j + 1] & (sat_flag | dnu_flag) == 0
+                gdq[i, j + 1][sat_or_dnu_notset & flag] |= jump_flag
+                row_below_gdq[i, j + 1][flagsave[0]] = jump_flag
+                row_above_gdq[i, j + 1][flagsave[-1]] = jump_flag
+                
+    # Flag n groups after jumps above the specified thresholds to
+    # account for the transient seen after ramp jumps.  Again, use
+    # boolean arrays; the propagation happens in a separate function.
+    
+    if after_jump_flag_n1 > 0 or after_jump_flag_n2 > 0:
+        for i in range(nints):
+            ejump = first_diffs[i] - median_diffs[np.newaxis, :]
+            jump_set = gdq[i] & jump_flag != 0
+            
+            bigjump = np.zeros(jump_set.shape, dtype=bool)
+            verybigjump = np.zeros(jump_set.shape, dtype=bool)
+
+            bigjump[1:] = (ejump >= after_jump_flag_e1) & jump_set[1:]
+            verybigjump[1:] = (ejump >= after_jump_flag_e2) & jump_set[1:]
+            
+            # Propagate flags forward
+            propagate_flags(bigjump, after_jump_flag_n1)
+            propagate_flags(verybigjump, after_jump_flag_n2)
+            
+            # Set the flags for pixels after these jumps that are not
+            # already flagged as saturated or do not use.
+            sat_or_dnu_notset = gdq[i] & (sat_flag | dnu_flag) == 0
+            addflag = (bigjump | verybigjump) & sat_or_dnu_notset
+            gdq[i][addflag] |= jump_flag
+            
     if "stddev" in locals():
         return gdq, row_below_gdq, row_above_gdq, num_primary_crs, stddev
 
@@ -521,6 +475,38 @@ def find_crs_old(
     return gdq, row_below_gdq, row_above_gdq, num_primary_crs, dummy
 
 
+def propagate_flags(boolean_flag, n_groups_flag):
+    """Propagate a boolean flag array npix groups along the first axis.
+    If the number of groups to propagate is not too large, or if a
+    high percentage of pixels are flagged, use boolean or on the
+    array.  Otherwise use np.where.  In both cases operate on the
+    array in-place.
+    Parameters
+    ----------
+    boolean_flag : 3D boolean array
+        Should be True where the flag is to be propagated.
+    n_groups_flag : int
+        Number of groups to propagate flags forward.
+    Returns
+    -------
+    None
+    """
+    ngroups = boolean_flag.shape[0]
+    jmax = min(n_groups_flag, ngroups - 2)
+    # Option A: iteratively propagate all flags forward by one
+    # group at a time.  Do this unless we have a lot of groups
+    # and cosmic rays are rare.
+    if (jmax <= 50 and jmax > 0) or np.mean(boolean_flag) > 1e-3:
+        for j in range(jmax):
+            boolean_flag[j + 1:] |= boolean_flag[j:-1]
+    # Option B: find the flags and propagate them individually.
+    elif jmax > 0:
+        igrp, icol, irow = np.where(boolean_flag)
+        for j in range(len(igrp)):
+            boolean_flag[igrp[j]:igrp[j] + n_groups_flag + 1, icol[j], irow[j]] = True
+    return
+
+
 def calc_med_first_diffs(in_first_diffs):
     """Calculate the median of `first diffs` along the group axis.
 

Test patch: diff --git a/tests/test_jump.py b/tests/test_jump.py
index c620b374..734b4075 100644
--- a/tests/test_jump.py
+++ b/tests/test_jump.py
@@ -10,8 +10,7 @@
     flag_large_events,
     point_inside_ellipse,
     find_first_good_group,
-    detect_jumps_data,
-    find_last_grp
+    detect_jumps_data
 )
 
 DQFLAGS = {
@@ -614,14 +613,3 @@ def test_calc_num_slices():
     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)