implement average and median support for bkg_spectrum (#253)

Author: gibsongreen

specreduce/background.py

@@ -330,7 +330,7 @@ def bkg_image(self, image=None):
                                   (image.shape[0], 1)) * image.unit,
                           spectral_axis=image.spectral_axis)
 
-    def bkg_spectrum(self, image=None):
+    def bkg_spectrum(self, image=None, bkg_statistic="sum"):
         """
         Expose the 1D spectrum of the background.
 
@@ -341,6 +341,13 @@ def bkg_spectrum(self, image=None):
             (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
         -------
@@ -351,12 +358,22 @@ def bkg_spectrum(self, image=None):
         """
         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(np.nansum, axis=self.crossdisp_axis)
+            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 = np.nansum(bkg_image.flux, axis=self.crossdisp_axis)
+            ext1d = statistic_function(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

@@ -1,6 +1,7 @@
 from astropy.nddata import NDData
 import astropy.units as u
 import numpy as np
+from numpy.testing import assert_allclose
 import pytest
 from specutils import Spectrum1D
 
@@ -80,6 +81,16 @@ def test_background(mk_test_img_raw, mk_test_spec_no_spectral_axis,
         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')
+    assert_allclose(bkg_spec_avg.mean().value, 14.5, rtol=0.5)
+
+    bkg_spec_median = bg1.bkg_spectrum(bkg_statistic='median')
+    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