Fix Background.sub_image for physical wavelengths

Author: ojustino

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)

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()