diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/tests/test_twodim_base.py | 34 | ||||
| -rw-r--r-- | numpy/lib/twodim_base.py | 4 |
2 files changed, 31 insertions, 7 deletions
diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py index 8d0275a25..022c45bd0 100644 --- a/numpy/lib/tests/test_twodim_base.py +++ b/numpy/lib/tests/test_twodim_base.py @@ -275,16 +275,40 @@ class TestTri(TestCase): assert_array_equal(tri(3, dtype=bool), out.astype(bool)) -def test_tril_triu(): +def test_tril_triu_ndim2(): for dtype in np.typecodes['AllFloat'] + np.typecodes['AllInteger']: a = np.ones((2, 2), dtype=dtype) b = np.tril(a) c = np.triu(a) - assert_array_equal(b, [[1, 0], [1, 1]]) - assert_array_equal(c, b.T) + yield assert_array_equal, b, [[1, 0], [1, 1]] + yield assert_array_equal, c, b.T # should return the same dtype as the original array - assert_equal(b.dtype, a.dtype) - assert_equal(c.dtype, a.dtype) + 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([ + [[1, 1], [1, 1]], + [[1, 1], [1, 0]], + [[1, 1], [0, 0]], + ], dtype=dtype) + a_tril_desired = np.array([ + [[1, 0], [1, 1]], + [[1, 0], [1, 0]], + [[1, 0], [0, 0]], + ], dtype=dtype) + a_triu_desired = np.array([ + [[1, 1], [0, 1]], + [[1, 1], [0, 0]], + [[1, 1], [0, 0]], + ], dtype=dtype) + a_triu_observed = np.triu(a) + a_tril_observed = np.tril(a) + yield assert_array_equal, a_triu_observed, a_triu_desired + yield assert_array_equal, a_tril_observed, a_tril_desired + yield assert_equal, a_triu_observed.dtype, a.dtype + yield assert_equal, a_tril_observed.dtype, a.dtype def test_mask_indices(): diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index 8c99f6804..336f23c64 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -430,7 +430,7 @@ def tril(m, k=0): """ m = asanyarray(m) - out = multiply(tri(m.shape[0], m.shape[1], k=k, dtype=m.dtype), m) + out = multiply(tri(m.shape[-2], m.shape[-1], k=k, dtype=m.dtype), m) return out @@ -457,7 +457,7 @@ def triu(m, k=0): """ m = asanyarray(m) - out = multiply((1 - tri(m.shape[0], m.shape[1], k - 1, dtype=m.dtype)), m) + out = multiply((1 - tri(m.shape[-2], m.shape[-1], k - 1, dtype=m.dtype)), m) return out |