JP-3829: Add flat handling to clean_flicker_noise (#9064)

changes/9064.clean_flicker_noise.rst

@@ -0,0 +1 @@
+Add optional flat handling for the ``clean_flicker_noise`` step to avoid fitting and removing flat structure.

docs/jwst/clean_flicker_noise/arguments.rst

@@ -20,17 +20,27 @@ the behavior of the processing.
   via `~photutils.background.Background2D`.
   If None, the background value is set to 0.0.
 
-``--background_box_size`` (list of int, default=(32,32))
+``--background_box_size`` (list of int, default=None)
   Box size for the data grid used by `~photutils.background.Background2D`
   when `background_method` = 'model'. For best results, use a
-  box size that evenly divides the input image shape.
+  box size 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`` (boolean, default=False)
   For NIRSpec, mask regions of the image defined by WCS bounding
   boxes for slits/slices, as well as any regions known to be
   affected by failed-open MSA shutters.  For MIRI imaging, mask
   regions of the detector not used for science.
 
+``--apply_flat_field`` (boolean, default=False)
+  If set, images are flat-corrected prior to fitting background
+  and noise levels.  A full-frame flat field image
+  (reference type FLAT) is required. For modes that do not provide
+  FLAT files via CRDS, including all NIRSpec modes, a manually
+  generated override flat is required to enable this option.
+  Use the `override_flat` parameter to provide an alternate flat image
+  as needed (see :ref:`overriding reference files <intro_override_reference_file>`).
+
 ``--n_sigma`` (float, default=2.0)
   The sigma-clipping threshold to use when searching for outliers
   and illuminated pixels to be excluded from use in the background

docs/jwst/clean_flicker_noise/main.rst

@@ -184,6 +184,9 @@ information on all referenced parameters.
       This will mask out regions of the detector under the metering
       structure.
 
+   #. If `apply_flat_field` is set and a flat file is available, divide the
+      draft rate data by the flat image.
+
    #. Iteratively sigma clip the data to get a center value (mean or median)
       and sigma value (standard deviation).
 
@@ -200,6 +203,9 @@ information on all referenced parameters.
 
    #. Make a diff image (current group – previous group) to correct.
 
+   #. If `apply_flat_field` is set and a flat file is available, divide the
+      diff image by the flat image.
+
    #. Fit and remove a background level, using the scene mask to identify
       background pixels.
 
@@ -231,7 +237,8 @@ information on all referenced parameters.
          detector channel.
 
    #. Restore the background level to the cleaned, background-subtracted
-      diff image.
+      diff image.  Also restore the flat structure if needed by multiplying the
+      cleaned diff by the flat image.
 
    #. Add the cleaned diff back to a cleaned version of the previous
       group image.

docs/jwst/user_documentation/parameters.rst

@@ -63,6 +63,7 @@ individual steps, you have two options:
 To do this using the Python pipeline interface, see .. _python_output_file:. To do
 this when using the command line interface, see .. _cli_output_file:.
 
+.. _intro_override_reference_file:
 
 Override Reference File
 -----------------------
@@ -134,8 +135,3 @@ file from CRDS when running a pipeline:
 
 If there is need to re-use a set of parameters often, parameters can be stored
 in **parameter files**. See :ref:`parameter_files` for more information.
-
-Pipeline/Step Parameters
-========================
-
-

jwst/clean_flicker_noise/clean_flicker_noise.py

@@ -15,6 +15,7 @@
 from jwst.flatfield import FlatFieldStep
 from jwst.clean_flicker_noise.lib import NSClean, NSCleanSubarray
 from jwst.lib.basic_utils import LoggingContext
+from jwst.lib.reffile_utils import ref_matches_sci, get_subarray_model
 from jwst.msaflagopen import MSAFlagOpenStep
 from jwst.ramp_fitting import RampFitStep
 
@@ -581,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):
@@ -1015,8 +1024,63 @@ def _standardize_parameters(
     return axis_to_correct, background_method, fit_by_channel, fc
 
 
+def _read_flat_file(input_model, flat_filename):
+    """
+    Read flat data from an input file path.
+
+    Flat data is assumed to be full frame.  Subarrays matching the input
+    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 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.
+
+    Parameters
+    ----------
+    input_model : `~jwst.datamodel.JwstDataModel`
+        The input data.
+    flat_filename : str
+        File path for a full-frame flat image.
+
+    Returns
+    -------
+    flat_data : array-like of float
+        A 2D flat image array matching the input data.
+    """
+    if flat_filename is None:
+        return None
+
+    # Open the provided flat as FlatModel
+    log.debug('Dividing by flat data prior to fitting')
+    flat = datamodels.FlatModel(flat_filename)
+
+    # Extract subarray from reference data, if necessary
+    if ref_matches_sci(input_model, flat):
+        flat_data = flat.data
+    else:
+        log.debug("Extracting matching subarray from flat")
+        sub_flat = get_subarray_model(input_model, flat)
+        flat_data = sub_flat.data
+        sub_flat.close()
+    flat.close()
+
+    # 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)
+    if np.any(bad_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
+
+
 def _make_processed_rate_image(
-        input_model, single_mask, input_dir, exp_type, mask_science_regions):
+        input_model, single_mask, input_dir, exp_type, mask_science_regions, flat):
     """
     Make a draft rate image and postprocess if needed.
 
@@ -1039,6 +1103,9 @@ def _make_processed_rate_image(
         If set, science regions should be masked, so run `assign_wcs`
         and `msaflagopen` for NIRSpec and `flatfield` for MIRI
         after creating the draft rate image.
+    flat : array-like of float or None
+        If not None, the draft rate will be divided by the flat array
+        before returning. The provided flat must match the rate shape.
 
     Returns
     -------
@@ -1063,6 +1130,13 @@ def _make_processed_rate_image(
             msaflagopen=flag_open, flat_dq=flat_dq,
             input_dir=input_dir)
 
+    # Divide by the flat if provided
+    if flat is not None:
+        # Make a copy if necessary so we don't change the input rate file
+        if image_model is input_model:
+            image_model = input_model.copy()
+        image_model.data /= flat
+
     return image_model
 
 
@@ -1192,7 +1266,7 @@ def _check_data_shapes(input_model, background_mask):
 
 def _clean_one_image(image, mask, background_method, background_box_size,
                      n_sigma, fit_method, detector, fc, axis_to_correct,
-                     fit_by_channel):
+                     fit_by_channel, flat):
     """
     Clean an image by fitting and removing background noise.
 
@@ -1223,6 +1297,9 @@ def _clean_one_image(image, mask, background_method, background_box_size,
     fit_by_channel : bool
         Option to fit noise in detector channels separately,
         used for `fit_method` = 'median'
+    flat : array-like of float or None
+        If not None, the image is divided by the flat before fitting background
+        and noise.  Flat data must match the shape of the image.
 
     Returns
     -------
@@ -1240,6 +1317,10 @@ def _clean_one_image(image, mask, background_method, background_box_size,
     # Make sure data is float32
     image = image.astype(np.float32)
 
+    # If provided, divide the image by the flat
+    if flat is not None:
+        image /= flat
+
     # Find and replace/mask NaNs
     nan_pix = np.isnan(image)
     image[nan_pix] = 0.0
@@ -1284,6 +1365,10 @@ def _clean_one_image(image, mask, background_method, background_box_size,
     # Restore the background level
     cleaned_image += background
 
+    # Restore the flat structure
+    if flat is not None:
+        cleaned_image *= flat
+
     # Restore NaNs in cleaned image
     cleaned_image[nan_pix] = np.nan
 
@@ -1313,7 +1398,8 @@ def _mask_unusable(mask, dq):
 def do_correction(input_model, input_dir=None, fit_method='median',
                   fit_by_channel=False, background_method='median',
                   background_box_size=None,
-                  mask_science_regions=False, n_sigma=2.0, fit_histogram=False,
+                  mask_science_regions=False, flat_filename=None,
+                  n_sigma=2.0, fit_histogram=False,
                   single_mask=True, user_mask=None, save_mask=False,
                   save_background=False, save_noise=False):
     """
@@ -1353,6 +1439,10 @@ def do_correction(input_model, input_dir=None, fit_method='median',
         affected by failed-open MSA shutters.  For MIRI imaging, mask
         regions of the detector not used for science.
 
+    flat_filename : str, optional
+        Path to a flat field image to apply to the data before fitting
+        noise/background.
+
     n_sigma : float, optional
         N-sigma rejection level for finding outliers.
 
@@ -1417,10 +1507,13 @@ def do_correction(input_model, input_dir=None, fit_method='median',
     axis_to_correct, background_method, fit_by_channel, fc = _standardize_parameters(
             exp_type, subarray, slowaxis, background_method, fit_by_channel)
 
+    # Read the flat file, if provided
+    flat = _read_flat_file(input_model, flat_filename)
+
     # Make a rate file if needed
     if user_mask is None:
         image_model = _make_processed_rate_image(
-            input_model, single_mask, input_dir, exp_type, mask_science_regions)
+            input_model, single_mask, input_dir, exp_type, mask_science_regions, flat)
     else:
         image_model = input_model
 
@@ -1477,7 +1570,7 @@ def do_correction(input_model, input_dir=None, fit_method='median',
             # Clean the image
             cleaned_image, background, success = _clean_one_image(
                 image, mask, background_method, background_box_size, n_sigma,
-                fit_method, detector, fc, axis_to_correct, fit_by_channel)
+                fit_method, detector, fc, axis_to_correct, fit_by_channel, flat)
 
             if not success:
                 # Cleaning failed for internal reasons - probably the

jwst/clean_flicker_noise/clean_flicker_noise_step.py

@@ -30,6 +30,7 @@ class CleanFlickerNoiseStep(Step):
         background_method = option('median', 'model', None, default='median') # Background fitting algorithm
         background_box_size = int_list(min=2, max=2, default=None)  # Background box size for modeled background
         mask_science_regions = boolean(default=False)  # Mask known science regions
+        apply_flat_field = boolean(default=False)  # Apply a flat correction before fitting background and noise
         n_sigma = float(default=2.0)  # Clipping level for non-background signal
         fit_histogram = boolean(default=False)  # Fit a value histogram to derive sigma
         single_mask = boolean(default=True)  # Make a single mask for all integrations
@@ -40,6 +41,8 @@ class CleanFlickerNoiseStep(Step):
         skip = boolean(default=True)  # By default, skip the step
     """
 
+    reference_file_types = ['flat']
+
     def process(self, input):
         """
         Fit and subtract 1/f background noise from a ramp data set.
@@ -63,17 +66,28 @@ 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
             boxes for slits/slices, as well as any regions known to be
             affected by failed-open MSA shutters.  For MIRI imaging, mask
             regions of the detector not used for science.
 
+        apply_flat_field : bool, optional
+            If set, images are flat-corrected prior to fitting background
+            and noise levels.  A full-frame flat field image
+            (reference type FLAT) is required. For modes that do not provide
+            FLAT files via CRDS, including all NIRSpec spectral modes, a manually
+            generated override flat is required to enable this option.
+            Use the `override_flat` parameter to provide an alternate flat image
+            as needed.
+
         n_sigma : float, optional
             Sigma clipping threshold to be used in detecting outliers in the image.
 
@@ -108,12 +122,28 @@ def process(self, input):
         # Open the input data model
         with datamodels.open(input) as input_model:
 
+            flat_filename = None
+            if self.apply_flat_field:
+                flat_filename = self.get_reference_file(input_model, 'flat')
+                exp_type = input_model.meta.exposure.type
+                if flat_filename == 'N/A':
+                    self.log.warning(f"Flat correction is not available for "
+                                     f"exposure type {exp_type} without a user-"
+                                     f"supplied flat.")
+                    flat_filename = None
+                else:
+                    self.log.info(f'Using FLAT reference file: {flat_filename}')
+
             result = clean_flicker_noise.do_correction(
-                input_model, self.input_dir, self.fit_method, self.fit_by_channel,
-                self.background_method, self.background_box_size,
-                self.mask_science_regions, self.n_sigma, self.fit_histogram,
-                self.single_mask, self.user_mask,
-                self.save_mask, self.save_background, self.save_noise)
+                input_model, input_dir=self.input_dir, fit_method=self.fit_method,
+                fit_by_channel=self.fit_by_channel,
+                background_method=self.background_method,
+                background_box_size=self.background_box_size,
+                mask_science_regions=self.mask_science_regions,
+                flat_filename=flat_filename, n_sigma=self.n_sigma,
+                fit_histogram=self.fit_histogram, single_mask=self.single_mask,
+                user_mask=self.user_mask, save_mask=self.save_mask,
+                save_background=self.save_background, save_noise=self.save_noise)
             output_model, mask_model, background_model, noise_model, status = result
 
             # Save the mask, if requested