Add a with_mask for LDOs

spectral_cube/lower_dimensional_structures.py

@@ -17,7 +17,8 @@
 from .utils import SliceWarning, BeamWarning, SmoothingWarning, FITSWarning
 from .cube_utils import convert_bunit
 from . import wcs_utils
-from .masks import BooleanArrayMask, MaskBase
+from .masks import (LazyMask, LazyComparisonMask, BooleanArrayMask, MaskBase,
+                    is_broadcastable_and_smaller)
 
 from .base_class import (BaseNDClass, SpectralAxisMixinClass,
                          SpatialCoordMixinClass, MaskableArrayMixinClass,
@@ -301,6 +302,55 @@ def _initial_set_mask(self, mask):
 
         self._mask = mask
 
+    def with_mask(self, mask, inherit_mask=True, wcs_tolerance=None):
+        """
+        Return a new LowerDimensionalObject instance that contains a composite
+        mask of the current LDO and the new ``mask``.  Values of the mask that
+        are ``True`` will be *included* (masks are analogous to numpy boolean
+        index arrays, they are the inverse of the ``.mask`` attribute of a numpy
+        masked array).
+
+        Parameters
+        ----------
+        mask : :class:`MaskBase` instance, or boolean numpy array
+            The mask to apply. If a boolean array is supplied,
+            it will be converted into a mask, assuming that
+            `True` values indicate included elements.
+
+        inherit_mask : bool (optional, default=True)
+            If True, combines the provided mask with the
+            mask currently attached to the cube
+
+        wcs_tolerance : None or float
+            The tolerance of difference in WCS parameters between the cube and
+            the mask.  Defaults to `self._wcs_tolerance` (which itself defaults
+            to 0.0) if unspecified
+
+        Returns
+        -------
+        new_ldo : :class:`LowerDimensionalObject`
+            A LDO with the new mask applied.
+
+        Notes
+        -----
+        This operation returns a view into the data, and not a copy.
+        """
+        if isinstance(mask, np.ndarray):
+            if not is_broadcastable_and_smaller(mask.shape, self._data.shape):
+                raise ValueError("Mask shape is not broadcastable to data shape: "
+                                 "%s vs %s" % (mask.shape, self._data.shape))
+            mask = BooleanArrayMask(mask, self._wcs, shape=self._data.shape)
+
+        if self._mask is not None and inherit_mask:
+            new_mask = self._mask & mask
+        else:
+            new_mask = mask
+
+        new_mask._validate_wcs(new_data=self._data, new_wcs=self._wcs,
+                               wcs_tolerance=wcs_tolerance or self._wcs_tolerance)
+
+        return self._new_thing_with(mask=new_mask, wcs_tolerance=wcs_tolerance)
+
 
 class Projection(LowerDimensionalObject, SpatialCoordMixinClass,
                  MaskableArrayMixinClass, BeamMixinClass):