[Numpy-discussion] Functions for indexing into certain parts of an array (2d)

[Keith Goodman]

On Sat, Jun 6, 2009 at 12:01 AM, Fernando Perez <fperez.net at 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