Permalink
Browse files

FIX: Make tril/triu return the same dtype as the original array.

  • Loading branch information...
1 parent f3e70d9 commit 1af2f37eae78986b370e2d44df6c5156945d69d8 @fabianp fabianp committed with rgommers Aug 31, 2011
Showing with 14 additions and 2 deletions.
  1. +12 −0 numpy/lib/tests/test_twodim_base.py
  2. +2 −2 numpy/lib/twodim_base.py
@@ -239,6 +239,18 @@ def test_dtype(self):
assert_array_equal(tri(3,dtype=bool),out.astype(bool))
+def test_tril_triu():
+ 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)
+ # should return the same dtype as the original array
+ assert_equal(b.dtype, a.dtype)
+ assert_equal(c.dtype, a.dtype)
+
+
def test_mask_indices():
# simple test without offset
iu = mask_indices(3, np.triu)
View
@@ -413,7 +413,7 @@ def tril(m, k=0):
"""
m = asanyarray(m)
- out = multiply(tri(m.shape[0], m.shape[1], k=k, dtype=int),m)
+ out = multiply(tri(m.shape[0], m.shape[1], k=k, dtype=m.dtype),m)
return out
def triu(m, k=0):
@@ -439,7 +439,7 @@ def triu(m, k=0):
"""
m = asanyarray(m)
- out = multiply((1 - tri(m.shape[0], m.shape[1], k - 1, int)), m)
+ out = multiply((1 - tri(m.shape[0], m.shape[1], k - 1, dtype=m.dtype)), m)
return out
# borrowed from John Hunter and matplotlib

0 comments on commit 1af2f37

Please sign in to comment.