JP-2075: fixed flux scaling issues in wfss_contam (#8416)

CHANGES.rst

@@ -121,6 +121,11 @@ tweakreg
 - Changed default settings for ``abs_separation`` parameter for the ``tweakreg``
   step to have a value compatible with the ``abs_tolerance`` parameter. [#8445]
 
+wfss_contam
+-----------
+
+- Fixed flux scaling issue in model contamination image by adding background
+  subtraction and re-scaling fluxes to respect wavelength oversampling. [#8416]
 
 1.14.0 (2024-03-29)
 ===================

jwst/wfss_contam/disperse.py

@@ -1,4 +1,5 @@
 import numpy as np
+import warnings
 
 from scipy.interpolate import interp1d
 
@@ -115,9 +116,10 @@ def flux(x):
 
     # Use a natural wavelength scale or the wavelength scale of the input SED/spectrum,
     # whichever is smaller, divided by oversampling requested
-    input_dlam = np.median(lams[1:] - lams[:-1])
-    if input_dlam < dw:
-        dlam = input_dlam / oversample_factor
+    if len(lams) > 1:
+        input_dlam = np.median(lams[1:] - lams[:-1])
+        if input_dlam < dw:
+            dlam = input_dlam / oversample_factor
     else:
         # this value gets used when we only have 1 direct image wavelength
         dlam = dw / oversample_factor
@@ -161,7 +163,9 @@ def flux(x):
     # values are naturally in units of physical fluxes, so we divide out
     # the sensitivity (flux calibration) values to convert to units of
     # countrate (DN/s).
-    counts = flux(lams) * areas / sens
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=RuntimeWarning, message="divide by zero")
+        counts = flux(lams) * areas / (sens * oversample_factor)
     counts[no_cal] = 0.  # set to zero where no flux cal info available
 
     return xs, ys, areas, lams, counts, ID

jwst/wfss_contam/observations.py

@@ -10,10 +10,53 @@
 
 import logging
 
+from photutils.background import Background2D, MedianBackground
+from astropy.stats import SigmaClip
+
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
 
 
+def background_subtract(data, box_size=None, filter_size=(3,3), sigma=3.0, exclude_percentile=30.0):
+    """
+    Simple astropy background subtraction
+
+    Parameters
+    ----------
+    data : np.ndarray
+        2D array of pixel values
+    box_size : tuple
+        Size of box in pixels to use for background estimation. 
+        If not set, defaults to 1/5 of the image size.
+    filter_size : tuple
+        Size of filter to use for background estimation
+    sigma : float
+        Sigma threshold for background clipping
+    exclude_percentile : float
+        Percentage of masked pixels above which box is excluded from background estimation
+
+    Returns
+    -------
+    data : np.ndarray
+        2D array of pixel values with background subtracted
+
+    Notes
+    -----
+    Improper background subtraction in input _i2d image leads to extra flux
+    in the simulated dispersed image, and was one cause of flux scaling issues
+    in a previous version.
+    """
+    if box_size is None:
+        box_size = (int(data.shape[0]/5), int(data.shape[1]/5))
+    sigma_clip = SigmaClip(sigma=sigma)
+    bkg_estimator = MedianBackground()
+    bkg = Background2D(data, box_size, filter_size=filter_size,
+                   sigma_clip=sigma_clip, bkg_estimator=bkg_estimator, 
+                   exclude_percentile=exclude_percentile)
+
+    return data - bkg.background
+
+
 class Observation:
     """This class defines an actual observation. It is tied to a single grism image."""
 
@@ -120,6 +163,7 @@ def create_pixel_list(self):
             log.info(f"Using direct image {dir_image_name}")
             with datamodels.open(dir_image_name) as model:
                 dimage = model.data
+                dimage = background_subtract(dimage)
 
                 if self.sed_file is None:
                     # Default pipeline will use sed_file=None, so we need to compute