Skip to content

Commit

Permalink
ENH: Introduce np.ma.compress_nd(), generalizes np.ma.compress_rowcols()
Browse files Browse the repository at this point in the history 
Provides a way to supress slices along an abitrary tuple of dimensions.
  • Loading branch information
@ddasilva
ddasilva committed May 4, 2015
1 parent 147c60f commit 883d052
Show file tree
Hide file tree
Showing 2 changed files with 180 additions and 27 deletions.
85 changes: 63 additions & 22 deletions numpy/ma/extras.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@
__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',
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
"""
Expand All @@ -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):
Expand All @@ -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):
Expand Down
122 changes: 117 additions & 5 deletions numpy/ma/tests/test_extras.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
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)
Expand Down Expand Up @@ -347,10 +347,122 @@ def test_contiguous(self):
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]])
Expand Down

0 comments on commit 883d052

Please sign in to comment.