diff --git a/CHANGES.rst b/CHANGES.rst index 1bbdaad1..06eb977a 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -4,7 +4,8 @@ New Features ^^^^^^^^^^^^ -- Added uncertainty propagation to ``specreduce.extract.BoxcarExtract``. +- Added uncertainty propagation to ``specreduce.extract.BoxcarExtract`` and + ``specreduce.extract.HorneExtract``. The extracted spectra have now proper uncertainties. Other changes ^^^^^^^^^^^^^ diff --git a/specreduce/extract.py b/specreduce/extract.py index 61ccd713..62c7cd39 100644 --- a/specreduce/extract.py +++ b/specreduce/extract.py @@ -696,25 +696,32 @@ def __call__( if bkgrd_prof is None and profile_type == "gaussian": bkgrd_prof = models.Polynomial1D(2) + # Store original uncertainty type BEFORE parsing (parsing converts to VarianceUncertainty) + if hasattr(image, "uncertainty") and image.uncertainty is not None: + orig_uncty_type = type(image.uncertainty) + else: + orig_uncty_type = VarianceUncertainty # default if variance passed separately + self.image = self._parse_image(image, variance, mask, unit, disp_axis) + variance = self.image.uncertainty.represent_as(VarianceUncertainty).array mask = self.image.mask.astype(bool) | (~np.isfinite(self.image.data)) unit = self.image.unit - img = self.image.data + flux = self.image.data - ncross = img.shape[crossdisp_axis] - ndisp = img.shape[disp_axis] + ncross = flux.shape[crossdisp_axis] + ndisp = flux.shape[disp_axis] # If the trace is not flat, shift the rows in each column # so the image is aligned along the trace: if not isinstance(trace_object, FlatTrace): - img = _align_along_trace( - img, trace_object.trace, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis + flux = _align_along_trace( + flux, trace_object.trace, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis ) if profile_type == "gaussian": fit_ext_kernel = self._fit_gaussian_spatial_profile( - img, disp_axis, crossdisp_axis, mask, bkgrd_prof + flux, disp_axis, crossdisp_axis, mask, bkgrd_prof ) if isinstance(trace_object, FlatTrace): mean_cross_pix = trace_object.trace @@ -740,14 +747,14 @@ def __call__( kx, ky = interp_degree_interpolated_profile interp_spatial_prof = self._fit_spatial_profile( - img, disp_axis, crossdisp_axis, mask, n_bins_interpolated_profile, kx, ky + flux, disp_axis, crossdisp_axis, mask, n_bins_interpolated_profile, kx, ky ) # add private attribute to save fit profile. should this be public? self._interp_spatial_prof = interp_spatial_prof xd_pixels = np.arange(ncross) - kernel_vals = np.zeros(img.shape) + kernel_vals = np.zeros(flux.shape) norms = np.full(ndisp, np.nan) valid = ~mask @@ -767,11 +774,20 @@ def __call__( norms[idisp] = trapezoid(fitted_col, dx=1)[0] with np.errstate(divide="ignore", invalid="ignore"): - num = np.sum(np.where(valid, img * kernel_vals / variance, 0.0), axis=crossdisp_axis) + num = np.sum(np.where(valid, flux * kernel_vals / variance, 0.0), axis=crossdisp_axis) den = np.sum(np.where(valid, kernel_vals**2 / variance, 0.0), axis=crossdisp_axis) - extraction = (num / den) * norms + extracted_flux = (num / den) * norms + extracted_variance = norms**2 / den + + spectrum_uncty = VarianceUncertainty( + extracted_variance * self.image.unit**2 + ).represent_as(orig_uncty_type) - return Spectrum(extraction * unit, spectral_axis=self.image.spectral_axis) + return Spectrum( + extracted_flux * unit, + spectral_axis=self.image.spectral_axis, + uncertainty=spectrum_uncty, + ) def _align_along_trace(img, trace_array, disp_axis=1, crossdisp_axis=0): diff --git a/specreduce/tests/test_extract.py b/specreduce/tests/test_extract.py index 9df36c6c..34c145d7 100644 --- a/specreduce/tests/test_extract.py +++ b/specreduce/tests/test_extract.py @@ -536,6 +536,107 @@ def test_horne_interpolated_nbins_fails(mk_test_img): ex.spectrum +def test_horne_uncertainty_propagation(): + """Test that HorneExtract propagates uncertainties correctly.""" + # Create an image with a Gaussian source and uniform variance + nrows, ncols = 20, 30 + input_variance = 4.0 + img = Spectrum( + np.zeros((nrows, ncols)) * u.DN, + uncertainty=VarianceUncertainty(np.full((nrows, ncols), input_variance) * u.DN**2), + spectral_axis=np.arange(ncols) * u.pix + ) + add_gaussian_source(img, amps=100.0, stddevs=2.0, means=10) + + trace = FlatTrace(img, 10) + extracted = HorneExtract(img, trace)() + + # Check uncertainty exists and is the correct type + assert extracted.uncertainty is not None + assert isinstance(extracted.uncertainty, VarianceUncertainty) + + # Verify uncertainty values are finite and positive + assert np.all(np.isfinite(extracted.uncertainty.array)) + assert np.all(extracted.uncertainty.array > 0) + + # Check units are correct (flux_unit^2) + assert extracted.uncertainty.unit == u.DN**2 + + # Calculate expected variance analytically. + # For optimal extraction: var_out = norms^2 / sum(kernel^2 / var_in) + # With uniform variance: var_out = var_in * norms^2 / sum(kernel^2) + # where kernel is the normalized spatial profile and norms = sum(kernel) = 1 + gauss = models.Gaussian1D(amplitude=100.0, mean=10.0, stddev=2.0) + kernel = gauss(np.arange(nrows)) + kernel_normalized = kernel / kernel.sum() + expected_variance = input_variance / np.sum(kernel_normalized**2) + + # All columns should have the same uncertainty (uniform source and variance) + assert np.allclose(extracted.uncertainty.array, expected_variance, rtol=1e-5) + + +def test_horne_uncertainty_stddev_type(): + """Test that StdDevUncertainty type is preserved in HorneExtract.""" + nrows, ncols = 20, 30 + img = Spectrum( + np.zeros((nrows, ncols)) * u.DN, + uncertainty=StdDevUncertainty(np.full((nrows, ncols), 2.0) * u.DN), + spectral_axis=np.arange(ncols) * u.pix + ) + add_gaussian_source(img, amps=100.0, stddevs=2.0, means=10) + + trace = FlatTrace(img, 10) + extracted = HorneExtract(img, trace)() + + assert isinstance(extracted.uncertainty, StdDevUncertainty) + assert np.all(np.isfinite(extracted.uncertainty.array)) + assert np.all(extracted.uncertainty.array > 0) + assert extracted.uncertainty.unit == u.DN + + +def test_horne_uncertainty_inverse_variance_type(): + """Test that InverseVariance type is preserved in HorneExtract.""" + nrows, ncols = 20, 30 + img = Spectrum( + np.zeros((nrows, ncols)) * u.DN, + uncertainty=InverseVariance(np.full((nrows, ncols), 0.25) / u.DN**2), + spectral_axis=np.arange(ncols) * u.pix + ) + add_gaussian_source(img, amps=100.0, stddevs=2.0, means=10) + + trace = FlatTrace(img, 10) + extracted = HorneExtract(img, trace)() + + assert isinstance(extracted.uncertainty, InverseVariance) + assert np.all(np.isfinite(extracted.uncertainty.array)) + assert np.all(extracted.uncertainty.array > 0) + assert extracted.uncertainty.unit == 1 / u.DN**2 + + +def test_horne_uncertainty_with_mask(): + """Test uncertainty propagation with masked pixels in HorneExtract.""" + nrows, ncols = 20, 30 + mask = np.zeros((nrows, ncols), dtype=bool) + # Mask some pixels at column 5 + mask[8:12, 5] = True + + img = Spectrum( + np.zeros((nrows, ncols)) * u.DN, + uncertainty=VarianceUncertainty(np.full((nrows, ncols), 4.0) * u.DN**2), + mask=mask, + spectral_axis=np.arange(ncols) * u.pix + ) + add_gaussian_source(img, amps=100.0, stddevs=2.0, means=10) + + trace = FlatTrace(img, 10) + extracted = HorneExtract(img, trace)() + + assert extracted.uncertainty is not None + # Uncertainty at masked column should be larger (fewer valid pixels) + # compared to unmasked columns + assert extracted.uncertainty.array[5] > extracted.uncertainty.array[0] + + class TestMasksExtract: def mk_flat_gauss_img(self, nrows=200, ncols=160, nan_slices=None, add_noise=True):