First version of PSF extraction, contains hacks

Co-authored-by: jemorrison <morrison@as.arizona.edu>

jwst/extract_1d/extract.py

@@ -16,13 +16,17 @@
 from jwst.lib.wcs_utils import get_wavelengths
 from jwst.extract_1d import extract1d, spec_wcs
 from jwst.extract_1d.apply_apcorr import select_apcorr
+from jwst.extract_1d.psf_profile import psf_profile
 
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
 
 WFSS_EXPTYPES = ['NIS_WFSS', 'NRC_WFSS', 'NRC_GRISM']
 """Exposure types to be regarded as wide-field slitless spectroscopy."""
 
+OPTIMAL_EXPTYPES = ['MIR_LRS-FIXEDSLIT']
+"""Exposure types for which optimal extraction is available."""
+
 ANY = "ANY"
 """Wildcard for slit name.
 
@@ -139,7 +143,8 @@ def read_apcorr_ref(refname, exptype):
 
 def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
                            smoothing_length, bkg_fit, bkg_order, use_source_posn,
-                           subtract_background):
+                           subtract_background, extraction_type, specwcs_ref_name,
+                           psf_ref_name):
     """Get extraction parameter values.
 
     Parameters
@@ -202,6 +207,17 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
     subtract_background : bool or None
         If False, all background parameters will be ignored.
 
+    extraction_type : str
+        Extraction type ('box' or 'optimal').  Optimal extraction is
+        only available if `specwcs_ref_name` and `psf_ref_name` are
+        not 'N/A'.
+
+    specwcs_ref_name : str
+        The name of the specwcs reference file, or "N/A".
+
+    psf_ref_name : str
+        The name of the PSF reference file, or "N/A".
+
     Returns
     -------
     extract_params : dict
@@ -228,6 +244,8 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
         extract_params['bkg_order'] = 0  # because no background sub.
         extract_params['subtract_background'] = False
         extract_params['extraction_type'] = 'box'
+        extract_params['specwcs_ref_name'] = 'N/A'
+        extract_params['psf_ref_name'] = 'N/A'
 
         if use_source_posn is None:
             extract_params['use_source_posn'] = False
@@ -326,9 +344,10 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
                         # If the user supplied a value, use that value.
                         extract_params['smoothing_length'] = smoothing_length
 
-                    # Default the extraction type to 'box': 'optimal'
-                    # is not yet supported.
-                    extract_params['extraction_type'] = 'box'
+                    # Set the extraction type to 'box' or 'optimal'
+                    extract_params['extraction_type'] = extraction_type
+                    extract_params['specwcs'] = specwcs_ref_name
+                    extract_params['psf'] = psf_ref_name
 
                     break
 
@@ -1206,10 +1225,27 @@ def define_aperture(input_model, slit, extract_params, exp_type):
 
             # Offet extract parameters by location - nominal
             shift_by_source_location(location, nominal_location, extract_params)
+    else:
+        middle_pix, middle_wl, location = None, None, None
 
     # Make a spatial profile, including source shifts if necessary
-    profile, lower_limit, upper_limit = box_profile(
-        data_shape, extract_params, wl_array, return_limits=True)
+    if extract_params['extraction_type'] == 'optimal':
+        if location is None:
+            nominal_profile = box_profile(data_shape, extract_params, wl_array,
+                                          label='nominal aperture')
+            middle_pix, location = aperture_center(nominal_profile)
+            if extract_params['dispaxis'] == HORIZONTAL:
+                middle_wl = wl_array[middle_pix, location]
+            else:
+                middle_wl = wl_array[location, middle_pix]
+
+        profile, lower_limit, upper_limit = psf_profile(
+            data_model, extract_params['psf'], extract_params['specwcs'],
+            middle_wl, location)
+
+    else:
+        profile, lower_limit, upper_limit = box_profile(
+            data_shape, extract_params, wl_array, return_limits=True)
 
     # Make sure profile weights are zero where wavelengths are invalid
     profile[~np.isfinite(wl_array)] = 0.0
