JP-3588: Update extract_1d_step.py to call pastasoss

jwst/extract_1d/extract_1d_step.py

@@ -426,8 +426,7 @@ def process(self, input):
                     return input_model
 
                 # Load reference files.
-                spectrace_ref_name = self.get_reference_file(input_model, 'spectrace')
-                wavemap_ref_name = self.get_reference_file(input_model, 'wavemap')
+                pastasoss_ref_name = self.get_reference_file(input_model, 'pastasoss') # Add support for "pastasoss" reffile
                 specprofile_ref_name = self.get_reference_file(input_model, 'specprofile')
                 speckernel_ref_name = self.get_reference_file(input_model, 'speckernel')
 
@@ -452,8 +451,7 @@ def process(self, input):
                 # Run the extraction.
                 result, ref_outputs, atoca_outputs = soss_extract.run_extract1d(
                     input_model,
-                    spectrace_ref_name,
-                    wavemap_ref_name,
+                    pastasoss_ref_name,
                     specprofile_ref_name,
                     speckernel_ref_name,
                     subarray,

jwst/extract_1d/soss_extract/soss_extract.py

@@ -13,6 +13,7 @@
 from .soss_syscor import make_background_mask, soss_background
 from .soss_solver import solve_transform, transform_wavemap, transform_profile, transform_coords
 from .atoca import ExtractionEngine, MaskOverlapError
+from .soss_wavemaps import get_soss_wavemaps
 from .atoca_utils import (ThroughputSOSS, WebbKernel, grid_from_map, mask_bad_dispersion_direction,
                           make_combined_adaptive_grid, get_wave_p_or_m, oversample_grid)
 from .soss_boxextract import get_box_weights, box_extract, estim_error_nearest_data
@@ -21,15 +22,12 @@
 log.setLevel(logging.DEBUG)
 
 
-def get_ref_file_args(ref_files, transform):
+def get_ref_file_args(ref_files):
     """Prepare the reference files for the extraction engine.
     Parameters
     ----------
     ref_files : dict
         A dictionary of the reference file DataModels.
-    transform :  array or list
-        A 3-element array describing the rotation and translation to apply
-        to the reference files in order to match the observation.
 
     Returns
     -------
@@ -39,32 +37,18 @@
     """
 
     # The wavelength maps for order 1 and 2.
-    wavemap_ref = ref_files['wavemap']
+    pastasoss_ref = ref_files['pastasoss']
+    pwcpos = pastasoss_ref.map[0].pwcpos        # Make sure the pastasoss_ref object has these attrs
+    subarray = pastasoss_ref.map[0].subarray    # Make sure the pastasoss_ref object has these attrs
+    pad = pastasoss_ref.map[0].padding          # Make sure the pastasoss_ref object has these attrs
 
-    ovs = wavemap_ref.map[0].oversampling
-    pad = wavemap_ref.map[0].padding
+    # Use pastasoss to generate the appropriate wavemap for the given PWCPOS
+    (wavemap_o1, wavemap_o2), (spectrace_o1, spectrace_o2) = get_soss_wavemaps(pwcpos=pwcpos, subarray=subarray, padding=True, padsize=pad)
 
-    wavemap_o1 = transform_wavemap(transform, wavemap_ref.map[0].data, ovs, pad)
-    wavemap_o2 = transform_wavemap(transform, wavemap_ref.map[1].data, ovs, pad)
-
-    # Make sure all pixels follow the expected direction of the dispersion
-    wavemap_o1, flag_o1 = mask_bad_dispersion_direction(wavemap_o1)
-    wavemap_o2, flag_o2 = mask_bad_dispersion_direction(wavemap_o2)
-
-    # Warn if not all pixels were corrected
-    msg_warning = 'Some pixels in order {} do not follow the expected dispersion axis'
-    if not flag_o1:
-        log.warning(msg_warning.format(1))
-    if not flag_o2:
-        log.warning(msg_warning.format(2))
-
-    # The spectral profiles for order 1 and 2.
+    # The spectral profiles for order 1 and 2 (no transform necessary using pastasoss)
     specprofile_ref = ref_files['specprofile']
-    ovs = specprofile_ref.profile[0].oversampling
-    pad = specprofile_ref.profile[0].padding
-
-    specprofile_o1 = transform_profile(transform, specprofile_ref.profile[0].data, ovs, pad, norm=False)
-    specprofile_o2 = transform_profile(transform, specprofile_ref.profile[1].data, ovs, pad, norm=False)
+    specprofile_o1 = specprofile_ref.profile[0].data
+    specprofile_o2 = specprofile_ref.profile[1].data
 
     # The throughput curves for order 1 and 2.
     spectrace_ref = ref_files['spectrace']
@@ -106,18 +90,18 @@
     return [wavemap_o1, wavemap_o2], [specprofile_o1, specprofile_o2], [throughput_o1, throughput_o2], [kernels_o1, kernels_o2]
 
 
-def get_trace_1d(ref_files, transform, order, cols=None):
+def get_trace_1d(ref_files, order, transform=None, cols=None):
     """Get the x, y, wavelength of the trace after applying the transform.
     Parameters
     ----------
     ref_files : dict
         A dictionary of the reference file DataModels.
-    transform : array or list
+    order : int
+        The spectral order for which to return the trace parameters.
+    transform : array or list, optional
         A 3-element list or array describing the rotation and translation
         to apply to the reference files in order to match the
         observation.
-    order : int
-        The spectral order for which to return the trace parameters.
     cols : array[int], optional
         The columns on the detector for which to compute the trace
         parameters. If not given, all columns will be computed.
@@ -133,21 +117,30 @@
     else:
         xtrace = cols
 
-    spectrace_ref = ref_files['spectrace']
+    pastasoss_ref = ref_files['pastasoss']
 
     # Read x, y, wavelength for the relevant order.
-    xref = spectrace_ref.trace[order - 1].data['X']
-    yref = spectrace_ref.trace[order - 1].data['Y']
-    waveref = spectrace_ref.trace[order - 1].data['WAVELENGTH']
+    xref = pastasoss_ref.trace[order - 1].data['X']
+    yref = pastasoss_ref.trace[order - 1].data['Y']
+    waveref = pastasoss_ref.trace[order - 1].data['WAVELENGTH']
+
+    # No transform necessary is using pastasoss
+    if transform is None:
+        xtrace = xref
+        ytrace = yref
+        wavetrace = waveref
+
+    # Transform the trace if using a static PWCPOS
+    else:
 
-    # Rotate and shift the positions based on transform.
-    angle, xshift, yshift = transform
-    xrot, yrot = transform_coords(angle, xshift, yshift, xref, yref)
+        # Rotate and shift the positions based on transform.
+        angle, xshift, yshift = transform
+        xrot, yrot = transform_coords(angle, xshift, yshift, xref, yref)
 
-    # Interpolate y and wavelength to the requested columns.
-    sort = np.argsort(xrot)
-    ytrace = np.interp(xtrace, xrot[sort], yrot[sort])
-    wavetrace = np.interp(xtrace, xrot[sort], waveref[sort])
+        # Interpolate y and wavelength to the requested columns.
+        sort = np.argsort(xrot)
+        ytrace = np.interp(xtrace, xrot[sort], yrot[sort])
+        wavetrace = np.interp(xtrace, xrot[sort], waveref[sort])
 
     return xtrace, ytrace, wavetrace
 
@@ -1015,18 +1008,16 @@
     return fluxes, fluxerrs, npixels
 
 
-def run_extract1d(input_model, spectrace_ref_name, wavemap_ref_name,
+def run_extract1d(input_model, pastasoss_ref_name,
                   specprofile_ref_name, speckernel_ref_name, subarray,
                   soss_filter, soss_kwargs):
     """Run the spectral extraction on NIRISS SOSS data.
     Parameters
     ----------
     input_model : DataModel
         The input DataModel.
-    spectrace_ref_name : str
-        Name of the spectrace reference file.
-    wavemap_ref_name : str
-        Name of the wavemap reference file.
+    pastasoss_ref_name : str
+        Name of the pastasoss reference file.
     specprofile_ref_name : str
         Name of the specprofile reference file.
     speckernel_ref_name : str
@@ -1052,14 +1043,12 @@
     order_str_2_int = {f'Order {order}': order for order in [1, 2, 3]}
 
     # Read the reference files.
-    spectrace_ref = datamodels.SpecTraceModel(spectrace_ref_name)
-    wavemap_ref = datamodels.WaveMapModel(wavemap_ref_name)
+    pastasoss_ref = datamodels.PastasossModel(pastasoss_ref_name)
     specprofile_ref = datamodels.SpecProfileModel(specprofile_ref_name)
     speckernel_ref = datamodels.SpecKernelModel(speckernel_ref_name)
 
     ref_files = dict()
-    ref_files['spectrace'] = spectrace_ref
-    ref_files['wavemap'] = wavemap_ref
+    ref_files['pastasoss'] = pastasoss_ref
     ref_files['specprofile'] = specprofile_ref
     ref_files['speckernel'] = speckernel_ref