TiltSolution.resample fixes

CHANGES.rst

@@ -38,6 +38,14 @@ API Changes
   on instantiation, and its removal has been rescheduled from v2.0 to v1.11.
   Use ``specreduce.wavecal1d.WavelengthCalibration1D`` instead. [#316]
 
+Bug Fixes
+^^^^^^^^^
+
+- ``TiltSolution.resample`` now propagates the input uncertainty, returned in the same
+  uncertainty class as the input, marks output bins that received a contribution from a
+  masked input pixel, and copies the input metadata to the resampled ``NDData``.
+  Previously all three were dropped. [#XXX]
+
 Other changes
 ^^^^^^^^^^^^^
 

docs/tilt_correction/tilt_correction.rst

@@ -262,4 +262,7 @@ flux-conserving resampling independently of the calibration workflow.
         corrected = ts.resample(science_frame, bin_edges=np.linspace(50, 950, 501))
 
     The ``resample`` method accepts a ``mask_treatment`` parameter with the same options as
-    the :class:`~specreduce.tilt_correction.TiltCorrection` constructor.
+    the :class:`~specreduce.tilt_correction.TiltCorrection` constructor. The returned
+    :class:`~astropy.nddata.NDData` carries the resampled uncertainty (in the same
+    uncertainty class as the input), a mask flagging every bin that received a contribution
+    from a masked input pixel, and a copy of the input metadata.

specreduce/tests/test_tilt_solution.py

@@ -1,10 +1,26 @@
+import astropy.units as u
 import numpy as np
 import pytest
+from astropy.io import fits
 from astropy.modeling import models
+from astropy.modeling.models import Shift, Polynomial2D
+from astropy.nddata import (
+    NDData,
+    CCDData,
+    StdDevUncertainty,
+    VarianceUncertainty,
+    InverseVariance,
+)
 
 from specreduce.tilt_solution import TiltSolution, diff_poly2d_x
 
 
+def _linear_ts(ny, nx, shift=0.0, disp_axis=1):
+    """A fit-free tilt solution: detector x = rectified x + shift, for every row."""
+    solution = Shift(0) & Shift(0) | Polynomial2D(1, c0_0=shift, c1_0=1.0)
+    return TiltSolution(solution, disp_axis=disp_axis, image_shape=(ny, nx))
+
+
 def test_diff_poly2d_x_valid_derivative():
     model = models.Polynomial2D(degree=2, c0_0=1, c1_0=2, c2_0=3, c0_1=4, c1_1=5, c0_2=6)
     derivative = diff_poly2d_x(model)
@@ -116,7 +132,124 @@ def test_resample(mk_default_tc, mk_arc_frames):
     tc = mk_default_tc
     tc.find_arc_lines(3.0, 5.0)
     tc.fit(4)
-    tc.solution.resample(arcs[0])
+    result = tc.solution.resample(arcs[0])
+    assert isinstance(result.uncertainty, StdDevUncertainty)
+    assert result.uncertainty.array.shape == result.data.shape
+
+
+@pytest.mark.parametrize(
+    "uncertainty_cls", [StdDevUncertainty, VarianceUncertainty, InverseVariance]
+)
+def test_resample_preserves_uncertainty_type(uncertainty_cls):
+    ny, nx = 4, 12
+    data = np.full((ny, nx), 10.0)
+    image = NDData(data * u.ct, uncertainty=uncertainty_cls(np.full((ny, nx), 4.0)))
+    result = _linear_ts(ny, nx).resample(image)
+    assert isinstance(result.uncertainty, uncertainty_cls)
+    assert result.uncertainty.array.shape == result.data.shape
+
+
+def test_resample_uncertainty_identity():
+    ny, nx = 4, 12
+    data = np.full((ny, nx), 10.0)
+    variance = np.arange(1.0, ny * nx + 1).reshape(ny, nx)
+    image = NDData(data * u.ct, uncertainty=VarianceUncertainty(variance))
+    result = _linear_ts(ny, nx).resample(image)
+    np.testing.assert_allclose(result.data, data)
+    np.testing.assert_allclose(result.uncertainty.array, variance)
+    assert result.uncertainty.unit == u.ct**2
+
+
+def test_resample_uncertainty_half_pixel_shift():
+    """
+    Each rectified bin takes half of two neighboring detector pixels, so the
+    variance is 0.25 + 0.25 of the input variance, not 0.5 + 0.5. The last bin
+    covers half of the last pixel only.
+    """
+    ny, nx = 3, 10
+    data = np.full((ny, nx), 10.0)
+    image = NDData(data * u.ct, uncertainty=StdDevUncertainty(np.full((ny, nx), 3.0)))
+    result = _linear_ts(ny, nx, shift=0.5).resample(image)
+    np.testing.assert_allclose(result.data[:, :-1], 10.0)
+    np.testing.assert_allclose(result.data[:, -1], 5.0)
+    np.testing.assert_allclose(result.uncertainty.array[:, :-1], np.sqrt(4.5), rtol=1e-9)
+    np.testing.assert_allclose(result.uncertainty.array[:, -1], 1.5, rtol=1e-9)
+
+
+@pytest.mark.parametrize(
+    "image",
+    [
+        NDData(np.full((4, 12), 10.0) * u.ct),
+        np.full((4, 12), 10.0),
+        np.full((4, 12), 10.0) * u.ct,
+    ],
+)
+def test_resample_without_uncertainty_returns_none(image):
+    result = _linear_ts(4, 12).resample(image)
+    assert result.uncertainty is None
+
+
+def test_resample_copies_meta():
+    ny, nx = 4, 12
+    data = np.full((ny, nx), 10.0)
+    ts = _linear_ts(ny, nx)
+
+    meta = {"OBJECT": "target", "HISTORY": ["step 1"]}
+    result = ts.resample(NDData(data * u.ct, meta=meta))
+    assert result.meta == meta
+    assert result.meta is not meta
+    result.meta["OBJECT"] = "changed"
+    result.meta["HISTORY"].append("step 2")
+    assert meta == {"OBJECT": "target", "HISTORY": ["step 1"]}
+
+    header = fits.Header([("OBJECT", "target"), ("EXPTIME", 30.0)])
+    result = ts.resample(CCDData(data, unit="ct", meta=header))
+    assert isinstance(result.meta, fits.Header)
+    assert result.meta["OBJECT"] == "target" and result.meta["EXPTIME"] == 30.0
+    result.meta["OBJECT"] = "changed"
+    assert header["OBJECT"] == "target"
+
+    assert len(ts.resample(NDData(data * u.ct)).meta) == 0
+
+
+def test_resample_propagates_mask():
+    ny, nx = 4, 12
+    data = np.full((ny, nx), 10.0)
+    mask = np.zeros((ny, nx), dtype=bool)
+    mask[2, 3] = True
+
+    result = _linear_ts(ny, nx).resample(NDData(data * u.ct, mask=mask), mask_treatment="apply")
+    assert result.mask.dtype == bool
+    np.testing.assert_array_equal(result.mask, mask)
+
+    # a half-pixel shift spreads the masked pixel over the two bins that overlap it
+    result = _linear_ts(ny, nx, shift=0.5).resample(
+        NDData(data * u.ct, mask=mask), mask_treatment="apply"
+    )
+    expected = np.zeros((ny, nx), dtype=bool)
+    expected[2, 2:4] = True
+    np.testing.assert_array_equal(result.mask, expected)
+
+    # fill treatments drop the mask before resampling
+    result = _linear_ts(ny, nx).resample(NDData(data * u.ct, mask=mask), mask_treatment="zero_fill")
+    assert not result.mask.any()
+
+
+def test_resample_disp_axis_0_propagates_arrays():
+    n = 8
+    data = np.arange(1.0, n * n + 1).reshape(n, n)
+    variance = 2.0 * data
+    mask = np.zeros((n, n), dtype=bool)
+    mask[1, 5] = True
+    image = NDData(data * u.ct, uncertainty=VarianceUncertainty(variance), mask=mask)
+
+    result = _linear_ts(n, n, disp_axis=0).resample(image)
+    np.testing.assert_allclose(result.data, data.T)
+    np.testing.assert_allclose(result.uncertainty.array, variance.T)
+    np.testing.assert_array_equal(result.mask, mask.T)
+
+    result = _linear_ts(n, n, disp_axis=0).resample(image, nbins=2 * n)
+    assert result.data.shape == result.uncertainty.array.shape == result.mask.shape == (2 * n, n)
 
 
 @pytest.mark.remote_data
@@ -195,14 +328,15 @@ def test_resample_disp_axis_0(mk_default_tc, mk_arc_frames):
     tc.fit(4)
 
     # Use a square crop so _parse_image works with disp_axis=0
-    from astropy.nddata import NDData
-    import astropy.units as u
     ny = arcs[0].data.shape[0]
-    square = NDData(arcs[0].data[:, :ny] * u.ct)
+    square = NDData(
+        arcs[0].data[:, :ny] * u.ct, uncertainty=StdDevUncertainty(np.full((ny, ny), 5.0))
+    )
 
     ts = tc.solution
     ts.disp_axis = 0
     result = ts.resample(square, nbins=ny)
+    assert result.uncertainty.array.shape == result.data.shape
     # With disp_axis=0, output should be transposed
     assert result.data.shape[0] == ny  # nbins along axis 0
     assert result.data.shape[1] == ny  # cdisp along axis 1

specreduce/tilt_solution.py

@@ -1,4 +1,5 @@
 import warnings
+from copy import deepcopy
 from functools import cached_property
 from typing import Sequence, Literal
 
@@ -7,7 +8,7 @@
 import numpy as np
 from astropy.modeling import models, fitting, Model
 from astropy.modeling.models import Identity, Mapping, Shift, Polynomial2D
-from astropy.nddata import NDData
+from astropy.nddata import NDData, VarianceUncertainty
 from astropy.utils.exceptions import AstropyUserWarning
 from gwcs import coordinate_frames
 from numpy import ndarray
@@ -296,14 +297,34 @@ def resample(
         Returns
         -------
         NDData
-            NDData instance containing the flux values resampled into the uniform grid
-            defined by ``nbins``, ``bounds``, or ``bin_edges``.
+            The flux resampled into the uniform grid defined by ``nbins``, ``bounds``, or
+            ``bin_edges``. If the input carries an uncertainty, the resampled uncertainty is
+            attached in the same uncertainty class as the input. The mask marks every output
+            bin that received a contribution from a masked input pixel, and the metadata is a
+            copy of the input metadata. The WCS is not propagated.
+
+        Notes
+        -----
+        Each output bin is a linear combination of detector pixels,
+        ``F = n * sum_j k_j f_j`` with ``k_j`` the fractional pixel overlap times the
+        Jacobian of the transformation and ``n`` the per-row flux-conservation factor.
+        The variance is propagated as ``Var = n**2 * sum_j k_j**2 var_j``, assuming
+        independent pixel noise and treating ``n`` and the Jacobian as deterministic.
         """
 
+        # The metadata and the presence of an uncertainty are read from the input itself:
+        # parse_image drops the metadata and fabricates a unit variance for bare arrays.
+        meta = deepcopy(getattr(flux, "meta", None))
+        has_uncertainty = getattr(flux, "uncertainty", None) is not None
+
         # TODO: In the future, we want to make sure that we don't copy the data unless absolutely
         # necessary.
         im = parse_image(flux, disp_axis=self.disp_axis, mask_treatment=mask_treatment)
         flux = im.flux.value
+        mask = im.mask.astype(bool)
+        if has_uncertainty:
+            uncertainty_type = type(im.uncertainty)
+            variance = im.uncertainty.represent_as(VarianceUncertainty).array
 
         ny, nx = flux.data.shape
         ypix = np.arange(ny)
@@ -316,6 +337,8 @@ def resample(
         bin_edge_w = bin_edges_det - bin_edge_ix
 
         resampled_flux = np.zeros((ny, nbins))
+        resampled_variance = np.zeros((ny, nbins)) if has_uncertainty else None
+        resampled_mask = np.zeros((ny, nbins), dtype=bool)
         weights = np.zeros((ny, nx))
 
         # Calculate the derivative of the tilt-corrected space -> detector space transformation with
@@ -343,9 +366,11 @@ def resample(
             # the tilt-corrected flux is the detector flux in that pixel, scaled by the width of the
             # tilt-corrected bin in detector coordinates and the derivative dtdx.
             if m.any():
-                resampled_flux[:, i] = (
-                    (bin_edges_det[:, i + 1] - bin_edges_det[:, i]) * flux[ys, i1] * dtdx[ys, i1]
-                )
+                k = (bin_edges_det[:, i + 1] - bin_edges_det[:, i]) * dtdx[ys, i1]
+                resampled_flux[:, i] = k * flux[ys, i1]
+                resampled_mask[:, i] = mask[ys, i1]
+                if has_uncertainty:
+                    resampled_variance[:, i] = k**2 * variance[ys, i1]
 
             # For rows where the tilt-corrected bin spans multiple detector pixels, calculate the
             # tilt-corrected flux as a weighted sum of the detector flux, multiplied by dtdx,
@@ -358,11 +383,28 @@ def resample(
                 w[(ixc > i1[:, None]) & (ixc < i2[:, None])] = 1
                 w[ys, i1 - imin] = 1.0 - bin_edge_w[:, i]
                 w[ys, i2 - imin] = bin_edge_w[:, i + 1]
-                resampled_flux[~m, i] = (flux[~m, imin:imax] * dtdx[~m, imin:imax] * w[~m]).sum(1)
+                k = dtdx[~m, imin:imax] * w[~m]
+                resampled_flux[~m, i] = (flux[~m, imin:imax] * k).sum(1)
+                resampled_mask[~m, i] = (mask[~m, imin:imax] & (w[~m] > 0)).any(1)
+                if has_uncertainty:
+                    resampled_variance[~m, i] = (variance[~m, imin:imax] * k**2).sum(1)
 
         # Apply the normalization factor to conserve flux
         resampled_flux *= n[:, None]
+        if has_uncertainty:
+            resampled_variance *= n[:, None] ** 2
+
         if self.disp_axis == 0:
             resampled_flux = resampled_flux.T
-
-        return NDData(resampled_flux * im.unit)
+            resampled_mask = resampled_mask.T
+            if has_uncertainty:
+                resampled_variance = resampled_variance.T
+
+        uncertainty = None
+        if has_uncertainty:
+            uncertainty = VarianceUncertainty(resampled_variance * im.unit**2).represent_as(
+                uncertainty_type
+            )
+        return NDData(
+            resampled_flux * im.unit, uncertainty=uncertainty, mask=resampled_mask, meta=meta
+        )