@@ -1383,8 +1419,9 @@ def extract_one_slit(data_model, integ, profile, bg_profile, extract_params):
 def create_extraction(input_model, slit, output_model,
                       extract_ref_dict, slitname, sp_order, smoothing_length,
                       bkg_fit, bkg_order, use_source_posn, exp_type,
-                      subtract_background, apcorr_ref_model, log_increment,
-                      save_profile):
+                      subtract_background, apcorr_ref_model,
+                      extraction_type, specwcs_ref_name, psf_ref_name,
+                      log_increment, save_profile):
     """Extract spectra from an input model and append to an output model.
 
     Input data, specified in the `slit` or `input_model`, should contain data
@@ -1463,6 +1500,14 @@ def create_extraction(input_model, slit, output_model,
     apcorr_ref_model : DataModel or None
         The aperture correction reference datamodel, containing the
         APCORR reference file data.
+    extraction_type : str
+        Extraction type ('box' or 'optimal').  Optimal extraction is
+        only available if `specwcs_ref_name` and `psf_ref_name` are
+        not 'N/A'.
+    specwcs_ref_name : str
+        The name of the specwcs reference file, or "N/A".
+    psf_ref_name : str
+        The name of the PSF reference file, or "N/A".
     log_increment : int
         If greater than 0 and the input data are multi-integration, a message
         will be written to the log every `log_increment` integrations.
@@ -1543,7 +1588,7 @@ def create_extraction(input_model, slit, output_model,
     extract_params = get_extract_parameters(
         extract_ref_dict, data_model, slitname, sp_order, input_model.meta,
         smoothing_length, bkg_fit,bkg_order, use_source_posn,
-        subtract_background
+        subtract_background, extraction_type, specwcs_ref_name, psf_ref_name
     )
 
     if extract_params['match'] == NO_MATCH:
@@ -1784,9 +1829,10 @@ def create_extraction(input_model, slit, output_model,
     return profile_model
 
 
-def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_length,
-                  bkg_fit, bkg_order, log_increment, subtract_background,
-                  use_source_posn, save_profile):
+def run_extract1d(input_model, extract_ref_name, apcorr_ref_name,
+                  specwcs_ref_name, psf_ref_name, extraction_type,
+                  smoothing_length, bkg_fit, bkg_order, log_increment,
+                  subtract_background, use_source_posn, save_profile):
     """Extract all 1-D spectra from an input model.
 
     Parameters
@@ -1798,7 +1844,18 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
         The name of the extract1d reference file, or "N/A".
 
     apcorr_ref_name : str
-        Name of the APCORR reference file. Default is None
+        Name of the APCORR reference file. Default is None.
+
+    specwcs_ref_name : str
+        The name of the specwcs reference file, or "N/A".
+
+    psf_ref_name : str
+        The name of the PSF reference file, or "N/A".
+
+    extraction_type : str
+        Extraction type ('box' or 'optimal').  Optimal extraction is
+        only available if `specwcs_ref_name` and `psf_ref_name` are
+        not 'N/A'.
 
     smoothing_length : int or None
         Width of a boxcar function for smoothing the background regions.
@@ -1863,6 +1920,14 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
     if apcorr_ref_name is not None and apcorr_ref_name != 'N/A':
         apcorr_ref_model = read_apcorr_ref(apcorr_ref_name, exp_type)
 
+    # Check for non-null specwcs and PSF reference files
+    if (exp_type not in OPTIMAL_EXPTYPES
+            or specwcs_ref_name == 'N/A' or psf_ref_name == 'N/A'):
+        if extraction_type != 'box':
+            log.warning(f'Optimal extraction is not available for EXP_TYPE {exp_type}')
+            log.warning('Defaulting to box extraction.')
+            extraction_type = 'box'
+
     # Set up the output model
     output_model = datamodels.MultiSpecModel()
     if hasattr(meta_source, "int_times"):
@@ -1944,8 +2009,9 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
                         input_model, slit, output_model,
                         extract_ref_dict, slitname, sp_order, smoothing_length,
                         bkg_fit, bkg_order, use_source_posn, exp_type,
