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
####
def mask_indices(n,mask_func,k=0):
"""Return the indices for an array, given a masking function like
tri{u,l}."""
m = np.ones((n,n),int)
a = mask_func(m,k)
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]])
"""
return mask_indices(n,np.tril,k)
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]])
"""
return mask_indices(n,np.triu,k)
_______________________________________________
Numpy-discussion mailing list
Numpy-discussion@scipy.org
http://mail.scipy.org/mailman/listinfo/numpy-discussion
