,

specreduce/extract.py

@@ -183,6 +183,9 @@ class BoxcarExtract(SpecreduceOperation):
     crossdisp_axis: int = 0
     # TODO: should disp_axis and crossdisp_axis be defined in the Trace object?
 
+    mask_treatment : str = 'filter'
+    _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
+
     @property
     def spectrum(self):
         return self.__call__()
@@ -204,6 +207,20 @@ def __call__(self, image=None, trace_object=None, width=None,
             dispersion axis [default: 1]
         crossdisp_axis : int, optional
             cross-dispersion axis [default: 0]
+        mask_treatment : string, optional
+            The method for handling masked or non-finite data. Choice of `filter`,
+            `omit`, or `zero-fill`. If `filter` is chosen, masked/non-finite data
+            will be filtered during the fit to each bin/column (along disp. axis) to
+            find the peak. If `omit` is chosen, columns along disp_axis with any
+            masked/non-finite data values will be fully masked (i.e, 2D mask is
+            collapsed to 1D and applied). If `zero-fill` is chosen, masked/non-finite
+            data will be replaced with 0.0 in the input image, and the mask will then
+            be dropped. For all three options, the input mask (optional on input
+            NDData object) will be combined with a mask generated from any non-finite
+            values in the image data. Also note that because binning is an option in
+            FitTrace, that masked data will contribute zero to the sum when binning
+            adjacent columns.
+            [default: ``filter``]
 
 
         Returns
@@ -218,7 +235,8 @@ def __call__(self, image=None, trace_object=None, width=None,
         disp_axis = disp_axis if disp_axis is not None else self.disp_axis
         crossdisp_axis = crossdisp_axis if crossdisp_axis is not None else self.crossdisp_axis
 
-        # handle image processing based on its type
+        # Parse image, including masked/nonfinite data handling based on
+        # choice of `mask_treatment`. returns a Spectrum1D
         self.image = self._parse_image(image)
 
         # TODO: this check can be removed if/when implemented as a check in FlatTrace
@@ -236,6 +254,9 @@ def __call__(self, image=None, trace_object=None, width=None,
             disp_axis,
             crossdisp_axis,
             self.image.shape)
+        import matplotlib.pyplot as plt
+        plt.imshow(wimg)
+        plt.show()
 
         # extract, assigning no weight to non-finite pixels outside the window
         # (non-finite pixels inside the window will still make it into the sum)

specreduce/tests/test_background.py

@@ -164,6 +164,21 @@ def mk_img(self, nrows=4, ncols=5, nan_slices=None):
 
         return img * u.DN
 
+    def test_fully_masked_column(self):
+        """
+        Test what happens when a full column is masked, not the entire
+        image. In this case, the background value for that fully-masked
+        column should be 0.0, with no error or warning raised.
+        """
+
+        img = np.ones((12, 12))
+        img[:, 0:1] = np.nan
+
+        bkg = Background(img, traces=FlatTrace(img, 6))
+
+        assert np.all(bkg.bkg_image().data[:, 0:1] == 0.0)
+
+
     def test_fully_masked(self):
         """
         Test that the appropriate error is raised by `Background` when image
@@ -241,6 +256,62 @@ def test_mask_treatment_bkg_img_spectrum(self, method, expected):
                                        np.tile(expected, (img_size, 1)))
 
             # test background spectrum matches 'expected' times the number of rows
-            # since this is a sum
+            # in cross disp axis, since this is a sum and all values in a col are
+            # the same.
             bk_spec = background.bkg_spectrum()
             np.testing.assert_allclose(bk_spec.flux.value, expected * img_size)
+
+    def test_sub_bkg_image(self):
+        """
+        Test that masked and nonfinite data is handled correctly when subtracting
+        background from image, for all currently implemented masking
+        options ('filter', 'omit', and 'zero-fill'). 
+        """
+
+        # make image, set some value to nan, which will be masked in the function
+        image = self.mk_img(nrows=12, ncols=12,
+                            nan_slices=[np.s_[5:10, 0], np.s_[7:12, 3],
+                                        np.s_[2, 7]])
+
+        # Calculate a background value using mask_treatment = 'filter'.
+        # For 'filter', the flag applies to how masked values are handled during
+        # calculation of background for each column, but nonfinite data will
+        # remain in input data array
+        background_filter = Background(image, mask_treatment='filter',
+                                       traces=FlatTrace(image, 6),
+                                       width=2)
+        subtracted_img_filter = background_filter.sub_image()
+
+        assert np.all(np.isfinite(subtracted_img_filter.data) == np.isfinite(image.data))
+
+        # Calculate a background value using mask_treatment = 'omit'. The input
+        # 2d mask is reduced to a 1d mask to mask out full columns in the
+        # presence of any nans - this means that (as tested above in
+        # `test_mask_treatment_bkg_img_spectrum`) those columns will have 0.0
+        # background. In this case, image.mask is expanded to mask full 
+        # columns - the image itself will not have full columns set to np.nan,
+        # so there are still valid background subtracted data values in this
+        # case, but the corresponding mask for that entire column will be masked.
+
+        background_omit = Background(image, mask_treatment='omit',
+                                     traces=FlatTrace(image, 6),
+                                     width=2)
+        subtracted_img_omit = background_omit.sub_image()
+
+        assert np.all(np.isfinite(subtracted_img_omit.data) == np.isfinite(image.data))
+
+        # Calculate a background value using mask_treatment = 'zero-fill'. Data
+        # values at masked locations are set to 0 in the image array, and the
+        # background value calculated for that column will be subtracted
+        # resulting in a negative value. The resulting background subtracted
+        # image should be fully finite and the mask should be zero everywhere
+        # (all unmasked)
+
+        background_zero_fill = Background(image, mask_treatment='zero-fill',
+                                          traces=FlatTrace(image, 6),
+                                          width=2)
+        subtracted_img_zero_fill = background_zero_fill.sub_image()
+
+        assert np.all(np.isfinite(subtracted_img_zero_fill.data))
+        assert np.all(subtracted_img_zero_fill.mask == 0)
+

specreduce/tests/test_extract.py

@@ -15,7 +15,8 @@ def add_gaussian_source(image, amps=2, stddevs=2, means=None):
 
     """ Modify `image.data` to add a horizontal spectrum across the image.
         Each column can have a different amplitude, stddev or mean position
-        if these are arrays (otherwise, constant across image)."""
+        if these are arrays (otherwise, constant across image).
+    """
 
     nrows, ncols = image.shape
 

specreduce/tests/test_tracing.py

@@ -241,9 +241,8 @@ def test_fit_trace_all_nan_cols(self, mask_treatment):
         truth = [2.5318835, 2.782069, 3.0322546, 3.2824402, 3.5326257,
                  3.7828113, 4.0329969, 4.2831824, 4.533368, 4.7835536,
                  5.0337391]
-        max_trace = FitTrace(img, peak_method='centroid')
-        np.testing.assert_allclose(truth, max_trace.trace,
-                                   mask_treatment=mask_treatment)
+        max_trace = FitTrace(img, peak_method='centroid', mask_treatment=mask_treatment)
+        np.testing.assert_allclose(truth, max_trace.trace)
 
 
     def test_warn_msg_fit_trace_all_nan_cols(self):