Add top level parameter for optimizing psf location

jwst/extract_1d/extract.py

@@ -144,9 +144,9 @@ def read_apcorr_ref(refname, exptype):
 
 def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
                            smoothing_length=None, bkg_fit=None, bkg_order=None,
-                           use_source_posn=None, subtract_background=None,
-                           extraction_type='box', specwcs_ref_name='N/A',
-                           psf_ref_name='N/A'):
+                           use_source_posn=None, optimize_psf_location=False,
+                           subtract_background=None, extraction_type='box',
+                           specwcs_ref_name='N/A', psf_ref_name='N/A'):
     """Get extraction parameter values.
 
     Parameters
@@ -155,21 +155,16 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
         For an extract1d reference file in JSON format, `ref_dict` will be
         the entire contents of the file.  If there is no extract1d reference
         file, `ref_dict` will be None.
-
     input_model : JWSTDataModel
         This can be either the input science file or one SlitModel out of
         a list of slits.
-
     slitname : str
         The name of the slit, or "ANY".
-
     sp_order : int
         The spectral order number.
-
     meta : ObjectNode
         The metadata for the actual input model, i.e. not just for the
         current slit, from input_model.meta.
-
     smoothing_length : int or None, optional
         Width of a boxcar function for smoothing the background regions.
         If None, the smoothing length will be retrieved from `ref_dict`, or
@@ -179,14 +174,12 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
         explicitly specified the value, so that value will be used.
         This argument is only used if background regions have been
         specified.
-
     bkg_fit : str or None, optional
         The type of fit to apply to background values in each
         column (or row, if the dispersion is vertical). The default
         `poly` results in a polynomial fit of order `bkg_order`. Other
         options are `mean` and `median`. If `mean` or `median` is selected,
         `bkg_order` is ignored.
-
     bkg_order : int or None, optional
         Polynomial order for fitting to each column (or row, if the
         dispersion is vertical) of background.  If None, the polynomial
@@ -198,25 +191,24 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
         specified the value, so that value will be used.
         This argument must be positive or zero, and it is only used if
         background regions have been specified.
-
     use_source_posn : bool or None, optional
         If True, the target and background positions specified in `ref_dict`
         (or a default target position) will be shifted to account for
         the actual source location in the data.
         If None, the value specified in `ref_dict` will be used, or it will
         be set to True if not found in `ref_dict`.
-
+    optimize_psf_location : bool
+        If True, and if `extraction_type` is 'optimal', then the source
+        location will be optimized, via iterative comparisons of the scene
+        model with the input data.
     subtract_background : bool or None, optional
         If False, all background parameters will be ignored.
-
     extraction_type : str, optional
         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, optional
         The name of the specwcs reference file, or "N/A".
-
     psf_ref_name : str, optional
         The name of the PSF reference file, or "N/A".
 
@@ -247,6 +239,7 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
         extract_params['subtract_background'] = False
         extract_params['extraction_type'] = 'box'
         extract_params['use_source_posn'] = False  # no source position correction
+        extract_params['optimize_psf_location'] = False
         extract_params['specwcs'] = 'N/A'
         extract_params['psf'] = 'N/A'
         extract_params['position_correction'] = 0
@@ -369,6 +362,7 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
                     extract_params['extraction_type'] = extraction_type
                     extract_params['specwcs'] = specwcs_ref_name
                     extract_params['psf'] = psf_ref_name
+                    extract_params['optimize_psf_location'] = optimize_psf_location
 
                     break
 
@@ -387,22 +381,9 @@ def log_initial_parameters(extract_params):
         return
 
     log.debug("Extraction parameters:")
-    log.debug(f"dispaxis = {extract_params['dispaxis']}")
-    log.debug(f"spectral order = {extract_params['spectral_order']}")
-    log.debug(f"initial xstart = {extract_params['xstart']}")
-    log.debug(f"initial xstop = {extract_params['xstop']}")
-    log.debug(f"initial ystart = {extract_params['ystart']}")
-    log.debug(f"initial ystop = {extract_params['ystop']}")
-    log.debug(f"extract_width = {extract_params['extract_width']}")
-    log.debug(f"initial src_coeff = {extract_params['src_coeff']}")
-    log.debug(f"initial bkg_coeff = {extract_params['bkg_coeff']}")
-    log.debug(f"bkg_fit = {extract_params['bkg_fit']}")
-    log.debug(f"bkg_order = {extract_params['bkg_order']}")
-    log.debug(f"smoothing_length = {extract_params['smoothing_length']}")
-    log.debug(f"independent_var = {extract_params['independent_var']}")
-    log.debug(f"use_source_posn = {extract_params['use_source_posn']}")
-    log.debug(f"subtract_background = {extract_params['subtract_background']}")
-    log.debug(f"extraction_type = {extract_params['extraction_type']}")
+
+    for key, value in extract_params.items():
+        log.debug(f"  {key} = {value}")
 
 
 def create_poly(coeff):
@@ -1250,7 +1231,8 @@ def define_aperture(input_model, slit, extract_params, exp_type):
     if extract_params['extraction_type'] == 'optimal':
         profiles, lower_limit, upper_limit = psf_profile(
             data_model, extract_params['psf'], extract_params['specwcs'],
-            middle_wl, location, wl_array)
+            middle_wl, location, wl_array,
+            optimize_shifts=extract_params['optimize_psf_location'])
         if len(profiles) > 1:
             profile, nod_profile = profiles
         else:
@@ -1881,8 +1863,8 @@ def run_extract1d(input_model, extract_ref_name="N/A", apcorr_ref_name=None,
                   specwcs_ref_name="N/A", psf_ref_name="N/A", extraction_type="box",
                   smoothing_length=None, bkg_fit=None, bkg_order=None,
                   log_increment=50, subtract_background=None,
-                  use_source_posn=None, save_profile=False,
-                  save_scene_model=False):
+                  use_source_posn=None, optimize_psf_location=True,
+                  save_profile=False, save_scene_model=False):
     """Extract all 1-D spectra from an input model.
 
     Parameters
