Better local estimate for filling in missing flat data

jwst/clean_flicker_noise/clean_flicker_noise.py

@@ -582,7 +582,15 @@ def background_level(image, mask, background_method='median',
             bkg_estimator = MedianBackground()
 
             if background_box_size is None:
-                background_box_size = [32, 32]
+                # use 32 x 32 if possible, otherwise take next largest box
+                # size that evenly divides the image (minimum 1)
+                background_box_size = []
+                recommended = np.arange(1, 33)
+                for i_size in image.shape:
+                    divides_evenly = (i_size % recommended == 0)
+                    background_box_size.append(recommended[divides_evenly][-1])
+                log.debug(f'Using box size {background_box_size}')
+
             box_division_remainder = (image.shape[0] % background_box_size[0],
                                       image.shape[1] % background_box_size[1])
             if not np.allclose(box_division_remainder, 0):
@@ -1024,7 +1032,8 @@ def _read_flat_file(input_model, flat_filename):
     data are extracted as needed.
 
     Only the flat image is returned: error and DQ arrays are ignored.
-    Any zeros or NaNs in the flat image are set to a median value before
+    Any zeros or NaNs in the flat image are set to a smoothed local average
+    value (via `background_level`, with background_method = 'model') before
     returning, to avoid impacting the background and noise fits near
     missing flat data.
 
@@ -1057,9 +1066,14 @@ def _read_flat_file(input_model, flat_filename):
         sub_flat.close()
     flat.close()
 
-    # Set any zeros or non-finite values in the flat data to a median value
-    median_flat = np.nanmedian(flat_data)
-    flat_data[(flat_data == 0) | ~np.isfinite(flat_data)] = median_flat
+    # Set any zeros or non-finite values in the flat data to a smoothed local value
+    bad_data = (flat_data == 0) | ~np.isfinite(flat_data)
+    smoothed_flat = background_level(flat_data, ~bad_data, background_method='model')
+    try:
+        flat_data[bad_data] = smoothed_flat[bad_data]
+    except IndexError:
+        # 2D model failed, median value returned instead
+        flat_data[bad_data] = smoothed_flat
 
     return flat_data
 

jwst/clean_flicker_noise/clean_flicker_noise_step.py

@@ -66,10 +66,12 @@ def process(self, input):
             `~photutils.background.Background2D`.  If None, the background
             value is 0.0.
 
-        background_box_size : tuple of int, optional
+        background_box_size : tuple of int or None, optional
             Box size for the data grid used by `Background2D` when
             `background_method` = 'model'. For best results, use a box size
-            that evenly divides the input image shape.
+            that evenly divides the input image shape.  If None, the largest
+            value between 1 and 32 that evenly divides the image dimension
+            is used.
 
         mask_science_regions : bool, optional
             For NIRSpec, mask regions of the image defined by WCS bounding

jwst/clean_flicker_noise/tests/test_clean_flicker_noise.py

@@ -416,6 +416,14 @@ def test_background_level(log_watcher):
     # model method with mismatched box size:
     # warns, but completes successfully
     log_watcher.message = "does not divide evenly"
+    background = cfn.background_level(
+        image, mask, background_method='model', background_box_size=(32, 32))
+    assert background.shape == shape
+    assert np.all(background == 1.0)
+    log_watcher.assert_seen()
+
+    # model method with None box size: picks the largest even divisor < 32
+    log_watcher.message = "box size [25, 25]"
     background = cfn.background_level(
         image, mask, background_method='model', background_box_size=None)
     assert background.shape == shape