bkg_spectrum to take func instead of string

CHANGES.rst

@@ -1,4 +1,4 @@
-1.5.1 (unreleased)
+1.5.2 (unreleased)
 ------------------
 
 New Features
@@ -7,8 +7,13 @@ New Features
 API Changes
 ^^^^^^^^^^^
 
+- ``bkg_statistic`` keyword in ``Background.bkg_spectrum()`` now takes callable
+  instead of string. Its default is also changed from ``"sum"`` (``np.nansum``)
+  to ``np.nanmedian`` to better suit background calculation. [#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

@@ -311,7 +311,7 @@ def bkg_image(self, image=None):
             spectral_axis=image.spectral_axis,
         )
 
-    def bkg_spectrum(self, image=None, bkg_statistic="sum"):
+    def bkg_spectrum(self, image=None, bkg_statistic=np.nanmedian):
         """
         Expose the 1D spectrum of the background.
 
@@ -322,39 +322,26 @@ def bkg_spectrum(self, image=None, bkg_statistic="sum"):
             (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`).
+        bkg_statistic : func, optional
+            Statistical method used to collapse the background image.
+            If you provide your own function, it must take an `~astropy.units.Quantity`
+            array as input and accept an ``axis`` argument.
 
         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)
+            return bkg_image.collapse(bkg_statistic, 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)
+            ext1d = bkg_statistic(bkg_image.flux, axis=self.crossdisp_axis)
             return Spectrum1D(ext1d, bkg_image.spectral_axis)
 
     def sub_image(self, image=None):

specreduce/tests/test_background.py

@@ -59,7 +59,7 @@ def test_background(
     assert np.allclose(sub4.flux, sub5.flux)
     assert np.allclose(sub5.flux, sub6.flux)
 
-    bkg_spec = bg1.bkg_spectrum()
+    bkg_spec = bg1.bkg_spectrum(bkg_statistic=np.nansum)
     assert isinstance(bkg_spec, Spectrum1D)
     sub_spec = bg1.sub_spectrum()
     assert isinstance(sub_spec, Spectrum1D)
@@ -79,22 +79,15 @@ def test_background(
         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.bkg_spectrum(bkg_statistic=np.nansum).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(bkg_statistic=np.nanmean)
     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(bkg_statistic=np.nanmedian)
     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
@@ -295,7 +288,7 @@ def test_mask_treatment_bkg_img_spectrum(self, method, expected):
             # 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.
-            bk_spec = background.bkg_spectrum()
+            bk_spec = background.bkg_spectrum(bkg_statistic=np.nansum)
             np.testing.assert_allclose(bk_spec.flux.value, expected * img_size)
 
     def test_sub_bkg_image(self):