Merge pull request #1 from hpparvi/masking_options

Masking options

specreduce/background.py

@@ -49,7 +49,7 @@ class Background(_ImageParser):
         cross-dispersion axis
         [default: 0]
     mask_treatment : string, optional
-        The method for handling masked or non-finite data. Choice of `filter`,
+        The method for handling masked or non-finite data. Choice of ``filter``,
         ``omit``, or ``zero-fill``. If ``filter`` is chosen, masked and non-finite
         data will not contribute to the background statistic that is calculated
         in each column along `disp_axis`. If `omit` is chosen, columns along
@@ -114,7 +114,8 @@ def __post_init__(self):
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. Any uncaught nonfinte data values will be
         # masked as well. Returns a Spectrum1D.
-        self.image = self._parse_image(self.image)
+        self.image = self._parse_image(self.image, disp_axis=self.disp_axis,
+                                       mask_treatment=self.mask_treatment)
 
         # always work with masked array, even if there is no masked
         # or nonfinite data, in case padding is needed. if not, mask will be

specreduce/core.py

@@ -28,22 +28,42 @@ class _ImageParser:
         - `~numpy.ndarray`
     """
 
-    def _parse_image(self, image, disp_axis=1):
+    # The '_valid_mask_treatment_methods' in the Background, Trace, and Extract
+    # classes is a subset of implemented methods.
+    implemented_mask_treatment_methods = 'filter', 'zero-fill', 'omit'
+
+    def _parse_image(self, image,
+                     disp_axis: int = 1,
+                     mask_treatment: str = 'filter') -> Spectrum1D:
         """
-        Convert all accepted image types to a consistently formatted
-        Spectrum1D object.
+        Convert all accepted image types to a consistently formatted Spectrum1D object.
 
         Parameters
         ----------
-        image : `~astropy.nddata.NDData`-like or array-like, required
+        image : `~astropy.nddata.NDData`-like or array-like
             The image to be parsed. If None, defaults to class' own
             image attribute.
-        disp_axis : int, optional
+        disp_axis
             The index of the image's dispersion axis. Should not be
             changed until operations can handle variable image
             orientations. [default: 1]
