Add unmasked_channels property to cache any(axis=(1,2)) calls on masks

radio-astro-tools · Nov 11, 2020 · 0d1f6f2 · 0d1f6f2
1 parent 2dead55
commit 0d1f6f2
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/spectral_cube/base_class.py b/spectral_cube/base_class.py
@@ -346,6 +346,13 @@ def spectral_axis(self):
         # spectral objects should be forced to implement this
         raise NotImplementedError
 
+    @property
+    @cached
+    def unmasked_channels(self):
+        if isinstance(self._data, da.Array):
+            return self._compute(da.any(self._mask_include, axis=(1, 2)))
+
+        return self.mask.include().any(axis=(1, 2))
 
 class MaskableArrayMixinClass(object):
     """