WavelengthSolution1D.resample bug fixes and improvements

CHANGES.rst

@@ -30,6 +30,26 @@ New Features
   (``AWAV-GRA``) and vacuum (``WAVE-GRI``) spectral axis types in FITS WCS
   export. [#316]
 
+Bug Fixes
+^^^^^^^^^
+
+- Fixed ``WavelengthSolution1D.resample()``, which multiplied the per-pixel flux by
+  the wavelength width of each pixel before dividing by the bin width. The output was
+  labelled as a flux density per wavelength unit but was numerically off by the local
+  dispersion, so a flat spectrum in counts per pixel stayed flat instead of falling where
+  the dispersion grows, and the integrated flux was not conserved.  [#XXX]
+
+- ``WavelengthSolution1D.resample()`` now bins spectra whose spectral axis does not start
+  at pixel zero correctly, supports solutions where the wavelength decreases with pixel
+  number (the output is always ascending in wavelength), propagates the input mask by
+  flagging every bin that received a contribution from a masked pixel, copies the input
+  metadata, returns no uncertainty instead of a fabricated zero-valued one when the
+  input has none, and rejects ``nbins=0`` as its error message already promised. [#XXX]
+
+- The spectrum returned by ``WavelengthSolution1D.resample()`` now carries the bin edges
+  actually used on its spectral axis, so ``spectral_axis.bin_edges`` is exact for
+  non-uniform ``bin_edges`` grids instead of being inferred from the bin centres. [#XXX]
+
 API Changes
 ^^^^^^^^^^^
 

specreduce/tests/test_wavesol1d.py

@@ -5,7 +5,7 @@
 from astropy.io import fits
 from astropy.modeling import models, fitting
 from astropy.modeling.polynomial import Polynomial1D
-from astropy.nddata import StdDevUncertainty
+from astropy.nddata import StdDevUncertainty, VarianceUncertainty
 from gwcs import coordinate_frames, wcs
 from astropy.utils.exceptions import AstropyUserWarning
 from astropy.wcs import WCS as astropy_WCS
@@ -114,18 +114,20 @@ def test_resample(mk_spectrum, mk_ws_with_transform, mk_ws_without_transform):
     assert len(resampled.flux) == 50
     assert resampled.flux.unit == u.count / u.angstrom
 
-    pix_edges = np.arange(spectrum.spectral_axis.size + 1) - 0.5
-    f0 = (spectrum.flux.value * np.diff(ws._p2w(pix_edges))).sum()
+    # the input is in counts per pixel, so integrating the output density over the
+    # wavelength bins must return the total input counts, up to the accuracy of the
+    # interpolated wavelength-to-pixel inverse at the outermost bin edges
+    f0 = spectrum.flux.value.sum()
     f1 = (resampled.flux.value * np.diff(resampled.spectral_axis.value)[0]).sum()
-    np.testing.assert_approx_equal(f0, f1, 5)
+    np.testing.assert_allclose(f0, f1, rtol=1e-4)
 
     resampled = ws.resample(spectrum, wlbounds=wav_bounds)
     resampled = ws.resample(spectrum, bin_edges=np.linspace(*wav_bounds, num=50))
 
     # Resample a spectrum without uncertainty
     spectrum.uncertainty = None
     resampled = ws.resample(spectrum, nbins=50)
-    assert resampled.uncertainty is not None
+    assert resampled.uncertainty is None
 
     ws = mk_ws_without_transform
     with pytest.raises(ValueError, match="Wavelength solution not set."):
@@ -134,6 +136,8 @@ def test_resample(mk_spectrum, mk_ws_with_transform, mk_ws_without_transform):
     ws = mk_ws_with_transform
     with pytest.raises(ValueError, match="Number of bins must be non-zero and positive"):
         ws.resample(mk_spectrum, nbins=-5)
+    with pytest.raises(ValueError, match="Number of bins must be non-zero and positive"):
+        ws.resample(mk_spectrum, nbins=0)
 
 
 def test_pix_to_wav(mk_ws_with_transform):
@@ -419,3 +423,164 @@ def test_from_asdf_raises_on_unsupported_transform(tmp_path):
     asdf.AsdfFile({"wavelength_solution": {"gwcs": w, "wave_air": False}}).write_to(path)
     with pytest.raises(ValueError, match="shift followed by a polynomial"):
         WavelengthSolution1D.from_asdf(path)
+
+
+def test_resample_propagates_variance(mk_ws_with_transform):
+    ws = mk_ws_with_transform
+    npix = pix_bounds[1]
+    sigma = 3.0
+    spectrum = Spectrum(
+        flux=np.ones(npix) * u.count,
+        spectral_axis=np.arange(npix) * u.pix,
+        uncertainty=StdDevUncertainty(np.full(npix, sigma)),
+    )
+    result = ws.resample(spectrum, nbins=10)
+    var = result.uncertainty.represent_as(VarianceUncertainty).array
+    dl = np.diff(result.spectral_axis.value)[0]
+
+    # The output is sum(w * f) / dl, so the variance of an interior bin is
+    # sum(w**2 * sigma**2) / dl**2 over the contributing pixels.
+    ibin = 5
+    edges = result.spectral_axis.value[ibin] + np.array([-0.5, 0.5]) * dl
+    x_edges = ws.wav_to_pix(edges) + 0.5
+    i1, i2 = np.floor(x_edges).astype(int)
+    w = np.zeros(npix)
+    w[i1 + 1 : i2] = 1.0
+    w[i1] = 1.0 - (x_edges[0] - i1)
+    w[i2] = x_edges[1] - i2
+    expected = (w**2 * sigma**2).sum() / dl**2
+    np.testing.assert_allclose(var[ibin], expected, rtol=1e-6)
+
+
+def _flat_spectrum(npix, flux=100.0, sigma=4.0, first_pixel=0):
+    return Spectrum(
+        flux=np.full(npix, flux) * u.count,
+        spectral_axis=(first_pixel + np.arange(npix)) * u.pix,
+        uncertainty=StdDevUncertainty(np.full(npix, sigma)),
+    )
+
+
+def test_resample_handles_offset_spectral_axis():
+    """A spectrum whose spectral axis does not start at pixel zero must be binned correctly."""
+    npix, first, dispersion = 100, 300, 2.4
+    ws = WavelengthSolution1D(
+        models.Shift(0) | models.Polynomial1D(1, c0=5000, c1=dispersion),
+        (0, first + npix),
+        u.angstrom,
+    )
+    flux = 100.0 + np.arange(npix)  # a slope, so misindexing is visible
+    spectrum = _flat_spectrum(npix, first_pixel=first)
+    spectrum.flux[:] = flux * u.count
+
+    result = ws.resample(spectrum)
+    np.testing.assert_allclose(result.flux.value, flux / dispersion, rtol=1e-9)
+    np.testing.assert_allclose(result.spectral_axis.value[0], ws.p2w(first), rtol=1e-12)
+    np.testing.assert_allclose(result.uncertainty.array, 4.0 / dispersion, rtol=1e-9)
+
+
+def test_resample_handles_decreasing_dispersion():
+    """Wavelength decreasing with pixel number gives an ascending, flux-conserving result."""
+    npix, dispersion = 100, 2.4
+    ws = WavelengthSolution1D(
+        models.Shift(0) | models.Polynomial1D(1, c0=9000, c1=-dispersion), (0, npix), u.angstrom
+    )
+    spectrum = _flat_spectrum(npix)
+    spectrum.flux[:] = (100.0 + np.arange(npix)) * u.count
+
+    result = ws.resample(spectrum)
+    assert np.all(np.diff(result.spectral_axis.value) > 0)
+    # pixel 0 is the reddest, so it lands in the last bin
+    np.testing.assert_allclose(result.flux.value, spectrum.flux.value[::-1] / dispersion, rtol=1e-9)
+    np.testing.assert_allclose(result.uncertainty.array, 4.0 / dispersion, rtol=1e-9)
+
+    result = ws.resample(spectrum, nbins=npix // 4)
+    dlam = np.diff(result.spectral_axis.value)[0]
+    np.testing.assert_allclose(
+        (result.flux.value * dlam).sum(), spectrum.flux.value.sum(), rtol=1e-9
+    )
+
+
+def test_resample_propagates_mask_and_meta():
+    npix, dispersion = 100, 2.4
+    ws = WavelengthSolution1D(
+        models.Shift(0) | models.Polynomial1D(1, c0=5000, c1=dispersion), (0, npix), u.angstrom
+    )
+    spectrum = _flat_spectrum(npix)
+    assert ws.resample(spectrum).mask is None
+
+    mask = np.zeros(npix, dtype=bool)
+    mask[40] = True
+    meta = {"OBJECT": "target", "HISTORY": ["extracted"]}
+    spectrum = Spectrum(
+        flux=spectrum.flux, spectral_axis=spectrum.spectral_axis, mask=mask, meta=meta
+    )
+
+    # one bin per pixel: exactly the masked pixel's bin is flagged
+    result = ws.resample(spectrum)
+    np.testing.assert_array_equal(result.mask, mask)
+
+    # bins half a pixel out of phase with the pixels: the two overlapping bins are flagged
+    edges = ws.p2w(np.arange(npix + 1))  # pixel centres act as bin edges
+    result = ws.resample(spectrum, bin_edges=edges)
+    expected = np.zeros(npix, dtype=bool)
+    expected[39:41] = True
+    np.testing.assert_array_equal(result.mask, expected)
+
+    assert result.meta == meta
+    assert result.meta is not meta
+    result.meta["HISTORY"].append("resampled")
+    assert meta["HISTORY"] == ["extracted"]
+
+
+def test_resample_keeps_explicit_bin_edges():
+    """Non-uniform bin edges must be carried on the output spectral axis as given."""
+    npix, dispersion = 100, 2.4
+    ws = WavelengthSolution1D(
+        models.Shift(0) | models.Polynomial1D(1, c0=5000, c1=dispersion), (0, npix), u.angstrom
+    )
+    spectrum = _flat_spectrum(npix)
+    lo, hi = ws.p2w(np.array([-0.5, npix - 0.5]))
+    edges = lo + (hi - lo) * np.linspace(0, 1, 9) ** 2  # widths growing along the axis
+
+    result = ws.resample(spectrum, bin_edges=edges)
+    np.testing.assert_allclose(result.spectral_axis.bin_edges.value, edges, rtol=1e-12)
+    np.testing.assert_allclose(
+        result.spectral_axis.value, 0.5 * (edges[:-1] + edges[1:]), rtol=1e-12
+    )
+    # integrating over the true bin widths conserves the counts
+    widths = np.diff(result.spectral_axis.bin_edges.value)
+    np.testing.assert_allclose(
+        (result.flux.value * widths).sum(), spectrum.flux.value.sum(), rtol=1e-9
+    )
+
+
+def test_resample_returns_flux_density_per_wavelength():
+    """
+    A flat spectrum in counts per pixel becomes counts per wavelength unit divided by
+    the local dispersion, and the integral over the bins conserves the total counts.
+    """
+    npix, flux, sigma = 100, 100.0, 4.0
+    spectrum = Spectrum(
+        flux=np.full(npix, flux) * u.count,
+        spectral_axis=np.arange(npix) * u.pix,
+        uncertainty=StdDevUncertainty(np.full(npix, sigma)),
+    )
+
+    # linear dispersion: with nbins == npix every bin is exactly one pixel wide
+    dispersion = 2.4
+    ws = WavelengthSolution1D(
+        models.Shift(0) | models.Polynomial1D(1, c0=5000, c1=dispersion), (0, npix), u.angstrom
+    )
+    result = ws.resample(spectrum)
+    assert result.flux.unit == u.count / u.angstrom
+    np.testing.assert_allclose(result.flux.value, flux / dispersion, rtol=1e-9)
+    np.testing.assert_allclose(result.uncertainty.array, sigma / dispersion, rtol=1e-9)
+
+    # non-linear dispersion: the density follows 1 / dldx and the counts are conserved
+    p2w = models.Shift(0) | models.Polynomial1D(2, c0=5000, c1=1.0, c2=0.02)
+    ws = WavelengthSolution1D(p2w, (0, npix), u.angstrom)
+    result = ws.resample(spectrum, nbins=npix)
+    bin_widths = np.diff(np.linspace(*p2w(np.array([-0.5, npix - 0.5])), npix + 1))
+    np.testing.assert_allclose((result.flux.value * bin_widths).sum(), flux * npix, rtol=1e-4)
+    dldx_at_centers = _diff_poly1d(p2w[1])(ws.wav_to_pix(result.spectral_axis.value))
+    np.testing.assert_allclose(result.flux.value, flux / dldx_at_centers, rtol=2e-2)

specreduce/wavesol1d.py

@@ -1,8 +1,8 @@
 import warnings
+from copy import deepcopy
 from functools import cached_property
 from pathlib import Path
 from typing import Callable
-from copy import deepcopy
 
 import asdf
 import astropy.units as u
@@ -18,7 +18,7 @@
 from scipy import optimize
 from scipy.interpolate import interp1d
 
-from specutils import Spectrum
+from specutils import Spectrum, SpectralAxis
 
 
 __all__ = ["WavelengthSolution1D"]
@@ -636,9 +636,22 @@ def resample(
     ) -> Spectrum:
         """Bin the given pixel-space 1D spectrum to a wavelength space conserving the flux.
 
-        This method bins a pixel-space spectrum to a wavelength space using the computed
-        pixel-to-wavelength and wavelength-to-pixel transformations and their derivatives with
-        respect to the spectral axis. The binning is exact and conserves the integrated flux.
+        The input flux is taken to be integrated per pixel (e.g. counts per pixel, as
+        produced by the extraction methods). Each wavelength bin receives the flux of the
+        pixels it overlaps, weighted by the overlapping fraction of each pixel, and the total
+        is divided by the bin width, so the output is a flux density per wavelength unit.
+        The bin edges are stored on the output spectral axis, so the integrated flux is
+        recovered exactly as ``(flux * np.diff(spectral_axis.bin_edges)).sum()`` for any
+        bin grid. The binning is exact and conserves the integrated flux. The variance is propagated
+        with the squared weights, assuming independent pixel noise. The spectral axis of the
+        input must be in pixels with unit spacing but need not start at zero, and the
+        wavelength may increase or decrease with pixel number; the output is always
+        ascending in wavelength.
+
+        Bins that are not aligned with the pixels share pixels with their neighbours, so
+        the resampled uncertainties, while correct for each bin, are correlated between
+        neighbouring bins. Sums or fits that treat the bins as independent underestimate
+        the uncertainty; resampling to a grid finer than the pixels makes this worse.
 
         Parameters
         ----------
@@ -662,63 +675,83 @@ def resample(
 
         Returns
         -------
-            1D spectrum binned to the specified wavelength bins.
+            1D spectrum binned to the specified wavelength bins, with the flux in units of
+            the input flux unit per wavelength unit and a spectral axis that carries the bin
+            edges actually used. The uncertainty, if the input has one,
+            is returned in the same uncertainty class as the input; the mask, if the input
+            has one, flags every bin that received a contribution from a masked pixel; and
+            the metadata is a copy of the input metadata.
         """
-        if nbins is not None and nbins < 0:
+        if nbins is not None and nbins <= 0:
             raise ValueError("Number of bins must be non-zero and positive.")
 
         if self._p2w is None:
             raise ValueError("Wavelength solution not set.")
 
         flux = spectrum.flux.value
         pixels = spectrum.spectral_axis.value
+        npix = flux.size
 
-        if spectrum.uncertainty is not None:
+        has_uncertainty = spectrum.uncertainty is not None
+        if has_uncertainty:
             ucty = spectrum.uncertainty.represent_as(VarianceUncertainty).array
             ucty_type = type(spectrum.uncertainty)
-        else:
-            ucty = np.zeros_like(flux)
-            ucty_type = VarianceUncertainty
-        npix = flux.size
-        nbins = npix if nbins is None else nbins
-        if wlbounds is None:
-            l1, l2 = self.p2w(pixels[[0, -1]] + np.array([-0.5, 0.5]))
-        else:
-            l1, l2 = wlbounds
+        has_mask = spectrum.mask is not None
+        mask = np.asarray(spectrum.mask, dtype=bool) if has_mask else None
 
+        nbins = npix if nbins is None else nbins
         if bin_edges is not None:
-            bin_edges_wav = np.asarray(bin_edges)
+            bin_edges_wav = np.sort(np.asarray(bin_edges, dtype=float))
             nbins = bin_edges_wav.size - 1
         else:
+            if wlbounds is None:
+                l1, l2 = sorted(self.p2w(pixels[[0, -1]] + np.array([-0.5, 0.5])))
+            else:
+                l1, l2 = sorted(wlbounds)
             bin_edges_wav = np.linspace(l1, l2, num=nbins + 1)
 
-        bin_edges_pix = np.clip(self.w2p(bin_edges_wav) + 0.5, 0, npix - 1e-12)
-        bin_edge_ix = np.floor(bin_edges_pix).astype(int)
-        bin_edge_w = bin_edges_pix - bin_edge_ix
-        bin_centers_wav = 0.5 * (bin_edges_wav[:-1] + bin_edges_wav[1:])
+        # Bin edges in array-index space, where pixel j spans [j, j + 1). The spectral axis
+        # need not start at zero, and the wavelength may decrease with pixel number, in
+        # which case the left and right edges of a bin swap places in pixel space.
+        x = np.clip(self.w2p(bin_edges_wav) + 0.5 - pixels[0], 0, npix - 1e-12)
+        x_left, x_right = np.minimum(x[:-1], x[1:]), np.maximum(x[:-1], x[1:])
+        i_left, i_right = np.floor(x_left).astype(int), np.floor(x_right).astype(int)
+
         flux_wl = np.zeros(nbins)
-        ucty_wl = np.zeros(nbins)
+        ucty_wl = np.zeros(nbins) if has_uncertainty else None
+        mask_wl = np.zeros(nbins, dtype=bool) if has_mask else None
         weights = np.zeros(npix)
 
-        dldx = np.diff(self.p2w(np.arange(pixels[0], pixels[-1] + 2) - 0.5))
-
         for i in range(nbins):
-            i1, i2 = bin_edge_ix[i : i + 2]
-            weights[:] = 0.0
+            i1, i2 = i_left[i], i_right[i]
             if i1 != i2:
+                weights[:] = 0.0
                 weights[i1 + 1 : i2] = 1.0
-                weights[i1] = 1 - bin_edge_w[i]
-                weights[i2] = bin_edge_w[i + 1]
+                weights[i1] = 1.0 - (x_left[i] - i1)
+                weights[i2] = x_right[i] - i2
                 sl = slice(i1, i2 + 1)
                 w = weights[sl]
-                flux_wl[i] = (w * flux[sl] * dldx[sl]).sum()
-                ucty_wl[i] = (w**2 * ucty[sl] * dldx[sl]).sum()
+                flux_wl[i] = (w * flux[sl]).sum()
+                if has_uncertainty:
+                    ucty_wl[i] = (w**2 * ucty[sl]).sum()
+                if has_mask:
+                    mask_wl[i] = (mask[sl] & (w > 0)).any()
             else:
-                fracw = bin_edges_pix[i + 1] - bin_edges_pix[i]
-                flux_wl[i] = fracw * flux[i1] * dldx[i1]
-                ucty_wl[i] = fracw**2 * ucty[i1] * dldx[i1]
+                fracw = x_right[i] - x_left[i]
+                flux_wl[i] = fracw * flux[i1]
+                if has_uncertainty:
+                    ucty_wl[i] = fracw**2 * ucty[i1]
+                if has_mask:
+                    mask_wl[i] = mask[i1] & (fracw > 0)
 
         bin_widths_wav = np.diff(bin_edges_wav)
         flux_wl = flux_wl / bin_widths_wav * spectrum.flux.unit / self.unit
-        ucty_wl = VarianceUncertainty(ucty_wl / bin_widths_wav**2).represent_as(ucty_type)
-        return Spectrum(flux_wl, bin_centers_wav * self.unit, uncertainty=ucty_wl)
+        if has_uncertainty:
+            ucty_wl = VarianceUncertainty(ucty_wl / bin_widths_wav**2).represent_as(ucty_type)
+        return Spectrum(
+            flux_wl,
+            SpectralAxis(bin_edges_wav * self.unit, bin_specification="edges"),
+            uncertainty=ucty_wl,
+            mask=mask_wl,
+            meta=deepcopy(spectrum.meta),
+        )