From 4b0bdd54cb8ccfaa8c7ca9a277b4e834d3988126 Mon Sep 17 00:00:00 2001 From: ojustino Date: Wed, 16 Nov 2022 13:33:37 -0500 Subject: [PATCH] Fix Background.sub_image for physical wavelengths --- specreduce/background.py | 18 ++++++++++++------ specreduce/tests/test_background.py | 22 +++++++++++++++++++++- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/specreduce/background.py b/specreduce/background.py index b1887db9..b6ca2586 100644 --- a/specreduce/background.py +++ b/specreduce/background.py @@ -5,7 +5,7 @@ import numpy as np from astropy.nddata import NDData -from astropy.units import UnitTypeError +from astropy import units as u from specutils import Spectrum1D from specreduce.core import _ImageParser @@ -258,7 +258,7 @@ def bkg_spectrum(self, image=None): try: return bkg_image.collapse(np.sum, axis=self.crossdisp_axis) - except UnitTypeError: + except u.UnitTypeError: # can't collapse with a spectral axis in pixels because # SpectralCoord only allows frequency/wavelength equivalent units... ext1d = np.sum(bkg_image.flux, axis=self.crossdisp_axis) @@ -280,10 +280,14 @@ def sub_image(self, image=None): """ image = self._parse_image(image) + # a compare_wcs argument is needed for Spectrum1D.subtract() in order to + # avoid a TypeError from SpectralCoord when image's spectral axis is in + # pixels. it is not needed when image's spectral axis has physical units + kwargs = ({'compare_wcs': None} if image.spectral_axis.unit == u.pix + else {}) + # https://docs.astropy.org/en/stable/nddata/mixins/ndarithmetic.html - # (compare_wcs argument needed to avoid TypeError from SpectralCoord - # when image's spectral axis is in pixels) - return image.subtract(self.bkg_image(image), compare_wcs=None) + return image.subtract(self.bkg_image(image), **kwargs) def sub_spectrum(self, image=None): """ @@ -306,7 +310,7 @@ def sub_spectrum(self, image=None): try: return sub_image.collapse(np.sum, axis=self.crossdisp_axis) - except UnitTypeError: + except u.UnitTypeError: # can't collapse with a spectral axis in pixels because # SpectralCoord only allows frequency/wavelength equivalent units... ext1d = np.sum(sub_image.flux, axis=self.crossdisp_axis) @@ -316,4 +320,6 @@ def __rsub__(self, image): """ Subtract the background from an image. """ + # NOTE: will not be called until specutils PR #988 is merged, released, + # and pinned here return self.sub_image(image) diff --git a/specreduce/tests/test_background.py b/specreduce/tests/test_background.py index e85fee69..f36cec71 100644 --- a/specreduce/tests/test_background.py +++ b/specreduce/tests/test_background.py @@ -18,6 +18,9 @@ image[j, ::] *= j image = Spectrum1D(image * u.DN, uncertainty=VarianceUncertainty(np.ones_like(image))) +image_um = Spectrum1D(image.flux, + spectral_axis=np.arange(image.data.shape[1]) * u.um, + uncertainty=VarianceUncertainty(np.ones_like(image.data))) def test_background(): @@ -29,13 +32,22 @@ def test_background(): trace = FlatTrace(image, trace_pos) bkg_sep = 5 bkg_width = 2 - # all the following should be equivalent: + + # all the following should be equivalent, whether image's spectral axis + # is in pixels or physical units: bg1 = Background(image, [trace-bkg_sep, trace+bkg_sep], width=bkg_width) bg2 = Background.two_sided(image, trace, bkg_sep, width=bkg_width) bg3 = Background.two_sided(image, trace_pos, bkg_sep, width=bkg_width) assert np.allclose(bg1.bkg_array, bg2.bkg_array) assert np.allclose(bg1.bkg_array, bg3.bkg_array) + bg4 = Background(image_um, [trace-bkg_sep, trace+bkg_sep], width=bkg_width) + bg5 = Background.two_sided(image_um, trace, bkg_sep, width=bkg_width) + bg6 = Background.two_sided(image_um, trace_pos, bkg_sep, width=bkg_width) + assert np.allclose(bg1.bkg_array, bg4.bkg_array) + assert np.allclose(bg1.bkg_array, bg5.bkg_array) + assert np.allclose(bg1.bkg_array, bg6.bkg_array) + # test that creating a one_sided background works Background.one_sided(image, trace, bkg_sep, width=bkg_width) @@ -51,6 +63,14 @@ def test_background(): # assert np.allclose(sub1.flux, sub2.flux) assert np.allclose(sub2.flux, sub3.flux) + # NOTE: uncomment sub4 test once Spectrum1D and Background subtraction works + # sub4 = image_um - bg4 + sub5 = bg4.sub_image(image_um) + sub6 = bg4.sub_image() + assert np.allclose(sub2.flux, sub5.flux) + # assert np.allclose(sub4.flux, sub5.flux) + assert np.allclose(sub5.flux, sub6.flux) + bkg_spec = bg1.bkg_spectrum() assert isinstance(bkg_spec, Spectrum1D) sub_spec = bg1.sub_spectrum()