Horne extraction error propagation

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
 ^^^^^^^^^^^^^

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):

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):