Repository: spacetelescope/jwst

Issue:
wavecorr offsets are ignored in spectral resampling
So `wavecorr` corrects for wavelength offset due to non-centering of sources in the MSA (or fixed) slit(let).  Due to the nature of fixed centers of the MSA, just about every source is offset to some degree.

But only the point sources are corrected, i.e. all the extra-galactic sources, if labelled EXTENDED, are not corrected at all.

Anyway, the correction is not an update to the `meta.wcs` GWCS object of each 2D extracted cutout in the _cal.fits file, but an updated wavelength array in that _cal file.  This updated wavelength array is not always then used downstream, specifically in `resample_spec`, where the wavelength array is ignored in the building of the output wcs and specifically in the resampling in `reproject` when producing the drizzle pixmap, where it just uses the GWCS objects to generate the pixel-to-pixel mapping.

https://github.com/spacetelescope/jwst/blob/12030b54c269b8eef22e18f8b7f7f0a565a7e55b/jwst/resample/resample_utils.py#L128

Not sure if there is an easy fix to this, given that the wavelength offset is interpolated from a model and only written to the wavelength array, and not currently part of the GWCS transform.  But it should probably be documented, both for point sources and for extended, point-like sources, i.e. high-redshift galaxies.

Discussion:
nden: Regression tests run: https://plwishmaster.stsci.edu:8081/job/RT/job/JWST-Developers-Pull-Requests/1364/
nden: The implementation in this PR matches the description in the note attached to the issue and the JIRA ticket.
The documentation needs to be updated to include the new algorithm.
One further change to consider is whether the wavelength correction should be a separate frame in the WCS pipeline. As it stands now, the only way to exclude this correction is to rerun the pipeline and skip wavecorr. If it's a separate frame in the WCS pipeline one could evaluate the WCS skipping that frame. However, this is a convenience and purely optional.
hayescr: The new commits address the previous comments, update the documentation, and add in the wavelength correction as a separate WCS frame, `wavecorr_frame` following the `slit_frame`.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) Report
Attention: Patch coverage is `82.97872%` with `8 lines` in your changes are missing coverage. Please review.
> Project coverage is 75.39%. Comparing base [(`2fb073e`)](https://app.codecov.io/gh/spacetelescope/jwst/commit/2fb073e828d06f18ef7bd163445007c463b45136?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) to head [(`c2cd640`)](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope).
> Report is 34 commits behind head on master.

| [Files](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) | Patch % | Lines |
|---|---|---|
| [jwst/wavecorr/wavecorr.py](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376?src=pr&el=tree&filepath=jwst%2Fwavecorr%2Fwavecorr.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope#diff-andzdC93YXZlY29yci93YXZlY29yci5weQ==) | 82.97% | [8 Missing :warning: ](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #8376      +/-   ##
==========================================
+ Coverage   75.31%   75.39%   +0.08%     
==========================================
  Files         474      476       +2     
  Lines       38965    41894    +2929     
==========================================
+ Hits        29345    31587    +2242     
- Misses       9620    10307     +687     
```

| [Flag](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) | Coverage Δ | | *Carryforward flag |
|---|---|---|---|
| [nightly](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) | `77.32% <ø> (-0.02%)` | :arrow_down: | Carriedforward from [2fb073e](https://app.codecov.io/gh/spacetelescope/jwst/commit/2fb073e828d06f18ef7bd163445007c463b45136?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) |

*This pull request uses carry forward flags. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/spacetelescope/jwst/pull/8376?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope).

nden: @hayescr Is this ready for review? If so do you want to take it out of Draft?
nden: This PR needs a change log entry.

Another regression tests job started: https://plwishmaster.stsci.edu:8081/job/RT/job/JWST-Developers-Pull-Requests/1414/
melanieclarke: I think these changes might require some changes to the flat_field step as well: in the `flat_for_nirspec_slit` function, there is some handling to produce a flat field with and without wavelength correction.  The wavecorr=False option will need to be updated, since it currently assumes the wavelength correction is not in the WCS.
hayescr: > I think these changes might require some changes to the flat_field step as well: in the `flat_for_nirspec_slit` function, there is some handling to produce a flat field with and without wavelength correction. The use_wavecorr=False option will need to be updated, since it currently assumes the wavelength correction is not in the WCS.
> 
> Edited to add: also in the get_wavelengths function in lib/wcs_utils.py

Agreed it looks like these instances primarily use the wcs to generate the uncorrected wavelengths for fixed slit background corrections (since the background will be a extended source), which will no longer work, so I can work on adding fixes here.  @nden and @melanieclarke, would it be better to reproduce uncorrected wavelengths from the wcs for these steps (by skipping the wavelength correction transforms), or to add a "uniform wavelength/uncorrected wavelength" extension that would be saved in the wavecorr step in case users wanted access these after-the-fact, similar to how the `flat_field`/`pathloss`/`photom` steps add the extended source corrections as extra extensions for point sources?

The `pathloss` step also calculates point source and extended source corrections for fixed slit observations (so it can apply the right corrections for point sources and the uniform background in `master_background_step`), but it is currently using the same wavelength for the point source and uniform pathloss corrections.  Should this be corrected as well (see the lines below)?

https://github.com/hayescr/jwst/blob/250a44a54618ff9278d9b8d715746ab914701f96/jwst/pathloss/pathloss.py#L979-L991
melanieclarke: @hayescr - For the pathloss correction wavelengths, I think you're right that it shouldn't be using the corrected wavelength array for both corrections.  It should probably use uncorrected wavelengths for uniform, corrected wavelengths for point source. But you may need to double check with Elena to make sure the pathloss correction isn't expecting uncorrected wavelengths for both.
hayescr: @nden this is ready for re-review.  I've updated the steps that depended on the old format of wavecorr.  

For `flatfield`, I refactored the way `flat_for_nirspec_slit` checks for wavelengths to have it use the updated `lib.wcs_utils.get_wavelength`, because the two were getting the wavelengths in a similar way (i.e., looking for a wavelength array first and then using the wcs to calculate wavelengths if an array wasn't available).  There are some log warnings/errors that `flat_for_nirspec_slit` included that `get_wavelength` does not, but I think that should be the only difference in actual behavior.  Let me know if it looks like I missed something though. 

I checked with Elena Manjavacas about the `pathloss` step for fixed slit:  it should be using the corrected wavelengths for point sources and uncorrected wavelengths for uniform sources, so I've pointed this step to also use `get_wavelength` for the two cases.

The `wavecorr` step is now marked as complete/skipped for each individual MOS slit separately, and this shows up in the 2d unrectified spectra files (e.g., the cal files), but one issue is that these appear to be overwritten in the resampled (s2d) spectra.  Perhaps the resampling step is populating the slit metadata with the multislit metadata unless it's told otherwise?
hayescr: @melanieclarke thanks for the review and comments on the change log.  I took a look at the example you listed, some comments below on each case:

PATHLOSS_PS is set to unity when wavecorr is off:  This is because this example is a FS observation and currently wavecorr is what calculates the source x and y position in the slit for FS point sources.  So with wavecorr off the source_xpos and source_ypos both default to 0 and by definition the path loss corrections at the center of the slit are unity.  This is not introduced by this PR, but because the source positions for point sources are used in a couple places (pathloss, extract_1d for use_source_pos, etc.) it might be good to move this out of wavecorr and somewhere else in the spec2 pipeline (assign_wcs perhaps?).

PATHLOSS_UN show some small differences when wavecorr is off:  in my testing of this, I see that the (non-NaN) differences in PATHLOSS_UN between wavecorr on and off is either 0 or a small number (~1.2e-7) that looks like a floating point error, does this match with what you found?

FLATFIELD_UN show some small differences when wavecorr is off:  I also see this; the fractional differences are at the level of < ~1e-5.  I think this comes from the difference in how the wavelengths are obtained.  With the update I made to have flat_field use get_wavelengths, the wavelengths used for interpolating the flat_field data come from the WCS when wavecorr is on and will come from the stored wavelength array when wavecorr is off.  If I force the flat_field to use the WCS for calculating the wavelengths in both cases I find no difference in the FLATFIELD_UN.  Comparing the array and WCS wavelengths there is a small difference (on the order of < 1e-7um -> though that is about 1e-4 of a pixel in this case), which looks to be coming from the fact that the wavelength arrays are float32 and recalculating the wavelengths from the WCS are float64.
hayescr: @melanieclarke yes I'll open a ticket to move the source position assignment.  Thanks for the catch on the missing wcs case.
melanieclarke: Thanks @hayescr. 

I started a new regression test run here:
https://plwishmaster.stsci.edu:8081/job/RT/job/JWST-Developers-Pull-Requests/1498/
melanieclarke: Except one thing - @hayescr, can you please fix the merge conflicts in the change log?

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 2169ae668..08c5e3fe3 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -108,6 +108,9 @@ flat_field
 - Update NIRSpec flatfield code for all modes to ensure SCI=ERR=NaN wherever the
   DO_NOT_USE flag is set in the DQ array. [#8463]
 
+- Updated the NIRSpec flatfield code to use the new format of the ``wavecorr`` 
+  wavelength zero-point corrections for point sources.  [#8376]
+
 general
 -------
 
@@ -116,6 +119,13 @@ general
 
 - Increase minimum required scipy. [#8441]
 
+lib
+---
+
+- Updated the ``wcs_utils.get_wavelength`` to use the new format
+  of the ``wavecorr`` wavelength zero-point corrections for point 
+  sources in NIRSpec slit data. [#8376]
+
 master_background_mos
 ---------------------
 
@@ -149,6 +159,13 @@ outlier_detection
   to detect outliers in TSO data, with user-defined
   rolling window width via the ``rolling_window_width`` parameter. [#8473]
 
+pathloss
+--------
+
+- Updated pathloss calculations for NIRSpec fixed slit mode to use the appropriate 
+  wavelengths for point and uniform sources if the ``wavecorr`` wavelength 
+  zero-point corrections for point sources have been applied. [#8376]
+
 photom
 ------
 
@@ -242,6 +259,12 @@ tweakreg
   message and skip ``tweakreg`` step when this condition is not satisfied and
   source confusion is possible during catalog matching. [#8476]
 
+wavecorr
+--------
+
+- Changed the NIRSpec wavelength correction algorithm to include it in slit WCS
+  models and resampling.  Fixed the sign of the wavelength corrections. [#8376]
+  
 wfss_contam
 -----------
 
diff --git a/docs/jwst/pathloss/description.rst b/docs/jwst/pathloss/description.rst
index 2e10ad4cb..65529b736 100644
--- a/docs/jwst/pathloss/description.rst
+++ b/docs/jwst/pathloss/description.rst
@@ -45,6 +45,10 @@ Once the 1-D correction arrays have been computed, both forms of the correction
 (point and uniform) are interpolated, as a function of wavelength, into
 the 2-D space of the slit or IFU data and attached to the output data model
 (extensions "PATHLOSS_PS" and "PATHLOSS_UN") as a record of what was computed.
+For fixed slit data, if the ``wavecorr`` step has been run to provide wavelength
+corrections to point sources, the corrected wavelengths will be used to
+calculate the point source pathloss, whereas the uncorrected wavelengths (appropriate 
+for uniform sources) will be used to calculate the uniform source pathlosses.
 The form of the 2-D correction (point or uniform) that's appropriate for the
 data is divided into the SCI and ERR arrays and propagated into the variance
 arrays of the science data.
diff --git a/docs/jwst/wavecorr/description.rst b/docs/jwst/wavecorr/description.rst
index 999ee6a6d..fa7e311d5 100644
--- a/docs/jwst/wavecorr/description.rst
+++ b/docs/jwst/wavecorr/description.rst
@@ -22,16 +22,22 @@ These are recorded in the "SRCXPOS" and "SRCYPOS" keywords in the SCI
 extension header of each slitlet in a FITS product.
 
 The ``wavecorr`` step loops over all slit instances in the input
-science product and applies a wavelength correction to slits that
-contain a point source. The point source determination is based on the
-value of the "SRCTYPE" keyword populated for each slit by the
+science product and updates the WCS models of slits that contain a point 
+source to include a wavelength correction. The point source determination is 
+based on the value of the "SRCTYPE" keyword populated for each slit by the
 :ref:`srctype <srctype_step>` step. The computation of the correction is
 based on the "SRCXPOS" value. A value of 0.0 indicates a perfectly centered
 source, and ranges from -0.5 to +0.5 for sources at the extreme edges
 of a slit. The computation uses calibration data from the ``WAVECORR``
-reference file. The correction is computed as a 2-D grid of
-wavelength offsets, which is applied to the original 2-D grid of
-wavelengths associated with each slit.
+reference file, which contains pixel shifts as a function of source position 
+and wavelength, and can be converted to wavelength shifts with the dispersion. 
+For each slit, the ``wavecorr`` step uses the average wavelengths and 
+dispersions in a slit (averaged across the cross-dispersion direction) to 
+calculate corresponding corrected wavelengths.  It then uses the average 
+wavelengths and their corrections to generate a transform that interpolates 
+between "center of slit" wavelengths and corrected wavelengths.  This 
+transformation is added to the slit WCS after the ``slit_frame`` and
+produces a new wavelength corrected slit frame, ``wavecorr_frame``.
 
 NIRSpec Fixed Slit (FS)
 -----------------------
@@ -47,8 +53,8 @@ the wavelength correction is only applied to the primary slit.
 
 The estimated position of the source within the primary slit (in the
 dispersion direction) is then used in the same manner as described above
-for MOS slitlets to compute offsets to be added to the nominal wavelength
-grid for the primary slit.
+for MOS slitlets to update the slit WCS and compute corrected wavelengths
+for the primary slit.
 
 Upon successful completion of the step, the status keyword "S_WAVCOR"
 is set to "COMPLETE".
diff --git a/jwst/flatfield/flat_field.py b/jwst/flatfield/flat_field.py
index 1c4c1db6e..01da31668 100644
--- a/jwst/flatfield/flat_field.py
+++ b/jwst/flatfield/flat_field.py
@@ -6,12 +6,11 @@ import logging
 import math
 
 import numpy as np
-from gwcs.wcstools import grid_from_bounding_box
 
 from stdatamodels.jwst import datamodels
 from stdatamodels.jwst.datamodels import dqflags
 
-from ..lib import reffile_utils
+from ..lib import reffile_utils, wcs_utils
 from ..assign_wcs import nirspec
 
 log = logging.getLogger(__name__)
@@ -1924,59 +1923,11 @@ def flat_for_nirspec_slit(slit, f_flat_model, s_flat_model, d_flat_model,
     # Get the wavelength at each pixel in the extracted slit data.
     # If the wavelength attribute exists and is populated, use it
     # in preference to the wavelengths returned by the wcs function.
-    got_wl_attribute = True
-    try:
-        wl = slit.wavelength.copy()  # a 2-D array
-    except AttributeError:
-        got_wl_attribute = False
-    if not got_wl_attribute or len(wl) == 0:
-        got_wl_attribute = False
-    return_dummy = False
 
-    # Has the use_wavecorr param been set?
-    if use_wavecorr is not None:
-        if use_wavecorr:
-            # Need to use the 2D wavelength array, because that's where
-            # the corrected wavelengths are stored
-            if got_wl_attribute:
-                # We've got the "wl" wavelength array we need
-                pass
-            else:
-                # Can't do the computation without the 2D wavelength array
-                log.error(f"The wavelength array for slit {slit.name} is not populated")
-                log.error("Skipping flat-field correction")
-                return_dummy = True
-        elif not use_wavecorr:
-            # Need to use the WCS object to create an uncorrected 2D wavelength array
-            if got_wcs:
-                log.info(f"Creating wavelength array from WCS for slit {slit.name}")
-                bb = slit.meta.wcs.bounding_box
-                grid = grid_from_bounding_box(bb)
-                wl = slit.meta.wcs(*grid)[2]
-                del grid
-            else:
-                # Can't create the uncorrected wavelengths without the WCS
-                log.error(f"Slit {slit.name} has no WCS object")
-                log.error("Skipping flat-field correction")
-                return_dummy = True
-    else:
-        # use_wavecorr was not specified, so use default processing
-        if not got_wl_attribute or np.nanmin(wl) == 0. and np.nanmax(wl) == 0.:
-            got_wl_attribute = False
-            log.warning(f"The wavelength array for slit {slit.name} has not been populated")
-            # Try to create it from the WCS
-            if got_wcs:
-                bb = slit.meta.wcs.bounding_box
-                grid = grid_from_bounding_box(bb)
-                wl = slit.meta.wcs(*grid)[2]
-                del grid
-            else:
-                log.warning("and this slit does not have a 'wcs' attribute")
-                log.warning("likely because assign_wcs has not been run.")
-                log.error("skipping ...")
-                return_dummy = True
-        else:
-            log.debug("Wavelengths are from the wavelength array.")
+    return_dummy = False
+    wl = wcs_utils.get_wavelengths(slit, use_wavecorr=use_wavecorr)
+    if wl is None:
+        return_dummy = True
 
     # Create and return a dummy flat as a placeholder, if necessary
     if return_dummy:
diff --git a/jwst/lib/wcs_utils.py b/jwst/lib/wcs_utils.py
index 5dbfceeff..3070756f1 100644
--- a/jwst/lib/wcs_utils.py
+++ b/jwst/lib/wcs_utils.py
@@ -41,19 +41,24 @@ def get_wavelengths(model, exp_type="", order=None, use_wavecorr=None):
         got_wavelength = False
         wl_array = None
 
-    # If we've been asked to use the corrected wavelengths stored in
-    # the wavelength array, return those wavelengths. Otherwise, the
-    # results computed from the WCS object (below) will be returned.
-    if use_wavecorr is not None:
-        if use_wavecorr:
-            return wl_array
-        else:
-            got_wavelength = False  # force wl computation below
 
     # Evaluate the WCS on the grid of pixel indexes, capturing only the
     # resulting wavelength values
     shape = model.data.shape
     grid = np.indices(shape[-2:], dtype=np.float64)
+    
+    # If we've been asked to use the uncorrected wavelengths we need to
+    # recalculate them from the wcs by skipping the transformation between
+    # the slit frame and the wavelength corrected slit frame.  If the wavecorr_frame
+    # is not in the wcs assume that the wavelength correction has not been applied.
+    if use_wavecorr is not None:
+        if (not use_wavecorr and hasattr(model.meta, "wcs")
+                and 'wavecorr_frame' in model.meta.wcs.available_frames):
+            wcs = model.meta.wcs
+            detector2slit = wcs.get_transform('detector', 'slit_frame')
+            wavecorr2world = wcs.get_transform("wavecorr_frame", "world")
+            wl_array = (detector2slit | wavecorr2world)(grid[1], grid[0])[2]
+            return wl_array
 
     # If no existing wavelength array, compute one
     if hasattr(model.meta, "wcs") and not got_wavelength:
diff --git a/jwst/pathloss/pathloss.py b/jwst/pathloss/pathloss.py
index 4ad2f551b..9586bc320 100644
--- a/jwst/pathloss/pathloss.py
+++ b/jwst/pathloss/pathloss.py
@@ -976,24 +976,41 @@ def _corrections_for_fixedslit(slit, pathloss, exp_type, source_type):
             wavelength_pointsource *= 1.0e6
             wavelength_uniformsource *= 1.0e6
 
-            wavelength_array = slit.wavelength
-
-            # Compute the point source pathloss 2D correction
-            pathloss_2d_ps = interpolate_onto_grid(
-                wavelength_array,
-                wavelength_pointsource,
-                pathloss_pointsource_vector)
-
-            # Compute the uniform source pathloss 2D correction
-            pathloss_2d_un = interpolate_onto_grid(
-                wavelength_array,
-                wavelength_uniformsource,
-                pathloss_uniform_vector)
-
             # Use the appropriate correction for this slit
             if is_pointsource(source_type or slit.source_type):
+                # calculate the point source corrected wavelengths and uncorrected wavelengths for the slit
+                wavelength_array_corr = get_wavelengths(slit, use_wavecorr=True)
+                wavelength_array_uncorr = get_wavelengths(slit, use_wavecorr=False)
+
+                # Compute the point source pathloss 2D correction
+                pathloss_2d_ps = interpolate_onto_grid(
+                    wavelength_array_corr,
+                    wavelength_pointsource,
+                    pathloss_pointsource_vector)
+
+                # Compute the uniform source pathloss 2D correction
+                pathloss_2d_un = interpolate_onto_grid(
+                    wavelength_array_uncorr,
+                    wavelength_uniformsource,
+                    pathloss_uniform_vector)
+                
                 pathloss_2d = pathloss_2d_ps
+                
             else:
+                wavelength_array = slit.wavelength
+
+                # Compute the point source pathloss 2D correction
+                pathloss_2d_ps = interpolate_onto_grid(
+                    wavelength_array,
+                    wavelength_pointsource,
+                    pathloss_pointsource_vector)
+
+                # Compute the uniform source pathloss 2D correction
+                pathloss_2d_un = interpolate_onto_grid(
+                    wavelength_array,
+                    wavelength_uniformsource,
+                    pathloss_uniform_vector)
+                
                 pathloss_2d = pathloss_2d_un
 
             # Save the corrections. The `data` portion is the correction used.
diff --git a/jwst/wavecorr/wavecorr.py b/jwst/wavecorr/wavecorr.py
index 0174eb1f5..46a92513b 100644
--- a/jwst/wavecorr/wavecorr.py
+++ b/jwst/wavecorr/wavecorr.py
@@ -30,6 +30,10 @@ Both ways give the same result, the current code uses the first one.
 import logging
 import numpy as np
 from gwcs import wcstools
+from gwcs import coordinate_frames as cf
+from astropy import units as u
+from astropy.modeling import tabular
+from astropy.modeling.mappings import Identity
 
 from stdatamodels.jwst import datamodels
 from stdatamodels.jwst.transforms import models as trmodels
@@ -89,16 +93,31 @@ def do_correction(input_model, wavecorr_file):
                         input_model.meta.cal_step.wavecorr = 'SKIPPED'
                         break
                     if _is_point_source(slit, exp_type):
-                        apply_zero_point_correction(slit, wavecorr_file)
-                        output_model.meta.cal_step.wavecorr = 'COMPLETE'
+                        completed = apply_zero_point_correction(slit, wavecorr_file)
+                        if completed:
+                            output_model.meta.cal_step.wavecorr = 'COMPLETE'
+                        else: # pragma: no cover
+                            log.warning(f'Corrections are not invertible for slit {slit.name}')
+                            log.warning('Skipping wavecorr correction')
+                            output_model.meta.cal_step.wavecorr = 'SKIPPED'
+
                         break
 
         # For MOS work on all slits containing a point source
         else:
             for slit in output_model.slits:
                 if _is_point_source(slit, exp_type):
-                    apply_zero_point_correction(slit, wavecorr_file)
-                    output_model.meta.cal_step.wavecorr = 'COMPLETE'
+                    completed = apply_zero_point_correction(slit, wavecorr_file)
+                    if completed:
+                        slit.meta.cal_step.wavecorr = 'COMPLETE'
+                    else: # pragma: no cover
+                        log.warning(f'Corrections are not invertible for slit {slit.name}')
+                        log.warning('Skipping wavecorr correction')
+                        slit.meta.cal_step.wavecorr = 'SKIPPED'
+                else:
+                    slit.meta.cal_step.wavecorr = 'SKIPPED'
+                    
+            output_model.meta.cal_step.wavecorr = 'COMPLETE'
 
     return output_model
 
@@ -112,6 +131,11 @@ def apply_zero_point_correction(slit, reffile):
         Slit data to be corrected.
     reffile : str
         The ``wavecorr`` reference file.
+        
+    Returns
+    -------
+    completed : bool
+        A flag to report whether the zero-point correction was added or skipped.
     """
     log.info(f'slit name {slit.name}')
     slit_wcs = slit.meta.wcs
@@ -127,67 +151,108 @@ def apply_zero_point_correction(slit, reffile):
         # For the MSA the aperture name is "MOS"
         aperture_name = "MOS"
 
-    lam = slit.wavelength.copy()
+    lam = slit.wavelength.copy() * 1e-6
     dispersion = compute_dispersion(slit.meta.wcs)
-    corr, dq_lam = compute_zero_point_correction(lam, reffile, source_xpos,
-                                                 aperture_name, dispersion)
-    # TODO: set a DQ flag to a TBD value for pixels where dq_lam == 0.
-    # The only purpose of dq_lam is to set that flag.
-
-    # Wavelength is in um, the correction is computed in meters.
-    slit.wavelength = slit.wavelength - corr * 10 ** 6
 
-
-def compute_zero_point_correction(lam, freference, source_xpos,
-                                  aperture_name, dispersion):
-    """ Compute the NIRSpec wavelength zero-point correction.
+    wave2wavecorr = calculate_wavelength_correction_transform(lam,
+                                                              dispersion,
+                                                              reffile,
+                                                              source_xpos,
+                                                              aperture_name)
+    # wave2wavecorr should not be None for real data
+    if wave2wavecorr is None: # pragma: no cover
+        completed = False
+        return completed
+    else:        
+        # Make a new frame to insert into the slit wcs object
+        slit_spatial = cf.Frame2D(name='slit_spatial', axes_order=(0, 1), 
+                                  unit=("", ""), axes_names=('x_slit', 'y_slit'))
+        spec = cf.SpectralFrame(name='spectral', axes_order=(2,), unit=(u.micron,),
+                                axes_names=('wavelength',))
+        wcorr_frame = cf.CompositeFrame(
+            [slit_spatial, spec], name='wavecorr_frame')
+        
+        # Insert the new transform into the slit wcs object
+        wave2wavecorr = Identity(2) & wave2wavecorr
+        slit_wcs.insert_frame('slit_frame', wave2wavecorr, wcorr_frame)
+        
+        # Update the stored wavelengths for the slit
+        slit.wavelength = compute_wavelength(slit_wcs)
+        
+        completed = True
+        return completed
+    
+
+def calculate_wavelength_correction_transform(lam, dispersion, freference, 
+                                              source_xpos, aperture_name):
+    """ Generate a WCS transform for the NIRSpec wavelength zero-point correction
+    and add it to the WCS for each slit.
 
     Parameters
     ----------
     lam : ndarray
-        Wavelength array.
+        Wavelength array [in m].
+    dispersion : ndarray
+        The pixel dispersion [in m].
     freference : str
         ``wavecorr`` reference file name.
     source_xpos : float
         X position of the source as a fraction of the slit size.
     aperture_name : str
         Aperture name.
-    dispersion : ndarray
-        The pixel dispersion [in m].
-
+        
     Returns
     -------
-    lambda_corr : ndarray
-        Wavelength correction.
-    lam : ndarray
-        Interpolated wavelengths. Extrapolated values are reset to 0.
-        This is returned so that the DQ array can be updated with a flag
-        which indicates that no zero-point correction was done.
+    model : `~astropy.modeling.tabular.Tabular1D`or None
+        A model which takes wavelength inputs and returns zero-point
+        corrected wavelengths.  Returns None if an invertible model
+        cannot be generated.
     """
+    # Open the zero point reference model
     with datamodels.WaveCorrModel(freference) as wavecorr:
         for ap in wavecorr.apertures:
             if ap.aperture_name == aperture_name:
                 log.info(f'Using wavelength zero-point correction for aperture {ap.aperture_name}')
                 offset_model = ap.zero_point_offset.copy()
-                # TODO: implement variance
-                # variance = ap.variance.copy()
-                # width = ap.width
                 break
         else:
             log.info(f'No wavelength zero-point correction found for slit {aperture_name}')
-
-    deltax = source_xpos
-    lam = lam.copy()
-    lam_no_nans = lam[~np.isnan(lam)]
+        
+    # Set lookup table to extrapolate at bounds to recover wavelengths
+    # beyond model bounds, particularly for the red and blue ends of
+    # prism observations.  fill_value = None sets the lookup tables
+    # to use the default extrapolation which is a linear extrapolation
+    # from scipy.interpolate.interpn
     offset_model.bounds_error = False
-    correction = offset_model(lam_no_nans * 10 ** -6, [deltax] * lam_no_nans.size)
-    lam[~np.isnan(lam)] = correction
-
-    # The correction for pixels outside the slit and wavelengths
-    # outside the wave_range is 0.
-    lam[np.isnan(lam)] = 0.
-    lambda_cor = dispersion * lam
-    return lambda_cor, lam
+    offset_model.fill_value = None
+        
+    # Average the wavelength and dispersion across 2D extracted slit and remove nans
+    # So that we have a 1D wavelength array for building a 1D lookup table wcs transform
+    lam_mean = np.nanmean(lam, axis=0)
+    disp_mean = np.nanmean(dispersion, axis=0)
+    nan_lams = np.isnan(lam_mean) | np.isnan(disp_mean)
+    lam_mean = lam_mean[~nan_lams]
+    disp_mean = disp_mean[~nan_lams]
+    
+    # Calculate the corrected wavelengths
+    pixel_corrections = offset_model(lam_mean, source_xpos)
+    lam_corrected = lam_mean + (pixel_corrections * disp_mean)
+    
+    # Check to make sure that the corrected wavelengths are monotonically increasing
+    if np.all(np.diff(lam_corrected) > 0):
+        # monotonically increasing        
+        # Build a look up table to transform between corrected and uncorrected wavelengths
+        wave2wavecorr = tabular.Tabular1D(points=lam_mean, 
+                                          lookup_table=lam_corrected, 
+                                          bounds_error=False, 
+                                          fill_value=None, 
+                                          name='wave2wavecorr')
+    
+        return wave2wavecorr
+    
+    else:
+        # output wavelengths are not monotonically increasing
+        return None
 
 
 def compute_dispersion(wcs, xpix=None, ypix=None):
@@ -218,6 +283,32 @@ def compute_dispersion(wcs, xpix=None, ypix=None):
     return (lamright - lamleft) * 10 ** -6
 
 
+def compute_wavelength(wcs, xpix=None, ypix=None):
+    """ Compute the pixel wavelength.
+
+    Parameters
+    ----------
+    wcs : `~gwcs.wcs.WCS`
+        The WCS object for this slit.
+    xpix : ndarray, float, optional
+    ypix : ndarray, float, optional
+        Compute the wavelength at the x, y pixels.
+        If not provided the dispersion is computed on a
+        grid based on ``wcs.bounding_box``.
+
+    Returns
+    -------
+    wavelength : ndarray
+        The wavelength [in microns].
+
+    """
+    if xpix is None or ypix is None:
+        xpix, ypix = wcstools.grid_from_bounding_box(wcs.bounding_box, step=(1, 1))
+        
+    _, _, lam = wcs(xpix, ypix)
+    return lam
+
+
 def _is_point_source(slit, exp_type):
     """
     Determine if a source is a point source.


Test patch:
diff --git a/jwst/lib/tests/test_wcs_utils.py b/jwst/lib/tests/test_wcs_utils.py
new file mode 100644
index 000000000..0cd731e77
--- /dev/null
+++ b/jwst/lib/tests/test_wcs_utils.py
@@ -0,0 +1,157 @@
+import numpy as np
+from numpy.testing import assert_allclose
+
+from astropy import units as u
+from astropy import coordinates as coord
+from astropy.modeling.models import Mapping, Identity, Shift, Scale
+from gwcs import wcstools, wcs
+from gwcs import coordinate_frames as cf
+
+from stdatamodels.jwst import datamodels
+from stdatamodels.jwst.transforms.models import NirissSOSSModel
+from jwst.lib.wcs_utils import get_wavelengths
+from jwst.assign_wcs import util
+
+
+def create_model():
+
+    det = cf.Frame2D(name="detector", axes_order=(0, 1))
+
+    sky = cf.CelestialFrame(name="sky", axes_order=(0, 1), reference_frame=coord.ICRS())
+    slit_spatial = cf.Frame2D(
+        name="slit_spatial",
+        axes_order=(0, 1),
+        unit=("", ""),
+        axes_names=("x_slit", "y_slit"),
+    )
+
+    spec = cf.SpectralFrame(
+        name="spectral", axes_order=(2,), unit=(u.micron,), axes_names=("wavelength",)
+    )
+    slit_frame = cf.CompositeFrame([slit_spatial, spec], name="slit_frame")
+    world = cf.CompositeFrame([sky, spec], name="world")
+
+    det2slit = Mapping((0, 1, 1)) | (Identity(2) & (Scale(0.5) | Shift(0.5)))
+    slit2sky = Identity(3)
+
+    slit_wcs = wcs.WCS([(det, det2slit), (slit_frame, slit2sky), (world, None)])
+
+    # compute wavelengths
+
+    data = np.full((10, 10), fill_value=5.0)
+
+    bounding_box = util.wcs_bbox_from_shape(data.shape)
+
+    x, y = wcstools.grid_from_bounding_box(bounding_box, step=(1, 1))
+    _, _, lam = slit_wcs(x, y)
+    lam = lam.astype(np.float32)
+    model = datamodels.SlitModel(data=data, wavelength=lam)
+    model.meta.wcs = slit_wcs
+
+    return model
+
+
+def create_mock_wl():
+    wl = np.arange(10.0)
+    wl = wl[:, np.newaxis]
+    wl = np.repeat(wl, 10, axis=1)
+    wl = (wl * 0.5) + 0.5
+    return wl
+
+
+def test_get_wavelengths():
+
+    # create a mock SlitModel
+    model = create_model()
+
+    # calculate what the wavelength array should be
+    wl_og = create_mock_wl()
+
+    # Test that the get wavelengths returns the wavelength grid
+    wl = get_wavelengths(model)
+    assert_allclose(wl, wl_og)
+
+    del model.wavelength
+
+    # Check that wavelengths can be generated from wcs when the
+    # wavelength attribute is unavailable
+    wl = get_wavelengths(model)
+    assert_allclose(wl, wl_og)
+
+    # Check that wavelengths are generated correctly when given a WFSS exp_type
+    wl = get_wavelengths(model, exp_type="NRC_TSGRISM")
+    assert_allclose(wl, wl_og)
+
+
+def test_get_wavelengths_soss():
+
+    # create a mock SlitModel
+    model = create_model()
+
+    del model.wavelength
+    model.meta.exposure.type = "NIS_SOSS"
+
+    wcs = model.meta.wcs
+    new_wcs = NirissSOSSModel(
+        [
+            1,
+        ],
+        [
+            wcs,
+        ],
+    )
+    model.meta.wcs = new_wcs
+
+    # calculate what the wavelength array should be
+    wl_og = create_mock_wl()
+
+    wl = get_wavelengths(model, order=1)
+    assert_allclose(wl, wl_og)
+
+
+def test_get_wavelength_wavecorr():
+
+    # create a mock SlitModel
+    model = create_model()
+
+    wl_og = create_mock_wl()
+
+    # Test use_wavecorr with no wavelength correction modificiation
+    # get_wavelengths should return the same wavelengths for use_wavecorr
+    # True and False
+
+    wl_corr = get_wavelengths(model, use_wavecorr=True)
+    assert_allclose(wl_corr, wl_og)
+
+    wl_uncorr = get_wavelengths(model, use_wavecorr=False)
+    assert_allclose(wl_corr, wl_uncorr)
+
+    # Update the model wcs to add a wavelength corrected slit frame
+    slit_spatial = cf.Frame2D(
+        name="slit_spatial",
+        axes_order=(0, 1),
+        unit=("", ""),
+        axes_names=("x_slit", "y_slit"),
+    )
+    spec = cf.SpectralFrame(
+        name="spectral", axes_order=(2,), unit=(u.micron,), axes_names=("wavelength",)
+    )
+    wcorr_frame = cf.CompositeFrame([slit_spatial, spec], name="wavecorr_frame")
+
+    # Insert the new transform into the slit wcs object
+    wave2wavecorr = Identity(2) & Shift(0.1)
+    model.meta.wcs.insert_frame("slit_frame", wave2wavecorr, wcorr_frame)
+
+    bounding_box = util.wcs_bbox_from_shape(model.data.shape)
+    x, y = wcstools.grid_from_bounding_box(bounding_box, step=(1, 1))
+    _, _, lam = model.meta.wcs(x, y)
+    model.wavelength = lam
+
+    # calculate what the corrected wavelength array should be
+    wl_corr_og = wl_og + 0.1
+
+    wl_corr = get_wavelengths(model, use_wavecorr=True)
+    assert_allclose(wl_corr, wl_corr_og)
+
+    wl_uncorr = get_wavelengths(model, use_wavecorr=False)
+    assert_allclose(wl_uncorr, wl_og)
diff --git a/jwst/wavecorr/tests/test_wavecorr.py b/jwst/wavecorr/tests/test_wavecorr.py
index e4a06cb62..5ff61fc24 100644
--- a/jwst/wavecorr/tests/test_wavecorr.py
+++ b/jwst/wavecorr/tests/test_wavecorr.py
@@ -27,32 +27,60 @@ def test_wavecorr():
     im = datamodels.ImageModel(hdul)
     im_wcs = AssignWcsStep.call(im)
     im_ex2d = Extract2dStep.call(im_wcs)
-    im_ex2d.slits[0].meta.wcs.bounding_box = ((-.5, 1432.5), (-.5, 37.5))
+    bbox = ((-.5, 1432.5), (-.5, 37.5))
+    im_ex2d.slits[0].meta.wcs.bounding_box = bbox
+    x, y = wcstools.grid_from_bounding_box(bbox)
+    ra, dec, lam_before = im_ex2d.slits[0].meta.wcs(x, y)
+    im_ex2d.slits[0].wavelength = lam_before
     im_src = SourceTypeStep.call(im_ex2d)
+
+    # the mock msa source is an extended source, change to point for testing
+    im_src.slits[0].source_type = 'POINT'
     im_wave = WavecorrStep.call(im_src)
 
     # test dispersion is of the correct order
     # there's one slit only
     slit = im_src.slits[0]
-    x, y = wcstools.grid_from_bounding_box(slit.meta.wcs.bounding_box)
     dispersion = wavecorr.compute_dispersion(slit.meta.wcs, x, y)
     assert_allclose(dispersion[~np.isnan(dispersion)], 1e-9, atol=1e-10)
-
-    # the difference in wavelength should be of the order of e-10 in um
-    assert_allclose(im_src.slits[0].wavelength - im_wave.slits[0].wavelength, 1e-10)
-
+    
+    # test that the wavelength is on the order of microns
+    wavelength = wavecorr.compute_wavelength(slit.meta.wcs, x, y)
+    assert_allclose(np.nanmean(wavelength), 2.5, atol=0.1)
+        
+    # Check that the stored wavelengths were corrected
+    abs_wave_correction = np.abs(im_src.slits[0].wavelength - im_wave.slits[0].wavelength)
+    assert_allclose(np.nanmean(abs_wave_correction), 0.00046, atol=0.0001)
+    
+    # Check that the slit wcs has been updated to provide corrected wavelengths
+    corrected_wavelength = wavecorr.compute_wavelength(im_wave.slits[0].meta.wcs, x, y)
+    assert_allclose(im_wave.slits[0].wavelength, corrected_wavelength)
+
+    # test the round-tripping on the wavelength correction transform
+    ref_name = im_wave.meta.ref_file.wavecorr.name
+    freference = datamodels.WaveCorrModel(
+        WavecorrStep.reference_uri_to_cache_path(ref_name, im.crds_observatory))
+
+    lam_uncorr = lam_before * 1e-6
+    wave2wavecorr = wavecorr.calculate_wavelength_correction_transform(
+        lam_uncorr, dispersion, freference, slit.source_xpos, 'MOS')
+    lam_corr = wave2wavecorr(lam_uncorr)
+    assert_allclose(lam_uncorr, wave2wavecorr.inverse(lam_corr))
+    
     # test on both sides of the shutter
     source_xpos1 = -.2
     source_xpos2 = .2
 
-    ra, dec, lam = slit.meta.wcs(x, y)
-    ref_name = im_wave.meta.ref_file.wavecorr.name
-    freference = datamodels.WaveCorrModel(WavecorrStep.reference_uri_to_cache_path(ref_name, im.crds_observatory))
-    zero_point1 = wavecorr.compute_zero_point_correction(lam, freference, source_xpos1, 'MOS', dispersion)
-    zero_point2 = wavecorr.compute_zero_point_correction(lam, freference, source_xpos2, 'MOS', dispersion)
-    diff_correction = np.abs(zero_point1[1] - zero_point2[1])
-    non_zero = diff_correction[diff_correction != 0]
-    assert_allclose(np.nanmean(non_zero), 0.75, atol=0.01)
+    wave_transform1 = wavecorr.calculate_wavelength_correction_transform(
+        lam_uncorr, dispersion, freference, source_xpos1, 'MOS')
+    wave_transform2 = wavecorr.calculate_wavelength_correction_transform(
+        lam_uncorr, dispersion, freference, source_xpos2, 'MOS')
+
+    zero_point1 = wave_transform1(lam_uncorr)
+    zero_point2 = wave_transform2(lam_uncorr)
+
+    diff_correction = np.abs(zero_point1 - zero_point2)
+    assert_allclose(np.nanmean(diff_correction), 8.0e-10, atol=0.1e-10)
 
 
 def test_ideal_to_v23_fs():
@@ -119,7 +147,64 @@ def test_skipped():
     outw = WavecorrStep.call(outs)
 
     source_pos = (0.004938526981283373, -0.02795306204991911)
-    assert_allclose((outw.slits[ind].source_xpos, outw.slits[ind].source_ypos), source_pos)
+    assert_allclose((outw.slits[ind].source_xpos, outw.slits[ind].source_ypos),source_pos)
+
+    # Test if the corrected wavelengths are not monotonically increasing
+
+    # This case is not expected with real data, test that no correction 
+    # transform is returned (a skip criterion) with simple case 
+    # of flipped wavelength solutions, which produces a monotonically 
+    # decreasing wavelengths
+    lam = np.tile(np.flip(np.arange(0.6, 5.5, 0.01)*1e-6), (22, 1))
+    disp = np.tile(np.full(lam.shape[-1], -0.01)*1e-6, (22, 1))
+
+    ref_name = outw.meta.ref_file.wavecorr.name
+    reffile = datamodels.WaveCorrModel(
+        WavecorrStep.reference_uri_to_cache_path(ref_name, im.crds_observatory))
+    source_xpos = 0.1
+    aperture_name = 'S200A1'
+    
+    transform = wavecorr.calculate_wavelength_correction_transform(
+        lam, disp, reffile, source_xpos, aperture_name)
+    assert transform is None
+    
+    
+def test_mos_slit_status():
+    """ Test conditions that are skipped for mos slitlets."""
+
+    hdul = create_nirspec_mos_file()
+    msa_meta = os.path.join(jwst.__path__[0], *['assign_wcs', 'tests', 'data', 'msa_configuration.fits'])
+    hdul[0].header['MSAMETFL'] = msa_meta
+    hdul[0].header['MSAMETID'] = 12
+    im = datamodels.ImageModel(hdul)
+    im_wcs = AssignWcsStep.call(im)
+    im_ex2d = Extract2dStep.call(im_wcs)
+    bbox = ((-.5, 1432.5), (-.5, 37.5))
+    im_ex2d.slits[0].meta.wcs.bounding_box = bbox
+    x, y = wcstools.grid_from_bounding_box(bbox)
+    ra, dec, lam_before = im_ex2d.slits[0].meta.wcs(x, y)
+    im_ex2d.slits[0].wavelength = lam_before
+    im_src = SourceTypeStep.call(im_ex2d)
+
+    # test the mock msa source as an extended source
+    im_src.slits[0].source_type = 'EXTENDED'
+    im_wave = WavecorrStep.call(im_src)
+    
+    # check that the step is recorded as completed
+    assert im_wave.meta.cal_step.wavecorr == 'COMPLETE'
+    
+    # check that the step is listed as skipped for extended mos sources
+    assert im_wave.slits[0].meta.cal_step.wavecorr == 'SKIPPED'
+    
+    # test the mock msa source as a point source
+    im_src.slits[0].source_type = 'POINT'
+    im_wave = WavecorrStep.call(im_src)
+    
+    # check that the step is recorded as completed
+    assert im_wave.meta.cal_step.wavecorr == 'COMPLETE'
+    
+    # check that the step is listed as complete for mos point sources
+    assert im_wave.slits[0].meta.cal_step.wavecorr == 'COMPLETE'
 
 
 def test_wavecorr_fs():
@@ -171,13 +256,31 @@ def test_wavecorr_fs():
     dispersion = wavecorr.compute_dispersion(slit.meta.wcs, x, y)
     assert_allclose(dispersion[~np.isnan(dispersion)], 1e-8, atol=1.04e-8)
 
+    # Check that the slit wavelengths are consistent with the slit wcs
+    corrected_wavelength = wavecorr.compute_wavelength(slit.meta.wcs, x, y)
+    assert_allclose(slit.wavelength, corrected_wavelength)
+
+    # test the roundtripping on the wavelength correction transform
+    ref_name = result.meta.ref_file.wavecorr.name
+    freference = datamodels.WaveCorrModel(WavecorrStep.reference_uri_to_cache_path(ref_name, im.crds_observatory))
+
+    lam_uncorr = lam_before * 1e-6
+    wave2wavecorr = wavecorr.calculate_wavelength_correction_transform(lam_uncorr, dispersion, 
+                                                                       freference, slit.source_xpos, 'S200A1')
+    lam_corr = wave2wavecorr(lam_uncorr)
+    assert_allclose(lam_uncorr, wave2wavecorr.inverse(lam_corr))
+
     # test on both sides of the slit center
     source_xpos1 = -.2
     source_xpos2 = .2
 
-    ref_name = result.meta.ref_file.wavecorr.name
-    freference = datamodels.WaveCorrModel(WavecorrStep.reference_uri_to_cache_path(ref_name, im.crds_observatory))
-    zero_point1 = wavecorr.compute_zero_point_correction(lam_before, freference, source_xpos1, 'S200A1', dispersion)
-    zero_point2 = wavecorr.compute_zero_point_correction(lam_before, freference, source_xpos2, 'S200A1', dispersion)
-    diff_correction = np.abs(zero_point1[1] - zero_point2[1])
-    assert_allclose(np.nanmean(diff_correction), 0.45, atol=0.01)
+    wave_transform1 = wavecorr.calculate_wavelength_correction_transform(
+        lam_uncorr, dispersion, freference, source_xpos1, 'S200A1')
+    wave_transform2 = wavecorr.calculate_wavelength_correction_transform(
+        lam_uncorr, dispersion, freference, source_xpos2, 'S200A1')
+
+    zero_point1 = wave_transform1(lam_uncorr)
+    zero_point2 = wave_transform2(lam_uncorr)
+
+    diff_correction = np.abs(zero_point1 - zero_point2)
+    assert_allclose(np.nanmean(diff_correction), 6.3e-9, atol=0.1e-9)