Merge pull request #4364 from argriffing/triu-broadcasting

charris · charris · commit 56eb28ed2957 · 2014-02-24T22:52:16.000-07:00
ENH: tril and triu broadcasting
diff --git a/doc/release/1.9.0-notes.rst b/doc/release/1.9.0-notes.rst
@@ -81,6 +81,10 @@ Dtype parameter added to `np.linspace` and `np.logspace`
 The returned data type from the `linspace` and `logspace` functions
 can now be specificed using the dtype parameter.
 
+More general `np.triu` and `np.tril` broadcasting
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+For arrays with `ndim` exceeding 2, these functions will now apply to the
+final two axes instead of raising an exception.
 
 `tobytes` alias for `tostring` method
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py
@@ -275,16 +275,40 @@ def test_dtype(self):
         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
@@ -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
 
 
