Repository: astropy/specutils

Issue:
JWST s3d loader broken for latest cube data
The specutils loader for JWST s3d files is broken for recent NIRSpec IFU cube data cal_ver 1.4.6.   It gives the error `ValueError: Spectral axis length (973) must be the same size or one greater (if specifying bin edges) than that of the last flux axis (950)`.  I think there's a mismatch between the input data and the loader code to extract WCS and convert it to a wavelength. 

To reproduce:
```
data = "jw01409-o018_t011_nirspec_prism-clear_s3d.fits"
ss = Spectrum1D.read(data)
```
 
Traceback:
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [24], in <cell line: 1>()
----> 1 ss=Spectrum1D.read(hdu.filename())

File ~/anaconda3/envs/vmpy3.9/lib/python3.9/site-packages/astropy/nddata/mixins/ndio.py:59, in NDDataRead.__call__(self, *args, **kwargs)
     58 def __call__(self, *args, **kwargs):
---> 59     return self.registry.read(self._cls, *args, **kwargs)

File ~/anaconda3/envs/vmpy3.9/lib/python3.9/site-packages/astropy/io/registry/core.py:199, in UnifiedInputRegistry.read(self, cls, format, cache, *args, **kwargs)
    195     format = self._get_valid_format(
    196         'read', cls, path, fileobj, args, kwargs)
    198 reader = self.get_reader(format, cls)
--> 199 data = reader(*args, **kwargs)
    201 if not isinstance(data, cls):
    202     # User has read with a subclass where only the parent class is
    203     # registered.  This returns the parent class, so try coercing
    204     # to desired subclass.
    205     try:

File ~/anaconda3/envs/vmpy3.9/lib/python3.9/site-packages/specutils/io/default_loaders/jwst_reader.py:531, in jwst_s3d_single_loader(filename, **kwargs)
    513 @data_loader(
    514     "JWST s3d", identifier=identify_jwst_s3d_fits, dtype=Spectrum1D,
    515     extensions=['fits'], priority=10,
    516 )
    517 def jwst_s3d_single_loader(filename, **kwargs):
    518     """
    519     Loader for JWST s3d 3D rectified spectral data in FITS format.
    520
   (...)
    529         The spectrum contained in the file.
    530     """
--> 531     spectrum_list = _jwst_s3d_loader(filename, **kwargs)
    532     if len(spectrum_list) == 1:
    533         return spectrum_list[0]

File ~/anaconda3/envs/vmpy3.9/lib/python3.9/site-packages/specutils/io/default_loaders/jwst_reader.py:618, in _jwst_s3d_loader(filename, **kwargs)
    615         mask_name = primary_header.get("MASKEXT", "DQ")
    616         mask = hdulist[mask_name].data.T
--> 618         spec = Spectrum1D(flux=flux, spectral_axis=wavelength, meta=meta,
    619                           uncertainty=err, mask=mask)
    620         spectra.append(spec)
    622 return SpectrumList(spectra)

File ~/anaconda3/envs/vmpy3.9/lib/python3.9/site-packages/specutils/spectra/spectrum1d.py:178, in Spectrum1D.__init__(self, flux, spectral_axis, wcs, velocity_convention, rest_value, redshift, radial_velocity, bin_specification, **kwargs)
    176         bin_specification = "edges"
    177     else:
--> 178         raise ValueError(
    179             "Spectral axis length ({}) must be the same size or one "
    180             "greater (if specifying bin edges) than that of the last "
    181             "flux axis ({})".format(spectral_axis.shape[0],
    182                                     flux.shape[-1]))
    184 # If a WCS is provided, check that the spectral axis is last and reorder
    185 # the arrays if not
    186 if wcs is not None and hasattr(wcs, "naxis"):

