fixed unit tests, improved output filename handling

jwst/badpix_selfcal/badpix_selfcal.py

@@ -2,6 +2,7 @@
 import numpy as np
 import jwst.datamodels as dm
 from jwst.outlier_detection.outlier_detection_ifu import medfilt
+from stdatamodels.jwst.datamodels.dqflags import pixel
 
 
 def badpix_selfcal(medbg: np.ndarray, 
@@ -62,6 +63,6 @@ def apply_flags(input_model: dm.IFUImageModel, flagged_indices: np.ndarray) -> d
 
     input_model.data[flagged_indices] = np.nan
     input_model.err[flagged_indices] = np.nan
-    input_model.dq[flagged_indices] = 1
+    input_model.dq[flagged_indices] = pixel["WARM"]
 
     return input_model

jwst/badpix_selfcal/badpix_selfcal_step.py

@@ -23,6 +23,7 @@ class BadpixSelfcalStep(Step):
     spec = """
     flagfrac = float(default=0.001)  #fraction of pixels to flag on each of low and high end
     kernel_size = integer(default=15)  #size of kernel for median filter
+    save_flagged_bkgd = boolean(default=False)  #save the flagged background exposures to fits
     skip = boolean(default=True)
     """
 
@@ -35,26 +36,33 @@ def process(self, input, bkg_list=None):
         input: JWST data model or association
             input science data to be corrected
 
-        bkg_list: list of background ImageModels.
-            When calwebb_spec2 is run, these are read from the association file
-            inside calwebb_spec2.py. When run as a standalone step, these are either
-            read directly from an input association file, or the user can provide
-            this list manually
+        bkg_list: list of background ImageModels or filenames.
 
         Returns
         -------
         output: JWST data model or association
             data model with CRs flagged
+
+        Notes
+        -----
+        If bkg_list is specified manually, it overrides the background exposures
+        in the assocation file.
+        If bkg_list is set to None and an association file is read in, the background exposures
+        are read from the association file.
+        If bkg_list is set to None and input is a single science exposure, the input exposure
+        will be used as the background exposure, i.e., true self-calibration of
+        bad pixels.
         """
         with dm.open(input) as input_data:
 
             if isinstance(input_data, dm.ModelContainer):
 
                 sci_models, bkg_list_asn = split_container(input_data)
                 if bkg_list is None:
-                    bkg_list = bkg_list_asn
+                    #keeping as ModelContainer causes problems figuring out names to save if save_results=True
+                    bkg_list = list(bkg_list_asn) 
                 else:
-                    log.warning("bkg_list provided as input, ignoring bkg_list from association file")
+                    log.warning("bkg_list provided directly as input, ignoring bkg_list from association file")
 
                 # in calwebb_spec2 there should be only a single science exposure in an association
                 input_sci = sci_models[0]
@@ -65,24 +73,33 @@ def process(self, input, bkg_list=None):
                 if bkg_list is None:
                     # true self-calibration on input data itself
                     bkg_list = [input_data]
+                else:
+                    bkg_list = [dm.open(bkg) for bkg in bkg_list]
 
             else:
                 raise TypeError("Input data is not a ModelContainer, ImageModel, or IFUImageModel.\
                                 Cannot continue.")
 
-            # collapse background dithers into a single background model
-            bkgd_3d = []
+        # collapse background dithers into a single background model
+        bkgd_3d = []
+        for i, background_model in enumerate(bkg_list):
+            bkgd_3d.append(background_model.data)
+        bkgd = np.nanmin(np.asarray(bkgd_3d), axis=0)
+        bad_indices = badpix_selfcal.badpix_selfcal(bkgd, self.flagfrac, self.kernel_size)
+
+        # apply the flags to the science data
+        input_sci = badpix_selfcal.apply_flags(input_sci, bad_indices)
+
+        # apply the flags to the background data
+        for i, background_model in enumerate(bkg_list):
+            bkg_list[i] = badpix_selfcal.apply_flags(background_model, bad_indices)
+
+        # save the flagged background exposures if desired
+        if self.save_flagged_bkgd:
             for i, background_model in enumerate(bkg_list):
-                bkgd_3d.append(background_model.data)
-            bkgd = np.nanmin(np.asarray(bkgd_3d), axis=0)
-            bad_indices = badpix_selfcal.badpix_selfcal(bkgd, self.flagfrac, self.kernel_size)
-
-            # apply the flags to the science data
-            input_sci = badpix_selfcal.apply_flags(input_sci, bad_indices)
+                stem = background_model.meta.filename.strip(".fits")
+                background_model.save(f"{stem}_{self.suffix}.fits")
 
-            # apply the flags to the background data
-            for i, background_model in enumerate(bkg_list):
-                bkg_list[i] = badpix_selfcal.apply_flags(background_model, bad_indices)
-
-            self.record_step_status(input_sci, "badpix_selfcal", success=True)
+        self.record_step_status(input_sci, "badpix_selfcal", success=True)
+
         return input_sci, bkg_list

jwst/lib/suffix.py

@@ -195,7 +195,8 @@
  'image2pipeline',
  'klip',
  'emicorr',
- 'emicorrstep'
+ 'emicorrstep',
+ 'badpixselfcalstep',
 }
 
 

jwst/stpipe/integration.py

@@ -37,6 +37,7 @@ def get_steps():
         ("jwst.step.AssignMTWcsStep", 'assign_mtwcs', False),
         ("jwst.step.AssignWcsStep", 'assign_wcs', False),
         ("jwst.step.BackgroundStep", 'background', False),
+        ("jwst.step.BadpixSelfcalStep", 'badpix_selfcal', False),
         ("jwst.step.BarShadowStep", 'barshadow', False),
         ("jwst.step.Combine1dStep", 'combine_1d', False),
         ("jwst.step.StackRefsStep", 'stack_refs', False),