+        mask_treatment
+            The method for handling masked or non-finite data. Choice of ``filter``,
+            ``omit``, or ``zero-fill``. If ``filter`` is chosen, masked and non-finite
+            data will not contribute to the background statistic that is calculated
+            in each column along `disp_axis`. 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.
+            [default: ``filter``]
+
+        Returns
+        -------
+        Spectrum1D
         """
-
         # would be nice to handle (cross)disp_axis consistently across
         # operations (public attribute? private attribute? argument only?) so
         # it can be called from self instead of via kwargs...
@@ -52,14 +72,35 @@ def _parse_image(self, image, disp_axis=1):
             # useful for Background's instance methods
             return self.image
 
-        img = self._get_data_from_image(image, disp_axis=disp_axis)
-
-        return img
+        return self._get_data_from_image(image, disp_axis=disp_axis,
+                                         mask_treatment=mask_treatment)
 
     @staticmethod
-    def _get_data_from_image(image, disp_axis=1):
-        """Extract data array from various input types for `image`.
-           Retruns `np.ndarray` of image data."""
+    def _get_data_from_image(image,
+                             disp_axis: int = 1,
+                             mask_treatment: str = 'filter') -> Spectrum1D:
+        """
+        Extract data array from various input types for `image`.
+
+        Parameters
+        ----------
+        image : array-like or Quantity
+            Input image from which data is extracted. This can be a 2D numpy
+            array, Quantity, or an NDData object.
+        disp_axis : int, optional
+            The dispersion axis of the image.
+        mask_treatment : str, optional
+            Treatment method for the mask:
+            - 'filter' (default): Return the unmodified input image and combined mask.
+            - 'zero-fill': Set masked values in the image to zero.
+            - 'omit': Mask all pixels along the cross dispersion axis if any value is masked.
+
+        Returns
+        -------
+        Spectrum1D
+        """
+        # This works only with 2D images.
+        crossdisp_axis = (disp_axis + 1) % 2
 
         if isinstance(image, u.quantity.Quantity):
             img = image.value
@@ -77,12 +118,15 @@ def _get_data_from_image(image, disp_axis=1):
         # handled as well. Note that input data may be modified if a fill value
         # is chosen to handle masked data. The returned image will always have
         # `image.mask` even if there are no nonfinte or masked values.
-        img, mask = self._mask_and_nonfinite_data_handling(image=img, mask=mask)
+        img, mask = _ImageParser._mask_and_nonfinite_data_handling(image=img,
+                                                                   mask=mask,
+                                                                   mask_treatment=mask_treatment,
+                                                                   crossdisp_axis=crossdisp_axis)
 
         # mask (handled above) and uncertainty are set as None when they aren't
         # specified upon creating a Spectrum1D object, so we must check whether
         # these attributes are absent *and* whether they are present but set as None
-        if getattr(image, 'uncertainty', None) is not None:
+        if hasattr(image, 'uncertainty'):
             uncertainty = image.uncertainty
         else:
             uncertainty = VarianceUncertainty(np.ones(img.shape))
@@ -94,26 +138,14 @@ def _get_data_from_image(image, disp_axis=1):
 
         img = Spectrum1D(img * unit, spectral_axis=spectral_axis,
                          uncertainty=uncertainty, mask=mask)
-
         return img
 
     @staticmethod
-    def _get_data_from_image(image):
-        """Extract data array from various input types for `image`.
-           Retruns `np.ndarray` of image data."""
-
-        if isinstance(image, u.quantity.Quantity):
-            img = image.value
-        if isinstance(image, np.ndarray):
-            img = image
-        else:  # NDData, including CCDData and Spectrum1D
-            img = image.data
-        return img
-
-    def _mask_and_nonfinite_data_handling(self, image, mask):
+    def _mask_and_nonfinite_data_handling(image, mask,
+                                          mask_treatment: str = 'filter',
+                                          crossdisp_axis: int = 0) -> tuple[np.ndarray, np.ndarray]:
         """
-        This function handles the treatment of masked and nonfinite data,
-        including input validation.
+        Handle the treatment of masked and nonfinite data.
 
         All operations in Specreduce can take in a mask for the data as
         part of the input NDData. Additionally, any non-finite values in the
@@ -124,31 +156,34 @@ def _mask_and_nonfinite_data_handling(self, image, mask):
         of masked and nonfinite data - filter, omit, and zero-fill.
         Depending on the step, all or a subset of these three options are valid.
 
