Add option to smooth data in level_match_step

This PR allows the user to pass a variable "smooth_fwhm" to level_match, which will smooth out the data before doing the level matching. This can either be passed as a value, which will be interpreted as pixels, or as "xarcsec", which will use an angular value

To avoid edge effects in the convolution, a mask is applied at the edges if this is used

Author: thomaswilliamsastro

CHANGES.rst

@@ -1,6 +1,7 @@
 1.3.1 (Unreleased)
 ==================
 
+- Add option to smooth data in ``level_match_step``
 - Fix backgrounds being wrongly moved in ``lv1_step``
 - Moved ``astrometric_align`` to a combined step, to avoid list ordering issues
 - Updated dependencies, added specific pins, dependabot, and CODEOWNERS

pjpipe/level_match/level_match_step.py

@@ -333,6 +333,7 @@ def __init__(
             recombine_lyot=False,
             combine_nircam_short=False,
             do_local_subtraction=True,
+            smooth_fwhm=None,
             sigma=3,
             npixels=3,
             dilate_size=7,
@@ -387,6 +388,11 @@ def __init__(
                 chips before matching in a mosaic. Defaults to False
             do_local_subtraction: Whether to do a sigma-clipped local median
                 subtraction. Defaults to True
+            smooth_fwhm: By setting this, you can smooth the data with a Gaussian
+                kernel with FWHM of this value before doing
+                the level matching. Can be specified either as a raw number of
+                pixels (just pass a number) or in arcsec (pass [int]arcsec).
+                Defaults to None, which will not smooth
             sigma: Sigma for sigma-clipping. Defaults to 3
             npixels: Pixels to grow for masking. Defaults to 5
             dilate_size: make_source_mask dilation size. Defaults to 7
@@ -424,6 +430,18 @@ def __init__(
             log.warning("Cannot do local subtraction for methods beyond simple offset. Switching off")
             do_local_subtraction = False
 
+        # If we've passed the FWHM as a string, parse that now
+        if smooth_fwhm is not None:
+            if isinstance(smooth_fwhm, str):
+
+                # Case 1: we have arcsec in there, so convert to astropy units
+                if "arcsec" in smooth_fwhm:
+                    smooth_fwhm = float(smooth_fwhm.replace("arcsec", "")) * u.arcsec
+
+                # Otherwise, just parse as pixels
+                else:
+                    smooth_fwhm = float(smooth_fwhm)
+
         self.in_dir = in_dir
         self.out_dir = out_dir
         self.step_ext = step_ext
@@ -436,6 +454,7 @@ def __init__(
         self.recombine_lyot = recombine_lyot
         self.combine_nircam_short = combine_nircam_short
         self.do_local_subtraction = do_local_subtraction
+        self.smooth_fwhm = smooth_fwhm
         self.sigma = sigma
         self.npixels = npixels
         self.dilate_size = dilate_size
@@ -2296,6 +2315,7 @@ def get_reproject(
                     do_sigma_clip=self.do_sigma_clip,
                     stacked_image=stacked_image,
                     reproject_func=self.reproject_func,
+                    smooth_fwhm=self.smooth_fwhm,
                 ),
                     "err": reproject_image(
                         i,
@@ -2305,6 +2325,7 @@ def get_reproject(
                         do_sigma_clip=self.do_sigma_clip,
                         stacked_image=stacked_image,
                         reproject_func=self.reproject_func,
+                        smooth_fwhm=self.smooth_fwhm,
                     ),
                 }
                 for i in file
@@ -2318,6 +2339,7 @@ def get_reproject(
                     do_sigma_clip=self.do_sigma_clip,
                     stacked_image=stacked_image,
                     reproject_func=self.reproject_func,
+                    smooth_fwhm=self.smooth_fwhm,
                 ),
                 "err": reproject_image(
                     file,
@@ -2327,6 +2349,7 @@ def get_reproject(
                     do_sigma_clip=self.do_sigma_clip,
                     stacked_image=stacked_image,
                     reproject_func=self.reproject_func,
+                    smooth_fwhm=self.smooth_fwhm,
                 ),
             }
 

pjpipe/utils/utils.py

@@ -6,8 +6,9 @@
 import os
 import warnings
 
+import astropy.units as u
 import numpy as np
-from astropy.convolution import convolve_fft
+from astropy.convolution import convolve_fft, Gaussian2DKernel
 from astropy.io import fits
 from astropy.nddata.bitmask import interpret_bit_flags, bitfield_to_boolean_mask
 from astropy.stats import sigma_clipped_stats, SigmaClip
@@ -19,6 +20,7 @@
 from reproject.mosaicking import find_optimal_celestial_wcs, reproject_and_coadd
 from reproject.mosaicking.subset_array import ReprojectedArraySubset
 from scipy.interpolate import RegularGridInterpolator
+from scipy.ndimage import binary_dilation
 from stdatamodels import util
 from stdatamodels.jwst import datamodels
 from stdatamodels.jwst.datamodels.dqflags import pixel
@@ -695,6 +697,7 @@ def reproject_image(
         optimal_wcs,
         optimal_shape,
         hdu_type="data",
+        smooth_fwhm=None,
         do_sigma_clip=False,
         stacked_image=False,
         do_level_data=False,
@@ -707,6 +710,9 @@ def reproject_image(
         optimal_wcs: Optimal WCS for input image stack
         optimal_shape: Optimal shape for input image stack
         hdu_type: Type of HDU. Can either be 'data', 'err', or 'var_rnoise'
+        smooth_fwhm: FWHM for the Gaussian kernel to smooth the data. Can
+            be specified either in astropy units or just a number, which will
+            be pixels
         do_sigma_clip: Whether to perform sigma-clipping or not.
             Defaults to False
         stacked_image: Stacked image or not? Defaults to False
@@ -834,6 +840,69 @@ def reproject_image(
         )
         data_reproj_small[dq_reproj_small == 1] = np.nan
 
+    # If we're smoothing, do that here
+    if smooth_fwhm is not None:
+
+        # Build the smoothing kernel. Need a FWHM in pixels
+        if isinstance(smooth_fwhm, u.Quantity):
+            pix_scale = w_in.proj_plane_pixel_scales()[0].to(u.arcsec)
+            smooth_fwhm_pix = smooth_fwhm / pix_scale
+        else:
+            smooth_fwhm_pix = smooth_fwhm
+
+        smooth_sig_pix = smooth_fwhm_pix / 2.355
+
+        # Make the kernel, remember we're going FWHM->sigma
+        k = Gaussian2DKernel(smooth_sig_pix)
+
+        # Because we need to square this sometimes, just get the array out
+        # and normalise
+        k = k.array
+
+        # For error and variance, we use the square of the kernel
+        if hdu_type in ["err", "var_rnoise"]:
+            k = k ** 2
+
+        k /= np.sum(k)
+
+        # We also want to build a dilation kernel for later, and keep track
+        # of non-valid values
+        dilate_struct = np.ones([int(smooth_sig_pix), int(smooth_sig_pix)])
+        valid_value_mask = np.logical_or(data_reproj_small == 0, ~np.isfinite(data_reproj_small))
+
+        # Do the convolution. This is different depending on the HDU type
+
+        # For data and variance, we can just do the straight convolution
+        if hdu_type in ["data", "var_rnoise"]:
+            data_reproj_small = convolve_fft(data_reproj_small,
+                                             kernel=k,
+                                             allow_huge=True,
+                                             preserve_nan=True,
+                                             fill_value=np.nan,
+                                             )
+
+        # For the error, we need the square of the error (and sqrt after)
+        elif hdu_type == "err":
+            data_reproj_small = np.sqrt(
+                convolve_fft(data_reproj_small ** 2,
+                             kernel=k,
+                             allow_huge=True,
+                             preserve_nan=True,
+                             fill_value=np.nan,
+                             )
+            )
+
+        else:
+            raise Warning(f"Unsure how to deal with hdu_type {hdu_type}")
+
+        # Finally, dilate the valid value mask by the sigma of the smoothing kernel and
+        # apply that as a mask
+        valid_value_dilate = binary_dilation(valid_value_mask,
+                                           structure=dilate_struct,
+                                           )
+
+        data_reproj_small[valid_value_dilate == 1] = np.nan
+
     # If we're sigma-clipping, reproject the mask. This needs to use
     # reproject_interp, so we can keep whole numbers
     if do_sigma_clip:
@@ -1060,11 +1129,11 @@ def level_data(
 
     for i in range(3):
         quad_1 = data[:, i * quadrant_size: (i + 1) * quadrant_size][
-                 :, quadrant_size - 20:
-                 ]
+            :, quadrant_size - 20:
+        ]
         dq_1 = dq_mask[:, i * quadrant_size: (i + 1) * quadrant_size][
-               :, quadrant_size - 20:
-               ]
+            :, quadrant_size - 20:
+        ]
         quad_2 = data[:, (i + 1) * quadrant_size: (i + 2) * quadrant_size][:, :20]
         dq_2 = dq_mask[:, (i + 1) * quadrant_size: (i + 2) * quadrant_size][:, :20]