few updates for new PSF model and option to write spatial profile

jwst/extract_1d/extract_1d_step.py

@@ -55,7 +55,11 @@ class Extract1dStep(Step):
         extract_1d reference file.
 
     psf_extraction : bool or None
-        If True, a optiminal psf extraction is performed. 
+        If True, a optimal psf extraction is performed. Currently this option
+        only works for MIRI LRS Fixed and Slitless data. 
+
+    save_spatial_profile: bool or None
+        If True and psf_extraction is True save the spatial profile to disk
 
     use_source_posn : bool or None
         If True, the source and background extraction positions specified in
@@ -163,6 +167,7 @@ class Extract1dStep(Step):
     log_increment = integer(default=50)  # increment for multi-integration log messages
     subtract_background = boolean(default=None)  # subtract background?
     psf_extraction = boolean(default=None) # Perform psf extraction
+    save_spatial_profile = boolean(default=False) # If doing psf extraction, then save the spatial profile to disk
     use_source_posn = boolean(default=None)  # use source coords to center extractions?
     center_xy = float_list(min=2, max=2, default=None)  # IFU extraction x/y center
     apply_apcorr = boolean(default=True)  # apply aperture corrections?
@@ -185,6 +190,7 @@ class Extract1dStep(Step):
     soss_modelname = output_file(default = None)  # Filename for optional model output of traces and pixel weights
     """
 
+    # TODO JEM - add PSF once that type is in CRDS
     reference_file_types = ['extract1d', 'apcorr', 'pastasoss', 'specprofile', 'speckernel', 'specwcs']
 
     def process(self, input):
@@ -243,10 +249,10 @@ def process(self, input):
         elif isinstance(input_model, ModelContainer):
             exp_type = input_model[0].meta.exposure.type
         else:
-            exp_type = None
+            # TODO JEM CAN I do this or will it break other modes
+            #exp_type = None to supprt PSF extraction commented out exp_type = None
+            exp_type = input_model.meta.exposure.type
 
-        print('************ exp_type ************', exp_type)
-        print(isinstance(input_model, ModelContainer)) 
         if self.ifu_rfcorr:
             if exp_type != "MIR_MRS":
                 self.log.warning("The option to apply a residual refringe correction is"
@@ -263,15 +269,22 @@ def process(self, input):
                                  f" not supported for {input_model.meta.exposure.type} data.")
 
         # Check if doing PSF extraction. Only works on LRS data at this time
+
         if self.psf_extraction:
-            # the input_model could be a container or a single model
-            # for now we are assuming a single model of LRS data 
-            print(type(input_model)) # cal file could have other nod position subtracted
-            print(input_model.meta.exposure.type)
-            #psf_ref_name = self.get_reference_file(input_model, 'psf')
-            psf_ref_name = '/Users/morrison/PSF_extract/MIRI_LRS_SLIT_EPSF_20240602_WAVE.fits'
-            specwcs_ref_name = self.get_reference_file(input_model, 'specwcs')
-            psf_extract.psf_extraction(input_model, psf_ref_name, specwcs_ref_name)
+            if exp_type != 'MIR_LRS-FIXEDSLIT':
+                  self.log.info ('PSF extraction is not possible with this mode %s', exp_type)
+                  self.psf_extraction = False
+            else: 
+                # the input_model could be a container or a single model
+
+                # once the reference file is in CRDS we can use get_reference_file. For
+                # how hard coding it because PSF is not a reference file type. 
+                #psf_ref_name = self.get_reference_file(input_model, 'psf')
+                # UNTIL WE GET THIS IN CRDS GET REFERENCE FILE FROM JEM 
+                psf_ref_name = 'MIRI_LRS_SLIT_EPSF_20240602_WAVE_1D.fits'
+                specwcs_ref_name = self.get_reference_file(input_model, 'specwcs')
+                psf_extract.psf_extraction(input_model, psf_ref_name, specwcs_ref_name,
+                                           self.save_spatial_profile)
 
         # ______________________________________________________________________
         # Do the extraction for ModelContainer - this might only be WFSS data

jwst/extract_1d/psf_extract.py

@@ -16,7 +16,7 @@
 from stdatamodels.properties import ObjectNode
 from stdatamodels.jwst import datamodels
 from stdatamodels import DataModel
-from stdatamodels.jwst.datamodels import dqflags, SlitModel, SpecModel,SpecwcsModel, PsfModel
+from stdatamodels.jwst.datamodels import dqflags, SlitModel, SpecModel,SpecwcsModel, MiriLrsPsfModel
 from jwst.datamodels import SourceModelContainer
 
 from ..assign_wcs.util import wcs_bbox_from_shape
@@ -25,27 +25,38 @@
 from . import extract1d
 from . import ifu
 from . import spec_wcs
-from . import utils
+from . import utils   # these are routines taken from extract.py and pulled out of the class definition. 
 
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
 
 
-
 HORIZONTAL = 1
 VERTICAL = 2
 """Dispersion direction, predominantly horizontal or vertical."""
 
-# This is intended to be larger than any possible distance (in pixels) between the target and any point in the image;
-# used by locn_from_wcs().
-HUGE_DIST = 1.e20
 
-def open_specwcs(specwcs_ref_name: str, exp_type):
-    print('in open specwcs', exp_type)
+def open_specwcs(specwcs_ref_name: str, exp_type: str):
+    """Open the specwcs reference file.
+
+    Parameters
+    ----------
+    specwcs_ref_name : str
+        The name of the specwcs reference file. This file contains information of the trace location.
+        For the MIRI LRS FIED-SlIT it is a fits file containing the x,y center of the trace. 
+    ext_type : str
+        The exposure type of the data
+
+    Returns
+    -------
+        Currently only works on MIRI LRS-FIXEDSLIT exposures. 
+        Return the center of the trace in x and y for a given wavelength
+
+    """
+
     if exp_type == 'MIR_LRS-FIXEDSLIT':
         # use fits to read file (the datamodel does not have all that is needed) 
         ref = fits.open(specwcs_ref_name)
-        print(' Number of rows in the table', ref[1].data.shape)
 
         with ref:
             lrsdata = np.array([d for d in ref[1].data])
@@ -54,11 +65,9 @@ def open_specwcs(specwcs_ref_name: str, exp_type):
             # These are 1-indexed in CDP-7 (i.e., SIAF convention) so must be converted to 0-indexed
             # for lrs_fixedslit
             zero_point = ref[0].header['imx'] - 1, ref[0].header['imy'] - 1
-            print('zero point', zero_point) 
 
         # In the lrsdata reference table, X_center,Y_center, wavelength  relative to zero_point
-        # x0,y0(ul) x1,y1 (ur) x2,y2(lr) x3,y3(ll) define corners of the box within which the distortion
-        # and wavelength calibration was derived
+
         xcen = lrsdata[:, 0]
         ycen = lrsdata[:, 1]
         wavetab = lrsdata[:, 2]
@@ -67,28 +76,43 @@ def open_specwcs(specwcs_ref_name: str, exp_type):
     return trace, wave_trace, wavetab
 
 
-def open_psf(psf_refname: str, exp_type):
+def open_psf(psf_refname: str, exp_type: str):
+    """Open the PSF reference file.
 
