ENH: Introduce np.ma.compress_nd(), generalizes np.ma.compress_rowcols()

Provides a way to supress slices along an abitrary tuple of dimensions.

2 files changed, 180 insertions, 27 deletions

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index d389099ae..51064e831 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -18,8 +18,8 @@ __date__ = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $'
 __all__ = ['apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d',
            'atleast_3d', 'average',
            'clump_masked', 'clump_unmasked', 'column_stack', 'compress_cols',
-           'compress_rowcols', 'compress_rows', 'count_masked', 'corrcoef',
-           'cov',
+           'compress_nd', 'compress_rowcols', 'compress_rows', 'count_masked',
+           'corrcoef', 'cov',
            'diagflat', 'dot', 'dstack',
            'ediff1d',
            'flatnotmasked_contiguous', 'flatnotmasked_edges',
@@ -716,6 +716,57 @@ def median(a, axis=None, out=None, overwrite_input=False):
 
 
 #..............................................................................
+def compress_nd(x, axis=None):
+    """Supress slices from multiple dimensions which contain masked values.
+    
+    Parameters
+    ----------
+    x : array_like, MaskedArray
+        The array to operate on. If not a MaskedArray instance (or if no array
+        elements are masked, `x` is interpreted as a MaskedArray with `mask`
+        set to `nomask`.
+    axis : tuple of ints or int, optional
+        Which dimensions to supress slices from can be configured with this
+        parameter.
+        - If axis is a tuple of ints, those are the axes to supress slices from.
+        - If axis is an int, then that is the only axis to supress slices from.
+        - If axis is None, all axis are selected.
+    
+    Returns
+    -------
+    compress_array : ndarray
+        The compressed array.        
+    """    
+    x = asarray(x)
+    m = getmask(x)
+    # Set axis to tuple of ints
+    if isinstance(axis, tuple):
+        axis = tuple(ax % x.ndim for ax in axis)
+    elif isinstance(axis, int):
+        axis = (axis % x.ndim,)
+    elif axis is None:
+        axis = tuple(range(x.ndim))
+    else:
+        raise ValueError('Invalid type for axis argument')        
+    # Check axis input
+    for ax in axis:
+        if not (0 <= ax < x.ndim):
+            raise ValueError('axis %d is out of range' % ax)
+    if not len(axis) == len(set(axis)):
+        raise ValueError('axis cannot have dupliate entries')
+    # Nothing is masked: return x
+    if m is nomask or not m.any():
+        return x._data
+    # All is masked: return empty
+    if m.all():
+        return nxarray([])
+    # Filter elements through boolean indexing
+    data = x._data
+    for ax in axis:
+        axes = tuple(list(range(ax)) + list(range(ax + 1, x.ndim)))
+        data = data[(slice(None),)*ax + (~m.any(axis=axes),)]
+    return data
+    
 def compress_rowcols(x, axis=None):
     """
     Suppress the rows and/or columns of a 2-D array that contain
@@ -767,26 +818,10 @@ def compress_rowcols(x, axis=None):
            [7, 8]])
 
     """
-    x = asarray(x)
-    if x.ndim != 2:
-        raise NotImplementedError("compress2d works for 2D arrays only.")
-    m = getmask(x)
-    # Nothing is masked: return x
-    if m is nomask or not m.any():
-        return x._data
-    # All is masked: return empty
-    if m.all():
-        return nxarray([])
-    # Builds a list of rows/columns indices
-    (idxr, idxc) = (list(range(len(x))), list(range(x.shape[1])))
-    masked = m.nonzero()
-    if not axis:
-        for i in np.unique(masked[0]):
-            idxr.remove(i)
-    if axis in [None, 1, -1]:
-        for j in np.unique(masked[1]):
-            idxc.remove(j)
-    return x._data[idxr][:, idxc]
+    if asarray(x).ndim != 2:
+        raise NotImplementedError("compress_rowcols works for 2D arrays only.")
+    return compress_nd(x, axis=axis)
+
 
 def compress_rows(a):
     """
@@ -800,6 +835,9 @@ def compress_rows(a):
     extras.compress_rowcols
 
     """
+    a = asarray(a)
+    if a.ndim != 2:
+        raise NotImplementedError("compress_rows works for 2D arrays only.")
     return compress_rowcols(a, 0)
 
 def compress_cols(a):
@@ -814,6 +852,9 @@ def compress_cols(a):
     extras.compress_rowcols
 
     """
+    a = asarray(a)
+    if a.ndim != 2:
+        raise NotImplementedError("compress_cols works for 2D arrays only.")
     return compress_rowcols(a, 1)
 
 def mask_rowcols(a, axis=None):
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index ee8e6bc18..b6749ae9e 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -29,7 +29,7 @@ from numpy.ma.extras import (
     cov, corrcoef, median, average,
     unique, setxor1d, setdiff1d, union1d, intersect1d, in1d, ediff1d,
     apply_over_axes, apply_along_axis,
-    compress_rowcols, mask_rowcols,
+    compress_nd, compress_rowcols, mask_rowcols,
     clump_masked, clump_unmasked,
     flatnotmasked_contiguous, notmasked_contiguous, notmasked_edges,
     masked_all, masked_all_like)
@@ -347,10 +347,122 @@ class TestNotMasked(TestCase):
         assert_equal(tmp[2][-2], slice(0, 6, None))
 
 
-class Test2DFunctions(TestCase):
-    # Tests 2D functions
-    def test_compress2d(self):
-        # Tests compress2d
+class TestCompressFunctions(TestCase):
+
+    def test_compress_nd(self):
+        # Tests compress_nd
+        x = np.array(list(range(3*4*5))).reshape(3, 4, 5)
+        m = np.zeros((3,4,5)).astype(bool)
+        m[1,1,1] = True
+        x = array(x, mask=m)
+
+        # axis=None
+        a = compress_nd(x)
+        assert_equal(a, [[[ 0,  2,  3 , 4],
+                          [10, 12, 13, 14],
+                          [15, 17, 18, 19]],
+                         [[40, 42, 43, 44],
+                          [50, 52, 53, 54],
+                          [55, 57, 58, 59]]])
+        
+        # axis=0
+        a = compress_nd(x, 0)
+        assert_equal(a, [[[ 0,  1,  2,  3,  4],
+                          [ 5,  6,  7,  8,  9],
+                          [10, 11, 12, 13, 14],
+                          [15, 16, 17, 18, 19]],
+                         [[40, 41, 42, 43, 44],
+                          [45, 46, 47, 48, 49],
+                          [50, 51, 52, 53, 54],
+                          [55, 56, 57, 58, 59]]])
+
+        # axis=1
+        a = compress_nd(x, 1)
+        assert_equal(a, [[[ 0,  1,  2,  3,  4],
+                          [10, 11, 12, 13, 14],
+                          [15, 16, 17, 18, 19]],
+                         [[20, 21, 22, 23, 24],
+                          [30, 31, 32, 33, 34],
+                          [35, 36, 37, 38, 39]],
+                         [[40, 41, 42, 43, 44],
+                          [50, 51, 52, 53, 54],
+                          [55, 56, 57, 58, 59]]])
+
+        a2 = compress_nd(x, (1,))
+        a3 = compress_nd(x, -2)
+        a4 = compress_nd(x, (-2,))
+        assert_equal(a, a2)
+        assert_equal(a, a3)
+        assert_equal(a, a4)
+
+        # axis=2
+        a = compress_nd(x, 2)
+        assert_equal(a, [[[ 0, 2,  3,  4],
+                          [ 5, 7,  8,  9],
+                          [10, 12, 13, 14],
+                          [15, 17, 18, 19]],
+                         [[20, 22, 23, 24],
+                          [25, 27, 28, 29],
+                          [30, 32, 33, 34],
+                          [35, 37, 38, 39]],
+                         [[40, 42, 43, 44],
+                          [45, 47, 48, 49],
+                          [50, 52, 53, 54],
+                          [55, 57, 58, 59]]])
+
+        a2 = compress_nd(x, (2,))
+        a3 = compress_nd(x, -1)
+        a4 = compress_nd(x, (-1,))
+        assert_equal(a, a2)
+        assert_equal(a, a3)
+        assert_equal(a, a4)
+
+        # axis=(0, 1)
+        a = compress_nd(x, (0, 1))
+        assert_equal(a, [[[ 0,  1,  2,  3,  4],
+                          [10, 11, 12, 13, 14],
+                          [15, 16, 17, 18, 19]],
+                         [[40, 41, 42, 43, 44],
+                          [50, 51, 52, 53, 54],
+                          [55, 56, 57, 58, 59]]])
+        a2 = compress_nd(x, (0, -2))
+        assert_equal(a, a2)
+
+        # axis=(1, 2)
+        a = compress_nd(x, (1, 2))
+        assert_equal(a, [[[ 0,  2,  3,  4],
+                          [10, 12, 13, 14],
+                          [15, 17, 18, 19]],
+                         [[20, 22, 23, 24],
+                          [30, 32, 33, 34],
+                          [35, 37, 38, 39]],
+                         [[40, 42, 43, 44],
+                          [50, 52, 53, 54],
+                          [55, 57, 58, 59]]])
+
+        a2 = compress_nd(x, (-2, 2))
+        a3 = compress_nd(x, (1, -1))
+        a4 = compress_nd(x, (-2, -1))
+        assert_equal(a, a2)
+        assert_equal(a, a3)
+        assert_equal(a, a4)
+
+        # axis=(0, 2)
+        a = compress_nd(x, (0, 2))
+        assert_equal(a, [[[ 0,   2,  3,  4],
+                          [ 5,   7,  8,  9],
+                          [10, 12, 13, 14],
+                          [15, 17, 18, 19]],
+                         [[40, 42, 43, 44],
+                          [45, 47, 48, 49],
+                          [50, 52, 53, 54],
+                          [55, 57, 58, 59]]])
+
+        a2 = compress_nd(x, (0, -1))
+        assert_equal(a, a2)
+
+    def test_compress_rowcols(self):
+        # Tests compress_rowcols
         x = array(np.arange(9).reshape(3, 3),
                   mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
         assert_equal(compress_rowcols(x), [[4, 5], [7, 8]])