-                        subtract_background, apcorr_ref_model, log_increment,
-                        save_profile)
+                        subtract_background, apcorr_ref_model,
+                        extraction_type, specwcs_ref_name, psf_ref_name,
+                        log_increment, save_profile)
                 except ContinueError:
                     pass
 
@@ -1979,8 +2045,9 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
                         input_model, slit, output_model,
                         extract_ref_dict, slitname, sp_order, smoothing_length,
                         bkg_fit, bkg_order, use_source_posn, exp_type,
-                        subtract_background, apcorr_ref_model, log_increment,
-                        save_profile)
+                        subtract_background, apcorr_ref_model,
+                        extraction_type, specwcs_ref_name, psf_ref_name,
+                        log_increment, save_profile)
                 except ContinueError:
                     pass
 

jwst/extract_1d/extract_1d_step.py

@@ -15,6 +15,12 @@ class Extract1dStep(Step):
 
     Attributes
     ----------
+    extraction_type : str
+        If 'box', a standard extraction is performed, summing over an
+        aperture box. If 'optimal', a PSF-based extraction is performed.
+        Currently, optimal extraction is only available for MIRI LRS Fixed
+        Slit data.
+
     use_source_posn : bool or None
         If True, the source and background extraction positions specified in
         the extract1d reference file (or the default position, if there is no
@@ -154,6 +160,7 @@ class Extract1dStep(Step):
     class_alias = "extract_1d"
 
     spec = """
+    extraction_type = option("box", "optimal", default="box") # Perform box or optimal extraction
     use_source_posn = boolean(default=None)  # use source coords to center extractions?
     apply_apcorr = boolean(default=True)  # apply aperture corrections?
     log_increment = integer(default=50)  # increment for multi-integration log messages
@@ -186,7 +193,8 @@ class Extract1dStep(Step):
     soss_modelname = output_file(default = None)  # Filename for optional model output of traces and pixel weights
     """
 
-    reference_file_types = ['extract1d', 'apcorr', 'pastasoss', 'specprofile', 'speckernel']
+    reference_file_types = ['extract1d', 'apcorr', 'pastasoss', 'specprofile',
+                            'speckernel', 'specwcs'] #, 'psf']
 
     def _get_extract_reference_files_by_mode(self, model, exp_type):
         """Get extraction reference files with special handling by exposure type."""
@@ -207,7 +215,15 @@ def _get_extract_reference_files_by_mode(self, model, exp_type):
         if apcorr_ref != 'N/A':
             self.log.info(f'Using APCORR file {apcorr_ref}')
 
-        return extract_ref, apcorr_ref
+        if exp_type in extract.OPTIMAL_EXPTYPES:
+            specwcs_ref = self.get_reference_file(model, 'specwcs')
+            #psf_ref = self.get_reference_file(model, 'psf')
+            psf_ref = "MIRI_LRS_SLIT_EPSF_20240602_WAVE_1D.fits"
+        else:
+            specwcs_ref = 'N/A'
+            psf_ref = 'N/A'
+
+        return extract_ref, apcorr_ref, specwcs_ref, psf_ref
 
     def _extract_soss(self, model):
         """Extract NIRISS SOSS spectra."""
@@ -381,8 +397,8 @@ def process(self, input):
             result = ModelContainer()
             for model in input_model:
                 # Get the reference file names
-                extract_ref, apcorr_ref = self._get_extract_reference_files_by_mode(
-                    model, exp_type)
+                extract_ref, apcorr_ref, specwcs_ref, psf_ref = \
+                    self._get_extract_reference_files_by_mode(model, exp_type)
 
                 profile = None
                 if isinstance(model, datamodels.IFUCubeModel):
@@ -394,6 +410,9 @@ def process(self, input):
                         model,
                         extract_ref,
                         apcorr_ref,
+                        specwcs_ref,
+                        psf_ref,
+                        self.extraction_type,
                         self.smoothing_length,
                         self.bkg_fit,
                         self.bkg_order,