ValueError: Spectral axis length (973) must be the same size or one greater (if specifying bin edges) than that of the last flux axis (950)
```

Comments:

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 026afba5..6a81e107 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -10,6 +10,8 @@ Bug Fixes
 - Fixed a bug with moment map orders greater than 1 not being able to handle
   cubes with non-square spatial dimensions. [#970]
 
+- Added a workaround for reading JWST IFUs with incorrect GWCS. [#973]
+
 Other Changes and Additions
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/specutils/analysis/width.py b/specutils/analysis/width.py
index 5d6bf897..609081fa 100644
--- a/specutils/analysis/width.py
+++ b/specutils/analysis/width.py
@@ -4,9 +4,11 @@ spectral features.
 """
 
 import numpy as np
+from astropy.nddata import StdDevUncertainty
 from astropy.stats.funcs import gaussian_sigma_to_fwhm
 from ..manipulation import extract_region
 from . import centroid
+from .uncertainty import _convert_uncertainty
 from .utils import computation_wrapper
 from scipy.signal import find_peaks, peak_widths
 
@@ -186,8 +188,15 @@ def _compute_gaussian_fwhm(spectrum, regions=None):
     This is a helper function for the above `gaussian_fwhm()` method.
     """
 
-    fwhm = _compute_gaussian_sigma_width(spectrum, regions) * gaussian_sigma_to_fwhm
+    sigma = _compute_gaussian_sigma_width(spectrum, regions)
 
+    fwhm = sigma * gaussian_sigma_to_fwhm
+    if sigma.uncertainty is not None:
+        fwhm.uncertainty = sigma.uncertainty * gaussian_sigma_to_fwhm
+    else:
+        fwhm.uncertainty = None
+
+    fwhm.uncertainty_type = 'std'
     return fwhm
 
 
@@ -201,9 +210,16 @@ def _compute_gaussian_sigma_width(spectrum, regions=None):
     else:
         calc_spectrum = spectrum
 
+    if spectrum.uncertainty is not None:
+        flux_uncert = _convert_uncertainty(calc_spectrum.uncertainty, StdDevUncertainty)
+    else:
+        # dummy value for uncertainties to avoid extra if-statements when applying mask
+        flux_uncert = np.zeros_like(calc_spectrum.flux)
+
     if hasattr(spectrum, 'mask') and spectrum.mask is not None:
         flux = calc_spectrum.flux[~spectrum.mask]
         spectral_axis = calc_spectrum.spectral_axis[~spectrum.mask]
+        flux_uncert = flux_uncert[~calc_spectrum.mask]
     else:
         flux = calc_spectrum.flux
         spectral_axis = calc_spectrum.spectral_axis
@@ -212,11 +228,36 @@ def _compute_gaussian_sigma_width(spectrum, regions=None):
 
     if flux.ndim > 1:
         spectral_axis = np.broadcast_to(spectral_axis, flux.shape, subok=True)
+        centroid_result_uncert = centroid_result.uncertainty
         centroid_result = centroid_result[:, np.newaxis]
+        centroid_result.uncertainty = centroid_result_uncert[:, np.newaxis] if centroid_result_uncert is not None else None  # noqa
 
     dx = (spectral_axis - centroid_result)
-    sigma = np.sqrt(np.sum((dx * dx) * flux, axis=-1) / np.sum(flux, axis=-1))
+    numerator = np.sum((dx * dx) * flux, axis=-1)
+    denom = np.sum(flux, axis=-1)
+    sigma2 = numerator / denom
+    sigma = np.sqrt(sigma2)
+    if centroid_result.uncertainty is not None:
+        # NOTE: until/unless disp_uncert is supported, dx_uncert == centroid_result.uncertainty
+        disp_uncert = 0.0 * spectral_axis.unit
+        dx_uncert = np.sqrt(disp_uncert**2 + centroid_result.uncertainty**2)
+        # dx_uncert = centroid_result.uncertainty
+
+        # uncertainty for each term in the numerator sum
+        num_term_uncerts = dx * dx * flux * np.sqrt(2*(dx_uncert/dx)**2 + (flux_uncert/flux)**2)
+        # uncertainty (squared) for the numerator, added in quadrature
+        num_uncertsq = np.sum(num_term_uncerts**2, axis=-1)
+        # uncertainty (squared) for the denomenator
+        denom_uncertsq = np.sum(flux_uncert**2)
+
+        sigma2_uncert = numerator/denom * np.sqrt(num_uncertsq * numerator**-2 +
+                                                  denom_uncertsq * denom**-2)
+
+        sigma.uncertainty = 0.5 * sigma2_uncert / sigma2 * sigma.unit
+    else:
+        sigma.uncertainty = None
 
+    sigma.uncertainty_type = 'std'
     return sigma
 
 
diff --git a/specutils/io/default_loaders/jwst_reader.py b/specutils/io/default_loaders/jwst_reader.py
index 8f9ee5f5..9bb895f8 100644
--- a/specutils/io/default_loaders/jwst_reader.py
+++ b/specutils/io/default_loaders/jwst_reader.py
@@ -7,6 +7,8 @@ from astropy.units import Quantity
 from astropy.table import Table
 from astropy.io import fits
 from astropy.nddata import StdDevUncertainty, VarianceUncertainty, InverseVariance
+from astropy.time import Time
+from astropy.wcs import WCS
 import asdf
 from gwcs.wcstools import grid_from_bounding_box
 
@@ -585,6 +587,24 @@ def _jwst_s3d_loader(filename, **kwargs):
 
             wavelength = Quantity(wavelength_array, unit=wavelength_unit)
 
+            # The GWCS is currently broken for some IFUs, here we work around that
+            wcs = None
+            if wavelength.shape[0] != flux.shape[-1]:
+                # Need MJD-OBS for this workaround
+                if 'MJD-OBS' not in hdu.header:
+                    for key in ('MJD-BEG', 'DATE-OBS'):  # Possible alternatives
+                        if key in hdu.header:
+                            if key.startswith('MJD'):
+                                hdu.header['MJD-OBS'] = hdu.header[key]
+                                break
+                            else:
+                                t = Time(hdu.header[key])
+                                hdu.header['MJD-OBS'] = t.mjd
+                                break
+                wcs = WCS(hdu.header)
+                # Swap to match the flux transpose
+                wcs = wcs.swapaxes(-1, 0)
+
             # Merge primary and slit headers and dump into meta
             slit_header = hdu.header
             header = primary_header.copy()
@@ -615,8 +635,11 @@ def _jwst_s3d_loader(filename, **kwargs):
             mask_name = primary_header.get("MASKEXT", "DQ")
             mask = hdulist[mask_name].data.T
 
-            spec = Spectrum1D(flux=flux, spectral_axis=wavelength, meta=meta,
-                              uncertainty=err, mask=mask)
+            if wcs is not None:
+                spec = Spectrum1D(flux=flux, wcs=wcs, meta=meta, uncertainty=err, mask=mask)
+            else:
+                spec = Spectrum1D(flux=flux, spectral_axis=wavelength, meta=meta,
+                                  uncertainty=err, mask=mask)
             spectra.append(spec)
 
     return SpectrumList(spectra)


Test patch:
diff --git a/specutils/tests/test_analysis.py b/specutils/tests/test_analysis.py
index 64561d90..81824440 100644
--- a/specutils/tests/test_analysis.py
+++ b/specutils/tests/test_analysis.py
@@ -469,6 +469,7 @@ def test_centroid(simulated_spectra):
     assert isinstance(spec_centroid, u.Quantity)
     assert np.allclose(spec_centroid.value, spec_centroid_expected.value)
     assert hasattr(spec_centroid, 'uncertainty')
+    # NOTE: value has not been scientifically validated
     assert quantity_allclose(spec_centroid.uncertainty, 3.91834165e-06*u.um, rtol=5e-5)
 
 
@@ -557,14 +558,19 @@ def test_gaussian_sigma_width():
 
     # Create a (centered) gaussian spectrum for testing
     mean = 5
-    frequencies = np.linspace(0, mean*2, 100) * u.GHz
+    frequencies = np.linspace(0, mean*2, 101)[1:] * u.GHz
     g1 = models.Gaussian1D(amplitude=5*u.Jy, mean=mean*u.GHz, stddev=0.8*u.GHz)
 
     spectrum = Spectrum1D(spectral_axis=frequencies, flux=g1(frequencies))
+    uncertainty = StdDevUncertainty(0.1*np.random.random(len(spectrum.flux))*u.mJy)
+    spectrum.uncertainty = uncertainty
 
     result = gaussian_sigma_width(spectrum)
 
     assert quantity_allclose(result, g1.stddev, atol=0.01*u.GHz)
+    assert hasattr(result, 'uncertainty')
+    # NOTE: value has not been scientifically validated!
+    assert quantity_allclose(result.uncertainty, 4.8190546890398186e-05*u.GHz, rtol=5e-5)
 
 
 def test_gaussian_sigma_width_masked():
@@ -573,7 +579,7 @@ def test_gaussian_sigma_width_masked():
 
     # Create a (centered) gaussian masked spectrum for testing
     mean = 5
-    frequencies = np.linspace(0, mean*2, 100) * u.GHz
+    frequencies = np.linspace(0, mean*2, 101)[1:] * u.GHz
     g1 = models.Gaussian1D(amplitude=5*u.Jy, mean=mean*u.GHz, stddev=0.8*u.GHz)
     uncertainty = StdDevUncertainty(0.1*np.random.random(len(frequencies))*u.Jy)
 
@@ -585,13 +591,16 @@ def test_gaussian_sigma_width_masked():
     result = gaussian_sigma_width(spectrum)
 
     assert quantity_allclose(result, g1.stddev, atol=0.01*u.GHz)
+    assert hasattr(result, 'uncertainty')
+    # NOTE: value has not been scientifically validated!
+    assert quantity_allclose(result.uncertainty, 0.06852821940808544*u.GHz, rtol=5e-5)
 
 
 def test_gaussian_sigma_width_regions():
 
     np.random.seed(42)
 
-    frequencies = np.linspace(100, 0, 10000) * u.GHz
+    frequencies = np.linspace(100, 0, 10000)[:-1] * u.GHz
     g1 = models.Gaussian1D(amplitude=5*u.Jy, mean=10*u.GHz, stddev=0.8*u.GHz)
     g2 = models.Gaussian1D(amplitude=5*u.Jy, mean=2*u.GHz, stddev=0.3*u.GHz)
     g3 = models.Gaussian1D(amplitude=5*u.Jy, mean=70*u.GHz, stddev=10*u.GHz)
@@ -654,15 +663,20 @@ def test_gaussian_fwhm():
 
     # Create a (centered) gaussian spectrum for testing
     mean = 5
-    frequencies = np.linspace(0, mean*2, 100) * u.GHz
+    frequencies = np.linspace(0, mean*2, 101)[1:] * u.GHz
     g1 = models.Gaussian1D(amplitude=5*u.Jy, mean=mean*u.GHz, stddev=0.8*u.GHz)
 
     spectrum = Spectrum1D(spectral_axis=frequencies, flux=g1(frequencies))
+    uncertainty = StdDevUncertainty(0.1*np.random.random(len(spectrum.flux))*u.mJy)
+    spectrum.uncertainty = uncertainty
 
     result = gaussian_fwhm(spectrum)
 
     expected = g1.stddev * gaussian_sigma_to_fwhm
     assert quantity_allclose(result, expected, atol=0.01*u.GHz)
+    assert hasattr(result, 'uncertainty')
+    # NOTE: value has not been scientifically validated!
+    assert quantity_allclose(result.uncertainty, 0.00011348006579851353*u.GHz, rtol=5e-5)
 
 
 def test_gaussian_fwhm_masked():
@@ -684,6 +698,9 @@ def test_gaussian_fwhm_masked():
 
     expected = g1.stddev * gaussian_sigma_to_fwhm
     assert quantity_allclose(result, expected, atol=0.01*u.GHz)
+    assert hasattr(result, 'uncertainty')
+    # NOTE: value has not been scientifically validated!
+    assert quantity_allclose(result.uncertainty, 0.16688079501948674*u.GHz, rtol=5e-5)
 
 
 @pytest.mark.parametrize('mean', range(3, 8))
@@ -795,7 +812,6 @@ def test_fwhm():
 
     spectrum = Spectrum1D(spectral_axis=wavelengths, flux=flux)
     result = fwhm(spectrum)
-
     assert quantity_allclose(result, 1.01 * u.um)
 
     # Highest point at the last point