From e0b8c5cc72d9b36f13650ad53bae0a500aebb652 Mon Sep 17 00:00:00 2001 From: Mark Wiebe Date: Tue, 23 Aug 2011 17:58:35 -0700 Subject: ENH: missingdata: Add maskna= flag to np.eye constructor --- numpy/lib/twodim_base.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) (limited to 'numpy/lib') 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): -- cgit v1.2.1