diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/_methods.py | 2 | ||||
| -rw-r--r-- | numpy/core/tests/test_maskna.py | 17 | ||||
| -rw-r--r-- | numpy/lib/twodim_base.py | 16 |
3 files changed, 24 insertions, 11 deletions
diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py index 3c05cf71d..33de34e1a 100644 --- a/numpy/core/_methods.py +++ b/numpy/core/_methods.py @@ -98,7 +98,7 @@ def _std(a, axis=None, dtype=None, out=None, ddof=0, skipna=skipna, keepdims=keepdims) if isinstance(ret, mu.ndarray): - um.sqrt(ret, out=ret) + ret = um.sqrt(ret, out=ret) else: ret = um.sqrt(ret) diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py index 638c52ba6..77ab22b29 100644 --- a/numpy/core/tests/test_maskna.py +++ b/numpy/core/tests/test_maskna.py @@ -1338,5 +1338,22 @@ def test_array_maskna_linspace_logspace(): assert_(not a.flags.maskna) assert_(b.flags.maskna) +def test_array_maskna_eye(): + # np.eye + + # By default there should be no NA mask + a = np.eye(3) + assert_(not a.flags.maskna) + + a = np.eye(3, maskna=True) + assert_(a.flags.maskna) + assert_(a.flags.ownmaskna) + assert_equal(a, np.eye(3)) + + a = np.eye(3, k=2, maskna=True) + assert_(a.flags.maskna) + assert_(a.flags.ownmaskna) + assert_equal(a, np.eye(3, k=2)) + if __name__ == "__main__": run_module_suite() diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index 9b3b50d04..12bba99a6 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -9,7 +9,7 @@ __all__ = ['diag','diagflat','eye','fliplr','flipud','rot90','tri','triu', from numpy.core.numeric import asanyarray, equal, subtract, arange, \ zeros, greater_equal, multiply, ones, asarray, alltrue, where, \ - empty + empty, diagonal def fliplr(m): """ @@ -166,7 +166,7 @@ def rot90(m, k=1): # k == 3 return fliplr(m.swapaxes(0,1)) -def eye(N, M=None, k=0, dtype=float): +def eye(N, M=None, k=0, dtype=float, maskna=False): """ Return a 2-D array with ones on the diagonal and zeros elsewhere. @@ -182,6 +182,8 @@ def eye(N, M=None, k=0, dtype=float): to a lower diagonal. dtype : data-type, optional Data-type of the returned array. + maskna : boolean + If this is true, the returned array will have an NA mask. Returns ------- @@ -207,14 +209,8 @@ def eye(N, M=None, k=0, dtype=float): """ if M is None: M = N - m = zeros((N, M), dtype=dtype) - if k >= M: - return m - if k >= 0: - i = k - else: - i = (-k) * M - m[:M-k].flat[i::M+1] = 1 + m = zeros((N, M), dtype=dtype, maskna=maskna) + diagonal(m, k)[...] = 1 return m def diag(v, k=0): |