Background estimation with a custom callable statistic

CHANGES.rst

@@ -1,4 +1,4 @@
-1.5.1 (unreleased)
+1.5.2 (unreleased)
 ------------------
 
 New Features
@@ -7,8 +7,16 @@ New Features
 API Changes
 ^^^^^^^^^^^
 
+- ``statistic`` argument in ``Background`` initializer can now be either ``"average"``,
+  ``"median"``, or a custom callable that takes a ``numpy.ma.MaskedArray`` masked array
+  as an input and accepts ``axis`` as an argument. [#253]
+- ``bkg_statistic`` keyword in ``Background.bkg_spectrum()`` is now deprecated and raises
+  a warning if set. The ``statistic`` argument in the ``Background`` initializer should
+  be used instead. [#253]
+
 Bug Fixes
 ^^^^^^^^^
+
 - When all-zero bin encountered in fit_trace with peak_method=gaussian, the bin peak
   will be set to NaN in this caseto work better with DogBoxLSQFitter. [#257]
 

specreduce/background.py

@@ -1,7 +1,9 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import warnings
+from collections.abc import Callable
 from dataclasses import dataclass, field
+from typing import Literal
 
 import numpy as np
 from astropy import units as u
@@ -39,7 +41,8 @@
     statistic
         Statistic to use when computing the background.  ``average`` will
         account for partial pixel weights, ``median`` will include all partial
-        pixels.
+        pixels. If you provide your own function, it must take a `~numpy.ma.MaskedArray`
+        masked array as input and accept an ``axis`` argument.
     disp_axis
         Dispersion axis.
     crossdisp_axis
@@ -70,7 +73,7 @@
     image: ImageLike
     traces: list = field(default_factory=list)
     width: float = 5
-    statistic: str = "average"
+    statistic: Literal["average", "median"] | Callable[..., np.ndarray] = "average"
     disp_axis: int = 1
     crossdisp_axis: int = 0
     mask_treatment: MaskingOption = "apply"
@@ -146,16 +149,19 @@
 
         self.bkg_wimage = bkg_wimage
 
-        if self.statistic == "average":
+        if isinstance(self.statistic, Callable):
+            img.mask = (self.bkg_wimage == 0) | self.image.mask
+            self._bkg_array = self.statistic(img, axis=self.crossdisp_axis)
+        elif self.statistic == "average":
             self._bkg_array = np.ma.average(img, weights=self.bkg_wimage, axis=self.crossdisp_axis)
-
         elif self.statistic == "median":
-            # combine where background weight image is 0 with image masked (which already
-            # accounts for non-finite data that wasn't already masked)
-            img.mask = np.logical_or(self.bkg_wimage == 0, self.image.mask)
+            img.mask = (self.bkg_wimage == 0) | self.image.mask
             self._bkg_array = np.ma.median(img, axis=self.crossdisp_axis)
         else:
-            raise ValueError("statistic must be 'average' or 'median'")
+            raise ValueError(
+                "statistic must be 'average', 'median', or a callable function that takes a masked"
+                "array as input and an axis argument."
+            )
 
     def _set_traces(self):
         """Determine `traces` from input. If an integer/float or list if int/float
@@ -216,10 +222,11 @@
             separation from ``trace_object`` for the background regions
         width : float
             width of each background aperture in pixels
-        statistic: string
-            statistic to use when computing the background.  'average' will
-            account for partial pixel weights, 'median' will include all partial
-            pixels.
+        statistic: string or Callable
+            Statistic to use when computing the background.  ``average`` will
+            account for partial pixel weights, ``median`` will include all partial
+            pixels. If you provide your own function, it must take a `~numpy.ma.MaskedArray`
+            masked array as input and accept an ``axis`` argument.
         disp_axis : int
             dispersion axis
         crossdisp_axis : int
@@ -265,10 +272,11 @@
             above the trace, negative below.
         width : float
             width of each background aperture in pixels
-        statistic: string
-            statistic to use when computing the background.  'average' will
-            account for partial pixel weights, 'median' will include all partial
-            pixels.
+        statistic: string or Callable
+            Statistic to use when computing the background.  ``average`` will
+            account for partial pixel weights, ``median`` will include all partial
+            pixels. If you provide your own function, it must take a `~numpy.ma.MaskedArray`
+            masked array as input and accept an ``axis`` argument.
         disp_axis : int
             dispersion axis
         crossdisp_axis : int
@@ -312,12 +320,9 @@
             kwargs = {}
         else:
             kwargs = {"spectral_axis_index": arr.ndim - 1}
-        return Spectrum(
-            arr * image.unit,
-            spectral_axis=image.spectral_axis, **kwargs
-        )
+        return Spectrum(arr * image.unit, spectral_axis=image.spectral_axis, **kwargs)
 
-    def bkg_spectrum(self, image=None, bkg_statistic="sum"):
+    def bkg_spectrum(self, image=None, bkg_statistic=None):
         """
         Expose the 1D spectrum of the background.
 
@@ -328,40 +333,21 @@
             (spatial) direction is axis 0 and dispersion (wavelength)
             direction is axis 1. If None, will extract the background
             from ``image`` used to initialize the class. [default: None]
-        bkg_statistic : {'average', 'median', 'sum'}, optional
-            Statistical method used to collapse the background image. [default: ``'sum'``]
-            Supported values are:
-
-            - ``'average'`` : Uses the mean (`numpy.nanmean`).
-            - ``'median'`` : Uses the median (`numpy.nanmedian`).
-            - ``'sum'`` : Uses the sum (`numpy.nansum`).
 
         Returns
         -------
         spec : `~specutils.Spectrum1D`
             The background 1-D spectrum, with flux expressed in the same
-            units as the input image (or u.DN if none were provided) and
+            units as the input image (or DN if none were provided) and
             the spectral axis expressed in pixel units.
         """
-        bkg_image = self.bkg_image(image)
-
-        if bkg_statistic == 'sum':
-            statistic_function = np.nansum
-        elif bkg_statistic == 'median':
-            statistic_function = np.nanmedian
-        elif bkg_statistic == 'average':
-            statistic_function = np.nanmean
-        else:
-            raise ValueError(f"Background statistic {bkg_statistic} is not supported. "
-                             "Please choose from: average, median, or sum.")
-
-        try:
-            return bkg_image.collapse(statistic_function, axis=self.crossdisp_axis)
-        except u.UnitTypeError:
-            # can't collapse with a spectral axis in pixels because
-            # SpectralCoord only allows frequency/wavelength equivalent units...
-            ext1d = statistic_function(bkg_image.flux, axis=self.crossdisp_axis)
-            return Spectrum(ext1d, bkg_image.spectral_axis)
+        if bkg_statistic is not None:
+            warnings.warn(
+                "'bkg_statistic' is deprecated and will be removed in a future release. "
+                "Please use the 'statistic' argument in the Background initializer instead.",
+                DeprecationWarning,
+            )
+        return Spectrum(self._bkg_array * self.image.unit, spectral_axis=self.image.spectral_axis)
 
     def sub_image(self, image=None):
         """
@@ -370,7 +356,7 @@
         Parameters
         ----------
         image : nddata-compatible image or None
-            image with 2-D spectral image data.  If None, will extract
+            image with 2-D spectral image data.  If None, will subtract
             the background from ``image`` used to initialize the class.
 
         Returns

specreduce/tests/test_background.py

@@ -73,28 +73,20 @@ def test_background(
     # the final 1D spectra.
     img[0, 0] = np.nan  # out of window
     img[trace_pos, 0] = np.nan  # in window
-    stats = ["average", "median"]
+    stats = ["average", "median", np.nanmean, np.nanmedian]
 
     for st in stats:
         bg = Background(img, trace - bkg_sep, width=bkg_width, statistic=st)
         assert np.isnan(bg.image.flux).sum() == 2
         assert np.isnan(bg._bkg_array).sum() == 0
         assert np.isnan(bg.bkg_spectrum().flux).sum() == 0
-        assert np.isnan(bg.sub_spectrum().flux).sum() == 0
 
-    bkg_spec_avg = bg1.bkg_spectrum(bkg_statistic="average")
+    bkg_spec_avg = bg1.bkg_spectrum()
     assert_allclose(bkg_spec_avg.mean().value, 14.5, rtol=0.5)
 
-    bkg_spec_median = bg1.bkg_spectrum(bkg_statistic="median")
+    bkg_spec_median = bg1.bkg_spectrum()
     assert_allclose(bkg_spec_median.mean().value, 14.5, rtol=0.5)
 
-    with pytest.raises(
-        ValueError,
-        match="Background statistic max is not supported. "
-        "Please choose from: average, median, or sum.",
-    ):
-        bg1.bkg_spectrum(bkg_statistic="max")
-
 
 def test_warnings_errors(mk_test_spec_no_spectral_axis):
     image = mk_test_spec_no_spectral_axis
@@ -288,15 +280,11 @@ def test_mask_treatment_bkg_img_spectrum(self, method, expected):
 
             # test background image matches 'expected'
             bk_img = background.bkg_image()
-            # change this and following assertions to assert_quantity_allclose once
-            # issue #213 is fixed
             np.testing.assert_allclose(bk_img.flux.value, np.tile(expected, (img_size, 1)))
 
-            # test background spectrum matches 'expected' times the number of rows
-            # in cross disp axis, since this is a sum and all values in a col are
-            # the same.
+            # test background spectrum matches 'expected'
             bk_spec = background.bkg_spectrum()
-            np.testing.assert_allclose(bk_spec.flux.value, expected * img_size)
+            np.testing.assert_allclose(bk_spec.flux.value, expected)
 
     def test_sub_bkg_image(self):
         """