diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/tests/test_twodim_base.py | 44 | ||||
| -rw-r--r-- | numpy/lib/twodim_base.py | 62 |
2 files changed, 80 insertions, 26 deletions
diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py index 022c45bd0..9e81cfe4b 100644 --- a/numpy/lib/tests/test_twodim_base.py +++ b/numpy/lib/tests/test_twodim_base.py @@ -286,6 +286,7 @@ def test_tril_triu_ndim2(): yield assert_equal, b.dtype, a.dtype yield assert_equal, c.dtype, a.dtype + def test_tril_triu_ndim3(): for dtype in np.typecodes['AllFloat'] + np.typecodes['AllInteger']: a = np.array([ @@ -324,16 +325,21 @@ def test_mask_indices(): def test_tril_indices(): # indices without and with offset il1 = tril_indices(4) - il2 = tril_indices(4, 2) + il2 = tril_indices(4, k=2) + il3 = tril_indices(4, m=5) + il4 = tril_indices(4, k=2, m=5) a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]) + b = np.arange(1, 21).reshape(4, 5) # indexing: yield (assert_array_equal, a[il1], array([1, 5, 6, 9, 10, 11, 13, 14, 15, 16])) + yield (assert_array_equal, b[il3], + array([1, 6, 7, 11, 12, 13, 16, 17, 18, 19])) # And for assigning values: a[il1] = -1 @@ -342,7 +348,12 @@ def test_tril_indices(): [-1, -1, 7, 8], [-1, -1, -1, 12], [-1, -1, -1, -1]])) - + b[il3] = -1 + yield (assert_array_equal, b, + array([[-1, 2, 3, 4, 5], + [-1, -1, 8, 9, 10], + [-1, -1, -1, 14, 15], + [-1, -1, -1, -1, 20]])) # These cover almost the whole array (two diagonals right of the main one): a[il2] = -10 yield (assert_array_equal, a, @@ -350,21 +361,32 @@ def test_tril_indices(): [-10, -10, -10, -10], [-10, -10, -10, -10], [-10, -10, -10, -10]])) + b[il4] = -10 + yield (assert_array_equal, b, + array([[-10, -10, -10, 4, 5], + [-10, -10, -10, -10, 10], + [-10, -10, -10, -10, -10], + [-10, -10, -10, -10, -10]])) class TestTriuIndices(object): def test_triu_indices(self): iu1 = triu_indices(4) - iu2 = triu_indices(4, 2) + iu2 = triu_indices(4, k=2) + iu3 = triu_indices(4, m=5) + iu4 = triu_indices(4, k=2, m=5) a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]) + b = np.arange(1, 21).reshape(4, 5) # Both for indexing: yield (assert_array_equal, a[iu1], array([1, 2, 3, 4, 6, 7, 8, 11, 12, 16])) + yield (assert_array_equal, b[iu3], + array([1, 2, 3, 4, 5, 7, 8, 9, 10, 13, 14, 15, 19, 20])) # And for assigning values: a[iu1] = -1 @@ -373,6 +395,12 @@ class TestTriuIndices(object): [5, -1, -1, -1], [9, 10, -1, -1], [13, 14, 15, -1]])) + b[iu3] = -1 + yield (assert_array_equal, b, + array([[-1, -1, -1, -1, -1], + [ 6, -1, -1, -1, -1], + [11, 12, -1, -1, -1], + [16, 17, 18, -1, -1]])) # These cover almost the whole array (two diagonals right of the # main one): @@ -382,20 +410,26 @@ class TestTriuIndices(object): [5, -1, -1, -10], [9, 10, -1, -1], [13, 14, 15, -1]])) + b[iu4] = -10 + yield (assert_array_equal, b, + array([[-1, -1, -10, -10, -10], + [6, -1, -1, -10, -10], + [11, 12, -1, -1, -10], + [16, 17, 18, -1, -1]])) class TestTrilIndicesFrom(object): def test_exceptions(self): assert_raises(ValueError, tril_indices_from, np.ones((2,))) assert_raises(ValueError, tril_indices_from, np.ones((2, 2, 2))) - assert_raises(ValueError, tril_indices_from, np.ones((2, 3))) + # assert_raises(ValueError, tril_indices_from, np.ones((2, 3))) class TestTriuIndicesFrom(object): def test_exceptions(self): assert_raises(ValueError, triu_indices_from, np.ones((2,))) assert_raises(ValueError, triu_indices_from, np.ones((2, 2, 2))) - assert_raises(ValueError, triu_indices_from, np.ones((2, 3))) + # assert_raises(ValueError, triu_indices_from, np.ones((2, 3))) class TestVander(object): diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index d168e0fca..5a0c0e7ee 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -11,10 +11,11 @@ __all__ = ['diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'rot90', 'tri', from numpy.core.numeric import ( asanyarray, subtract, arange, zeros, greater_equal, multiply, ones, - asarray, where, + asarray, where, dtype as np_dtype, less ) + def fliplr(m): """ Flip array in the left/right direction. @@ -372,6 +373,7 @@ def tri(N, M=None, k=0, dtype=float): dtype : dtype, optional Data type of the returned array. The default is float. + Returns ------- tri : ndarray of shape (N, M) @@ -393,8 +395,14 @@ def tri(N, M=None, k=0, dtype=float): """ if M is None: M = N - m = greater_equal(subtract.outer(arange(N), arange(M)), -k) - return m.astype(dtype) + + m = greater_equal.outer(arange(N), arange(-k, M-k)) + + # Avoid making a copy if the requested type is already bool + if np_dtype(dtype) != np_dtype(bool): + m = m.astype(dtype) + + return m def tril(m, k=0): @@ -430,8 +438,7 @@ def tril(m, k=0): """ m = asanyarray(m) - out = multiply(tri(m.shape[-2], m.shape[-1], k=k, dtype=m.dtype), m) - return out + return multiply(tri(*m.shape[-2:], k=k, dtype=bool), m, dtype=m.dtype) def triu(m, k=0): @@ -457,8 +464,7 @@ def triu(m, k=0): """ m = asanyarray(m) - out = multiply((1 - tri(m.shape[-2], m.shape[-1], k - 1, dtype=m.dtype)), m) - return out + return multiply(~tri(*m.shape[-2:], k=k-1, dtype=bool), m, dtype=m.dtype) # Originally borrowed from John Hunter and matplotlib @@ -757,17 +763,24 @@ def mask_indices(n, mask_func, k=0): return where(a != 0) -def tril_indices(n, k=0): +def tril_indices(n, k=0, m=None): """ - Return the indices for the lower-triangle of an (n, n) array. + Return the indices for the lower-triangle of an (n, m) array. Parameters ---------- n : int - The row dimension of the square arrays for which the returned + The row dimension of the arrays for which the returned indices will be valid. k : int, optional Diagonal offset (see `tril` for details). + m : int, optional + .. versionadded:: 1.9.0 + + The column dimension of the arrays for which the returned + arrays will be valid. + By default `m` is taken equal to `n`. + Returns ------- @@ -827,7 +840,7 @@ def tril_indices(n, k=0): [-10, -10, -10, -10]]) """ - return mask_indices(n, tril, k) + return where(tri(n, m, k=k, dtype=bool)) def tril_indices_from(arr, k=0): @@ -853,14 +866,14 @@ def tril_indices_from(arr, k=0): .. versionadded:: 1.4.0 """ - if not (arr.ndim == 2 and arr.shape[0] == arr.shape[1]): - raise ValueError("input array must be 2-d and square") - return tril_indices(arr.shape[0], k) + if arr.ndim != 2: + raise ValueError("input array must be 2-d") + return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1]) -def triu_indices(n, k=0): +def triu_indices(n, k=0, m=None): """ - Return the indices for the upper-triangle of an (n, n) array. + Return the indices for the upper-triangle of an (n, m) array. Parameters ---------- @@ -869,6 +882,13 @@ def triu_indices(n, k=0): be valid. k : int, optional Diagonal offset (see `triu` for details). + m : int, optional + .. versionadded:: 1.9.0 + + The column dimension of the arrays for which the returned + arrays will be valid. + By default `m` is taken equal to `n`. + Returns ------- @@ -930,12 +950,12 @@ def triu_indices(n, k=0): [ 12, 13, 14, -1]]) """ - return mask_indices(n, triu, k) + return where(~tri(n, m, k=k-1, dtype=bool)) def triu_indices_from(arr, k=0): """ - Return the indices for the upper-triangle of a (N, N) array. + Return the indices for the upper-triangle of arr. See `triu_indices` for full details. @@ -960,6 +980,6 @@ def triu_indices_from(arr, k=0): .. versionadded:: 1.4.0 """ - if not (arr.ndim == 2 and arr.shape[0] == arr.shape[1]): - raise ValueError("input array must be 2-d and square") - return triu_indices(arr.shape[0], k) + if arr.ndim != 2: + raise ValueError("input array must be 2-d") + return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1]) |