Merge pull request #159 from ojustino/background-handle-nans

Prevent (most) NaN propagation in `Background` and `BoxcarExtract`

Author: ojustino

CHANGES.rst

@@ -10,6 +10,14 @@ API Changes
 Bug Fixes
 ^^^^^^^^^
 
+- Output 1D spectra from Background no longer include NaNs. Output 1D
+spectra from BoxcarExtract no longer include NaNs when none are present
+in the extraction window. NaNs in the window will still propagate to
+BoxcarExtract's extracted 1D spectrum. [#159]
+
+- Backgrounds using median statistic properly ignore zero-weighted pixels
+[#159]
+
 
 1.3.0 (2022-12-05)
 ------------------
@@ -48,6 +56,7 @@ Bug Fixes
   function after change to upstream API. This will make specreduce
   be compatible with numpy 1.24 or later. [#155]
 
+
 1.2.0 (2022-10-04)
 ------------------
 

specreduce/background.py

@@ -133,14 +133,20 @@ def _to_trace(trace):
 
         self.bkg_wimage = bkg_wimage
 
+        # mask user-highlighted and invalid values (if any) before taking stats
+        or_mask = (np.logical_or(~np.isfinite(self.image.data), self.image.mask)
+                   if self.image.mask is not None
+                   else ~np.isfinite(self.image.data))
+
         if self.statistic == 'average':
-            self._bkg_array = np.average(self.image.data,
-                                         weights=self.bkg_wimage,
-                                         axis=self.crossdisp_axis)
+            image_ma = np.ma.masked_array(self.image.data, mask=or_mask)
+            self._bkg_array = np.ma.average(image_ma,
+                                            weights=self.bkg_wimage,
+                                            axis=self.crossdisp_axis).data
         elif self.statistic == 'median':
-            med_image = self.image.data.copy()
-            med_image[np.where(self.bkg_wimage) == 0] = np.nan
-            self._bkg_array = np.nanmedian(med_image, axis=self.crossdisp_axis)
+            med_mask = np.logical_or(self.bkg_wimage == 0, or_mask)
+            image_ma = np.ma.masked_array(self.image.data, mask=med_mask)
+            self._bkg_array = np.ma.median(image_ma, axis=self.crossdisp_axis).data
         else:
             raise ValueError("statistic must be 'average' or 'median'")
 
@@ -232,7 +238,7 @@ def bkg_image(self, image=None):
 
         Returns
         -------
-        Spectrum1D object with same shape as ``image``.
+        `~specutils.Spectrum1D` object with same shape as ``image``.
         """
         image = self._parse_image(image)
         return Spectrum1D(np.tile(self._bkg_array,
@@ -261,11 +267,11 @@ def bkg_spectrum(self, image=None):
         bkg_image = self.bkg_image(image)
 
         try:
-            return bkg_image.collapse(np.sum, axis=self.crossdisp_axis)
+            return bkg_image.collapse(np.nansum, 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.sum(bkg_image.flux, axis=self.crossdisp_axis)
+            ext1d = np.nansum(bkg_image.flux, axis=self.crossdisp_axis)
             return Spectrum1D(ext1d, bkg_image.spectral_axis)
 
     def sub_image(self, image=None):
@@ -280,7 +286,7 @@ def sub_image(self, image=None):
 
         Returns
         -------
-        array with same shape as ``image``
+        `~specutils.Spectrum1D` object with same shape as ``image``.
         """
         image = self._parse_image(image)
 
@@ -313,11 +319,11 @@ def sub_spectrum(self, image=None):
         sub_image = self.sub_image(image=image)
 
         try:
-            return sub_image.collapse(np.sum, axis=self.crossdisp_axis)
+            return sub_image.collapse(np.nansum, 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.sum(sub_image.flux, axis=self.crossdisp_axis)
+            ext1d = np.nansum(sub_image.flux, axis=self.crossdisp_axis)
             return Spectrum1D(ext1d, spectral_axis=sub_image.spectral_axis)
 
     def __rsub__(self, image):

specreduce/extract.py

@@ -236,8 +236,10 @@ def __call__(self, image=None, trace_object=None, width=None,
             crossdisp_axis,
             self.image.shape)
 
-        # extract
-        ext1d = np.sum(self.image.data * wimg, axis=crossdisp_axis)
+        # extract, assigning no weight to non-finite pixels outside the window
+        # (non-finite pixels inside the window will still make it into the sum)
+        image_windowed = np.where(wimg, self.image.data*wimg, 0)
+        ext1d = np.sum(image_windowed, axis=crossdisp_axis)
         return Spectrum1D(ext1d * self.image.unit,
                           spectral_axis=self.image.spectral_axis)
 

specreduce/tests/test_background.py

@@ -13,11 +13,11 @@
 # Test image is comprised of 30 rows with 10 columns each. Row content
 # is row index itself. This makes it easy to predict what should be the
 # value extracted from a region centered at any arbitrary Y position.
-image = np.ones(shape=(30, 10))
-for j in range(image.shape[0]):
-    image[j, ::] *= j
-image = Spectrum1D(image * u.DN,
-                   uncertainty=VarianceUncertainty(np.ones_like(image)))
+img = np.ones(shape=(30, 10))
+for j in range(img.shape[0]):
+    img[j, ::] *= j
+image = Spectrum1D(img * u.DN,
+                   uncertainty=VarianceUncertainty(np.ones_like(img)))
 image_um = Spectrum1D(image.flux,
                       spectral_axis=np.arange(image.data.shape[1]) * u.um,
                       uncertainty=VarianceUncertainty(np.ones_like(image.data)))
@@ -28,7 +28,7 @@ def test_background():
     # Try combinations of extraction center, and even/odd
     # extraction aperture sizes.
     #
-    trace_pos = 15.0
+    trace_pos = 15
     trace = FlatTrace(image, trace_pos)
     bkg_sep = 5
     bkg_width = 2
@@ -72,10 +72,25 @@ def test_background():
     assert isinstance(bkg_spec, Spectrum1D)
     sub_spec = bg1.sub_spectrum()
     assert isinstance(sub_spec, Spectrum1D)
+
     # test that width==0 results in no background
     bg = Background.two_sided(image, trace, bkg_sep, width=0)
     assert np.all(bg.bkg_image().flux == 0)
 
+    # test that any NaNs in input image (whether in or outside the window) don't
+    # propagate to _bkg_array (which affects bkg_image and sub_image methods) or
+    # the final 1D spectra.
+    img[0, 0] = np.nan  # out of window
+    img[trace_pos, 0] = np.nan  # in window
+    stats = ['average', 'median']
+
+    for st in stats:
+        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.sub_spectrum().flux).sum() == 0
+
 
 def test_warnings_errors():
     # image.shape (30, 10)