From 7f5a29766b6ef1b626259a3165e8f4178de2c27b Mon Sep 17 00:00:00 2001 From: Ned Molter Date: Wed, 22 May 2024 13:47:55 -0400 Subject: [PATCH] fixed unit tests, improved output filename handling --- jwst/badpix_selfcal/badpix_selfcal.py | 3 +- jwst/badpix_selfcal/badpix_selfcal_step.py | 57 ++++++++++++++-------- jwst/lib/suffix.py | 3 +- jwst/stpipe/integration.py | 1 + 4 files changed, 42 insertions(+), 22 deletions(-) diff --git a/jwst/badpix_selfcal/badpix_selfcal.py b/jwst/badpix_selfcal/badpix_selfcal.py index 277677a712c..925067a1df8 100644 --- a/jwst/badpix_selfcal/badpix_selfcal.py +++ b/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 diff --git a/jwst/badpix_selfcal/badpix_selfcal_step.py b/jwst/badpix_selfcal/badpix_selfcal_step.py index bb2cedca410..e1eb014eccd 100644 --- a/jwst/badpix_selfcal/badpix_selfcal_step.py +++ b/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,16 +36,22 @@ 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: @@ -52,9 +59,10 @@ def process(self, input, bkg_list=None): 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 diff --git a/jwst/lib/suffix.py b/jwst/lib/suffix.py index cab5be43a94..3fff82ba9fb 100644 --- a/jwst/lib/suffix.py +++ b/jwst/lib/suffix.py @@ -195,7 +195,8 @@ 'image2pipeline', 'klip', 'emicorr', - 'emicorrstep' + 'emicorrstep', + 'badpixselfcalstep', } diff --git a/jwst/stpipe/integration.py b/jwst/stpipe/integration.py index b43f4abff54..597a6afe3e8 100644 --- a/jwst/stpipe/integration.py +++ b/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),