+    Parameters
+    ----------
+    psf_ref_name : str
+        The name of the psf reference file. 
+    ext_type : str
+        The exposure type of the data
+
+    Returns
+    -------
+        Currenly only works on MIRI LRS-FIXEDSLIT exposures. 
+        Returns the EPSF model.
+
+    """
+
     psf_model = None
     if exp_type == 'MIR_LRS-FIXEDSLIT':
-        psf_model = PsfModel(psf_refname)
-        print(psf_model.meta.psf.center_col)
-        print(psf_model.meta.psf.subpix)
-        print(psf_model.data)
-        print(psf_model.wave) 
+        psf_model = MiriLrsPsfModel(psf_refname)
+        # The information we read in from PSF file is:
+        # center_col: psf_model.meta.psf.center_col
+        # super sample factor: psf_model.meta.psf.subpix)
+        # psf : psf_model.data (2d)
+        # wavelength of PSF planes: psf_model.wave 
     return psf_model 
 
 
 
+# TODO JEM  THIS CLASS IS NOT BEING USED YET. JUST HOOKS FOR LATER
 class PSFExtractData():
     """ Class for PSF extraction for each source"""
 
     def __init__(self,
                  input_model: DataModel,
                  slit: Union[DataModel, None] = None,
-                 dispaxis: int = HORIZONTAL,
-                 use_source_posn: Union[bool, None] = None):
-
+                 dispaxis: int = HORIZONTAL):
 
         """
         Parameters
