NumPy k-th diagonal indices

Question

I'd like to do arithmetics with k-th diagonal of a numpy.array. I need those indices. For example, something like:

>>> a = numpy.eye(2)
>>> a[numpy.diag_indices(a, k=-1)] = 5
>>> a
array([[ 1.,  0.],
       [ 5.,  1.]])

Unfortunately, diag_indices only returns the indices comprising the main diagonal, so at the moment I am doing:

a += numpy.diag([5], -1)

But that doesn't seem as nice or robust. :-)

Is there a way in numpy to get indices for other than the main diagonal?

Answer 1

A bit late, but this version also works for k = 0 (and does not alter the arrays, so does not need to make a copy).

def kth_diag_indices(a, k):
    rows, cols = np.diag_indices_from(a)
    if k < 0:
        return rows[-k:], cols[:k]
    elif k > 0:
        return rows[:-k], cols[k:]
    else:
        return rows, cols