On Sat, Jun 6, 2009 at 12:01 AM, Fernando Perez <fperez....@gmail.com> wrote:
> def diag_indices(n,ndim=2): > """Return the indices to index into a diagonal. > > Examples > -------- > >>> a = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]]) > >>> a > array([[ 1, 2, 3, 4], > [ 5, 6, 7, 8], > [ 9, 10, 11, 12], > [13, 14, 15, 16]]) > >>> di = diag_indices(4) > >>> a[di] = 100 > >>> a > array([[100, 2, 3, 4], > [ 5, 100, 7, 8], > [ 9, 10, 100, 12], > [ 13, 14, 15, 100]]) > """ > idx = np.arange(n) > return (idx,)*ndim I often set the diagonal to zero. Now I can make a fill_diag function. What do you think of passing in the array a instead of n and ndim (diag_indices_list_2 below)? from numpy import arange def diag_indices(n, ndim=2): idx = arange(n) return (idx,)*ndim def diag_indices_list(n, ndim=2): idx = range(n) return (idx,)*ndim def diag_indices_list_2(a): idx = range(a.shape[0]) return (idx,) * a.ndim >> a = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]]) >> n = 4 >> ndim = 2 >> >> timeit diag_indices(n, ndim) 1000000 loops, best of 3: 1.76 µs per loop >> >> timeit diag_indices_list(n, ndim) 1000000 loops, best of 3: 1.03 µs per loop >> >> timeit diag_indices_list_2(a) 1000000 loops, best of 3: 1.21 µs per loop _______________________________________________ Numpy-discussion mailing list Numpy-discussion@scipy.org http://mail.scipy.org/mailman/listinfo/numpy-discussion
