diff --git a/CHANGES.rst b/CHANGES.rst index 0464fbdc..656d0e42 100644 --- a/CHANGES.rst +++ b/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 ^^^^^^^^^^^^^ diff --git a/docs/tilt_correction/tilt_correction.rst b/docs/tilt_correction/tilt_correction.rst index 6398243f..2672f36b 100644 --- a/docs/tilt_correction/tilt_correction.rst +++ b/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. diff --git a/specreduce/tests/test_tilt_solution.py b/specreduce/tests/test_tilt_solution.py index c59edb62..53ee35fc 100644 --- a/specreduce/tests/test_tilt_solution.py +++ b/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 diff --git a/specreduce/tilt_solution.py b/specreduce/tilt_solution.py index f146bb8f..c056d04d 100644 --- a/specreduce/tilt_solution.py +++ b/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 + )