+        Parameters
+        ----------
+        image : array-like
+            The input image data array that may contain nonfinite values.
+        mask : array-like of bool or None
+            An optional Boolean mask array. Nonfinite values in the image will be added
+            to this mask.
+        mask_treatment : str
+            Specifies how to handle masked data:
+            - 'filter' (default): Returns the unmodified input image and combined mask.
+            - 'zero-fill': Sets masked values in the image to zero.
+            - 'omit': Masks entire columns or rows if any value is masked.
+        crossdisp_axis : int
+            Axis along which to collapse the 2D mask into a 1D mask for treatment 'omit'.
         """
-
-        # valid options depend on Specreduce step, and are set accordingly there
-        # for steps that this isn't implemeted for yet, default to 'filter',
-        # which will return unmodified input image and mask
-        mask_treatment = getattr(self, 'mask_treatment', 'filter')
-
-        # make sure chosen option is valid. if _valid_mask_treatment_methods
-        # is not an attribue, proceed with 'filter' to return back inupt data
-        # and mask that is combined with nonfinite data.
-        if mask_treatment is not None:  # None in operations where masks aren't relevant (FlatTrace)
-            valid_mask_treatment_methods = getattr(self, '_valid_mask_treatment_methods', ['filter'])  # noqa
-            if mask_treatment not in valid_mask_treatment_methods:
-                raise ValueError(f"`mask_treatment` must be one of {valid_mask_treatment_methods}")
+        if mask_treatment not in _ImageParser.implemented_mask_treatment_methods:
+            raise ValueError("`mask_treatment` must be one of "
+                             f"{_ImageParser.implemented_mask_treatment_methods}")
 
         # make sure there is always a 'mask', even when all data is unmasked and finite.
         if mask is not None:
-            mask = self.image.mask
             # always mask any previously uncaught nonfinite values in image array
             # combining these with the (optional) user-provided mask on `image.mask`
             mask = np.logical_or(mask, ~np.isfinite(image))
         else:
             mask = ~np.isfinite(image)
 
-        # if mask option is the default 'filter' option, or None,
+        # if mask option is the default 'filter' option,
         # nothing needs to be done. input mask (combined with nonfinite data)
         # remains with data as-is.
 
@@ -165,23 +200,16 @@ def _mask_and_nonfinite_data_handling(self, image, mask):
 
             # masked array with no masked values, so accessing image.mask works
             # but we don't need the actual mask anymore since data has been set to 0
-            mask = np.zeros(image.shape)
+            mask = np.zeros(image.shape, dtype=bool)
 
         elif mask_treatment == 'omit':
             # collapse 2d mask (after accounting for addl non-finite values in
             # data) to a 1d mask, along dispersion axis, to fully mask columns
             # that have any masked values.
 
-            # must have a crossdisp_axis specified to use 'omit' optoin
-            if hasattr(self, 'crossdisp_axis'):
-                crossdisp_axis = self.crossdisp_axis
-            if hasattr(self, '_crossdisp_axis'):
-                crossdisp_axis = self._crossdisp_axis
-
             # create a 1d mask along crossdisp axis - if any column has a single nan,
             # the entire column should be masked
-            reduced_mask = np.logical_or.reduce(mask,
-                                                axis=crossdisp_axis)
+            reduced_mask = np.logical_or.reduce(mask, axis=crossdisp_axis)
 
             # back to a 2D mask
             shape = (image.shape[0], 1) if crossdisp_axis == 0 else (1, image.shape[1])
@@ -193,6 +221,7 @@ def _mask_and_nonfinite_data_handling(self, image, mask):
 
         return image, mask
 
+
 @dataclass
 class SpecreduceOperation(_ImageParser):
     """

specreduce/extract.py

@@ -213,7 +213,8 @@ def __call__(self, image=None, trace_object=None, width=None,
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`, which for BoxcarExtract can be filter, zero-fill, or
         # omit. non-finite data will be masked, always. Returns a Spectrum1D.
-        self.image = self._parse_image(image)
+        self.image = self._parse_image(image, disp_axis=disp_axis,
+                                       mask_treatment=self.mask_treatment)
 
         # # _parse_image returns a Spectrum1D. convert this to a masked array
         # # for ease of calculations here (even if there is no masked data).

specreduce/tests/test_tracing.py

@@ -252,7 +252,7 @@ def test_array_trace_masking(self):
 
         # ensure correct warning is raised when entire trace is masked.
         trace_arr = np.ma.MaskedArray([1, 2, np.nan, 4, 5], mask=[1, 1, 0, 1, 1])
-        with pytest.raises(UserWarning, match=r'Entire trace array is masked.'):
+        with pytest.warns(UserWarning, match=r'Entire trace array is masked.'):
             array_trace = ArrayTrace(img, trace_arr)
 
     def test_fit_trace_fully_masked_image(self):

specreduce/tracing.py

@@ -287,7 +287,8 @@ def __post_init__(self):
 
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. returns a Spectrum1D
-        self.image = self._parse_image(self.image)
+        self.image = self._parse_image(self.image, disp_axis=self._disp_axis,
+                                       mask_treatment=self.mask_treatment)
 
         # _parse_image returns a Spectrum1D. convert this to a masked array
         # for ease of calculations here (even if there is no masked data).