diff --git a/specreduce/background.py b/specreduce/background.py index 242eaf6..7e5d262 100644 --- a/specreduce/background.py +++ b/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): diff --git a/specreduce/tests/test_background.py b/specreduce/tests/test_background.py index 388da5c..122ac30 100644 --- a/specreduce/tests/test_background.py +++ b/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