diff options
| author | Fabian Pedregosa <fabian.pedregosa@inria.fr> | 2011-08-31 18:16:26 +0200 |
|---|---|---|
| committer | Ralf Gommers <ralf.gommers@googlemail.com> | 2011-08-31 18:37:37 +0200 |
| commit | 1af2f37eae78986b370e2d44df6c5156945d69d8 (patch) | |
| tree | cafc7b5261da2873beae793182e72d96a7d14208 /numpy | |
| parent | f3e70d9c0ea5db74033be9f437bb8e275839c9dc (diff) | |
| download | numpy-1af2f37eae78986b370e2d44df6c5156945d69d8.tar.gz | |
FIX: Make tril/triu return the same dtype as the original array.
This should fix:
* http://projects.scipy.org/numpy/ticket/1848
* http://projects.scipy.org/scipy/ticket/1449
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/lib/tests/test_twodim_base.py | 12 | ||||
| -rw-r--r-- | numpy/lib/twodim_base.py | 4 |
2 files changed, 14 insertions, 2 deletions
diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py index 37165f672..7e7900090 100644 --- a/numpy/lib/tests/test_twodim_base.py +++ b/numpy/lib/tests/test_twodim_base.py @@ -239,6 +239,18 @@ class TestTri(TestCase): 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) diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index 12bba99a6..6fb348275 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -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 |