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

Fernando Perez fperez.net at gmail.com
Sat Jun 6 03:01:45 EDT 2009

```Howdy,

I'm finding myself often having to index *into* arrays to set values.
As best I can see, in numpy functions/methods like diag or tri{u,l}
provide for the extraction of values from arrays, but I haven't found
their counterparts for generating the equivalent indices.

Their implementations are actually quite trivial once you think of it,
but I wonder if there's any benefit in having this type of machinery
in numpy itself. I had to think of how to do it more than once, so I
ended up writing up a few utilities for it.

Below are my naive versions I have for internal use.  If there's any
interest, I'm happy to document them to compliance as a patch.

Cheers,

f

####

"""Return the indices for an array, given a masking function like
tri{u,l}."""
m = np.ones((n,n),int)
return np.where(a != 0)

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

def tril_indices(n,k=0):
"""Return the indices for the lower-triangle of an (n,n) array.

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]])
>>> dl = tril_indices(4)
>>> a[dl] = -1
>>> a
array([[-1,  2,  3,  4],
[-1, -1,  7,  8],
[-1, -1, -1, 12],
[-1, -1, -1, -1]])
>>> dl = tril_indices(4,2)
>>> a[dl] = -10
>>> a
array([[-10, -10, -10,   4],
[-10, -10, -10, -10],
[-10, -10, -10, -10],
[-10, -10, -10, -10]])
"""

def triu_indices(n,k=0):
"""Return the indices for the upper-triangle of an (n,n) array.

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]])
>>> du = triu_indices(4)
>>> a[du] = -1
>>> a
array([[-1, -1, -1, -1],
[ 5, -1, -1, -1],
[ 9, 10, -1, -1],
[13, 14, 15, -1]])
>>> du = triu_indices(4,2)
>>> a[du] = -10
>>> a
array([[ -1,  -1, -10, -10],
[  5,  -1,  -1, -10],
[  9,  10,  -1,  -1],
[ 13,  14,  15,  -1]])
"""