def eye(N, M=None, k=0, dtype=float, maskna=False): """ Return a 2-D array with ones on the diagonal and zeros elsewhere. Parameters ---------- N : int Number of rows in the output. M : int, optional Number of columns in the output. If None, defaults to `N`. k : int, optional Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper diagonal, and a negative value 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 ------- I : ndarray of shape (N,M) An array where all elements are equal to zero, except for the `k`-th diagonal, whose values are equal to one. See Also -------- identity : (almost) equivalent function diag : diagonal 2-D array from a 1-D array specified by the user. Examples -------- >>> np.eye(2, dtype=int) array([[1, 0], [0, 1]]) >>> np.eye(3, k=1) array([[ 0., 1., 0.], [ 0., 0., 1.], [ 0., 0., 0.]]) """ if M is None: M = N m = zeros((N, M), dtype=dtype, maskna=maskna) diagonal(m, k)[...] = 1 return m
def diag(v, k=0): """ Extract a diagonal or construct a diagonal array. See the more detailed documentation for ``numpy.diagonal`` if you use this function to extract a diagonal and wish to write to the resulting array; whether it returns a copy or a view depends on what version of numpy you are using. Parameters ---------- v : array_like If `v` is a 2-D array, return a copy of its `k`-th diagonal. If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th diagonal. k : int, optional Diagonal in question. The default is 0. Use `k>0` for diagonals above the main diagonal, and `k<0` for diagonals below the main diagonal. Returns ------- out : ndarray The extracted diagonal or constructed diagonal array. See Also -------- diagonal : Return specified diagonals. diagflat : Create a 2-D array with the flattened input as a diagonal. trace : Sum along diagonals. triu : Upper triangle of an array. tril : Lower triangle of an array. Examples -------- >>> x = np.arange(9).reshape((3,3)) >>> x array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) >>> np.diag(x) array([0, 4, 8]) >>> np.diag(x, k=1) array([1, 5]) >>> np.diag(x, k=-1) array([3, 7]) >>> np.diag(np.diag(x)) array([[0, 0, 0], [0, 4, 0], [0, 0, 8]]) """ v = asanyarray(v) s = v.shape if len(s) == 1: n = s[0] + abs(k) res = zeros((n, n), v.dtype) if k >= 0: i = k else: i = (-k) * n res[:n - k].flat[i::n + 1] = v return res elif len(s) == 2: return diagonal(v, k) else: raise ValueError("Input must be 1- or 2-d.")
def diag(v, k=0): """ Extract a diagonal or construct a diagonal array. See the more detailed documentation for ``numpy.diagonal`` if you use this function to extract a diagonal and wish to write to the resulting array; whether it returns a copy or a view depends on what version of numpy you are using. Parameters ---------- v : array_like If `v` is a 2-D array, return a copy of its `k`-th diagonal. If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th diagonal. k : int, optional Diagonal in question. The default is 0. Use `k>0` for diagonals above the main diagonal, and `k<0` for diagonals below the main diagonal. Returns ------- out : ndarray The extracted diagonal or constructed diagonal array. See Also -------- diagonal : Return specified diagonals. diagflat : Create a 2-D array with the flattened input as a diagonal. trace : Sum along diagonals. triu : Upper triangle of an array. tril : Lower triangle of an array. Examples -------- >>> x = np.arange(9).reshape((3,3)) >>> x array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) >>> np.diag(x) array([0, 4, 8]) >>> np.diag(x, k=1) array([1, 5]) >>> np.diag(x, k=-1) array([3, 7]) >>> np.diag(np.diag(x)) array([[0, 0, 0], [0, 4, 0], [0, 0, 8]]) """ v = asanyarray(v) s = v.shape if len(s) == 1: n = s[0]+abs(k) res = zeros((n, n), v.dtype) if k >= 0: i = k else: i = (-k) * n res[:n-k].flat[i::n+1] = v return res elif len(s) == 2: return diagonal(v, k) else: raise ValueError("Input must be 1- or 2-d.")