@@ -109,7 +133,6 @@ def __init__(self,
         self.dispaxis = dispaxis
         self.wcs = None
 
-
         if slit is None:
             if hasattr(input_model.meta, 'wcs'):
                 self.wcs = input_model.meta.wcs
@@ -120,11 +143,28 @@ def __init__(self,
             log.warning("WCS function not found in input.")
 
 
+def psf_extraction(input_model, psf_ref_name, specwcs_ref_name, save_spatial_profile):
+    """Open the PSF reference file.
 
-
-
-def psf_extraction(input_model, psf_ref_name, specwcs_ref_name):
+    Parameters
+    ----------
+    input_model : data model
+        This can be either the input science file or one SlitModel out of
+        a list of slits.
+    psf_ref_name : str
+        PSF reference filename
+    specwcs_ref_name : str
+        Reference file containing information on trace of spectra
+
+    save_spatial_profile : bool
+        Save the spatial profile. 
+    Returns
+    -------
+        Currenly only works on MIRI LRS-FIXEDSLIT exposures. 
+        spatial profile
 
+    """
+
     exp_type = input_model.meta.exposure.type
     trace, wave_trace, wavetab = open_specwcs(specwcs_ref_name, exp_type)
     psf_model = open_psf(psf_ref_name, exp_type)
@@ -147,7 +187,7 @@ def psf_extraction(input_model, psf_ref_name, specwcs_ref_name):
 
     for model in input_models:
 
-        # currently for LRS-FIXEDSLIT data it is assume there is 1 source and the source is located
+        # currently for LRS-FIXEDSLIT data it is assumes there is 1 source and the source is located
         # at the dither_position
 
         if exp_type == 'MIR_LRS-FIXEDSLIT': # assume we can use the WCS to find source location
@@ -158,26 +198,27 @@ def psf_extraction(input_model, psf_ref_name, specwcs_ref_name):
             # Set up the profiles that are to be determined
             nsource = 1
             profile_2d = []
-            profile_bg = []  # Currently not sure if this is needed
+            profile_bg = []  # Currently not sure if this is needed (at least for MIRI LRS data) 
                              # Calwebb_spec2 subtracts the other nod position automatically 
 
-            # Only one source 
+            # Only one source determine the location using the WCS
             middle, locn_w, locn_x = utils.locn_from_wcs(
                 dispaxis,
                 wcs, 
                 model,
                 None, None, None)
 
-            # Match the  reference trace and corresponding wavelength to the data
+            # Perform fit of reference trace and corresponding wavelength
             # The wavelength for the reference trace does not exactly line up exactly with the data
 
             cs = CubicSpline(wavetab, trace)
             cen_shift = cs(locn_w)
             shift = locn_x - cen_shift
 
-            print('Location of source in observed spectrum', locn_x, ' at wavelength', locn_w)
-            print('In the reference trace the location is ', cen_shift)
-            print('Shift to apply to ref trace', shift)
+            log.info ('Location of source in observed spectrum %s at wavelength %s', locn_x, locn_w)
+
+            log.info('For this wavelength: %s  the reference trace the location is %s ', locn_w, cen_shift)
+            log.info('Shift to apply to ref trace %s', shift)
 
             # For LRS-Fix data  we want to cut out the slit from the full array
             slit = []
@@ -187,31 +228,58 @@ def psf_extraction(input_model, psf_ref_name, specwcs_ref_name):
             cutout = model.data[int(np.ceil(y0)):int(np.ceil(y1)), int(np.round(x0)):int(np.round(x1))]
             slit.append(cutout)
 
-            # adjust the trace to for slit region 
+            # adjust the trace to the slit region 
             trace_cutout = trace - bbox[0][0]
             wtrace_cutout = wave_trace - bbox[1][0]
-            trace_shift = trace_cutout + xshift
-
-            # xtrace_shift - for each wavelength in the PSF this is the shift in x to apply to the PSF image to shift it
+            trace_shift = trace_cutout + shift
+            psf_wave = psf_model.wave
+
+            # trace_shift - for each wavelength in the PSF this is the shift in x to apply to the PSF image to shift it
             #                 to fall on the source.
-            # wavetab : is the wavelengh cooresponding the the trace. This wavelength may not match exactly to the the PSF.wave
-
+            # wavetab : is the wavelengh cooresponding the the trace. This wavelength may not match exactly to the the PSF.
+            # 
             # determine what the shifts per row are for the the wavelengths given by the model PSF 
 
-            PSFinterp = interpolate.interp1d(wavetab, xtrace_shift,fill_value="extrapolate")
-            psf_shift = PSFinterp(psf_model.wave)
-            psf_shift =  psf_model.meta.center_col - (psf_shift * psf_subpix)
+            psf_subpix = psf_model.meta.psf.subpix
+
+            PSFinterp = interpolate.interp1d(wavetab, trace_shift,fill_value="extrapolate")
+            psf_shift = PSFinterp(psf_wave)
+            psf_shift =  psf_model.meta.psf.center_col - (psf_shift * psf_subpix)
 
             # if the PSF is an ePSF
             _x, _y = np.meshgrid(np.arange(cutout.shape[1]), np.arange(cutout.shape[0]))
-            _x = _x*psf_subpix + psf_shift[:, np.newaxis]
-            sprofile = ndimage.map_coordinates(psf, [_y, _x], order=1)
+
+            # Convert 1D psf to 2 d with second dimension = 1
+            psf2 = psf_shift[:, np.newaxis]
+            _x = _x*psf_subpix + psf2
+
+            sprofile = ndimage.map_coordinates(psf_model.data, [_y, _x], order=1)
             profile_2d.append(sprofile)
-
-            print(sprofile)
+
+            # TODO JEM - this might be just for testing. If we keep writing spatial profile
+            # we will need a datamodel module for it  rather than using FITS write
+            # Also need to change how we form the name if we keep writing spatial profile
+            if save_spatial_profile:
+                file_out = model.meta.filename.replace('.fits', '_spatial_profile.fits')
+                hdu = fits.PrimaryHDU()
+                hdu.header['SUBPIX'] = psf_model.meta.psf.subpix
+                hdu.header['CENTCOL'] = psf_model.meta.psf.center_col
+                hdu.header['CENTWAVE'] = locn_w
+                hdu.header['CENTLOC'] = locn_x
+
+                hdu.header['REFSHFT'] = float(cen_shift)
+                hdu.header['TRCSHFT'] = shift
+                hdu1 = fits.ImageHDU(sprofile, name='SPROFILE')
+                hdu_final = fits.HDUList([hdu, hdu1])
+
+                print(file_out) 
+                hdu_final.writeto(file_out, overwrite=True)
 
             # set up the data to be passed to extract1d
-            result = util.setup_data(model)
+            # TODO JEM - setup data - gathers the data that extract1d will need.
+            # we do not have to do this in seperate module and we need to know what
+            # Tim B routine needs - so this will need to be updated. 
+            result = utils.setup_data(model)
             data, dq, var_poisson, var_rnoise, weights = result