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

[Fernando Perez]

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)