@@ -1924,6 +1906,10 @@ def run_extract1d(input_model, extract_ref_name="N/A", apcorr_ref_name=None,
         If True, the target and background positions specified in the
         reference file (or the default position, if there is no reference
         file) will be shifted to account for source position offset.
+    optimize_psf_location : bool
+        If True, and if `extraction_type` is 'optimal', then the source
+        location will be optimized, via iterative comparisons of the scene
+        model with the input data.
     save_profile : bool
         If True, the spatial profiles created for the input model will be returned
         as ImageModels. If False, the return value is None.
@@ -2034,6 +2020,7 @@ def run_extract1d(input_model, extract_ref_name="N/A", apcorr_ref_name=None,
                    smoothing_length=smoothing_length,
                    bkg_fit=bkg_fit, bkg_order=bkg_order,
                    use_source_posn=use_source_posn,
+                   optimize_psf_location=optimize_psf_location,
                    subtract_background=subtract_background)
             except ContinueError:
                 continue
@@ -2073,6 +2060,7 @@ def run_extract1d(input_model, extract_ref_name="N/A", apcorr_ref_name=None,
                         smoothing_length=smoothing_length,
                         bkg_fit=bkg_fit, bkg_order=bkg_order,
                         use_source_posn=use_source_posn,
+                        optimize_psf_location=optimize_psf_location,
                         subtract_background=subtract_background)
                 except ContinueError:
                     pass
@@ -2114,6 +2102,7 @@ def run_extract1d(input_model, extract_ref_name="N/A", apcorr_ref_name=None,
                         smoothing_length=smoothing_length,
                         bkg_fit=bkg_fit, bkg_order=bkg_order,
                         use_source_posn=use_source_posn,
+                        optimize_psf_location=optimize_psf_location,
                         subtract_background=subtract_background)
                 except ContinueError:
                     pass

jwst/extract_1d/extract_1d_step.py

@@ -169,6 +169,7 @@ class Extract1dStep(Step):
 
     extraction_type = option("box", "optimal", default="box") # Perform box or optimal extraction
     use_source_posn = boolean(default=None)  # use source coords to center extractions?
+    optimize_psf_location = boolean(default=True)  # For optimal extraction, optimize source location
     smoothing_length = integer(default=None)  # background smoothing size
     bkg_fit = option("poly", "mean", "median", None, default=None)  # background fitting type
     bkg_order = integer(default=None, min=0)  # order of background polynomial fit
@@ -443,6 +444,7 @@ def process(self, input):
                         self.log_increment,
                         self.subtract_background,
                         self.use_source_posn,
+                        self.optimize_psf_location,
                         self.save_profile,
                         self.save_scene_model
                     )

jwst/extract_1d/tests/test_extract.py

@@ -59,6 +59,7 @@ def extract_defaults():
                'src_coeff': None,
                'subtract_background': False,
                'use_source_posn': False,
+               'optimize_psf_location': False,
                'xstart': 0,
                'xstop': 49,
                'ystart': 0,