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)
>>>         


        

Answer 1

Here's a way:

1. Create index value arrays.
2. Get the daigonal index values you want.
3. Thats it! :)

Like this:

>>> import numpy as np
>>> rows, cols = np.indices((3,3))
>>> row_vals = np.diag(rows, k=-1)
>>> col_vals = np.diag(cols, k=-1)
>>> z = np.zeros((3,3))
>>> z[row_vals, col_vals]=1
>>> z
array([[ 0.,  0.,  0.],
       [ 1.,  0.,  0.],
       [ 0.,  1.,  0.]])