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

Answer 2

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

Answer 3

The indices of the k'th diagonal of a can be computed with

def kth_diag_indices(a, k):
    rowidx, colidx = np.diag_indices_from(a)
    colidx = colidx.copy()  # rowidx and colidx share the same buffer

    if k > 0:
        colidx += k
    else:
        rowidx -= k
    k = np.abs(k)

    return rowidx[:-k], colidx[:-k]

Demo:

>>> a
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])
>>> a[kth_diag_indices(a, 1)]
array([ 1,  7, 13, 19])
>>> a[kth_diag_indices(a, 2)]
array([ 2,  8, 14])
>>> a[kth_diag_indices(a, -1)]
array([ 5, 11, 17, 23])

Answer 4

So since np.diag_indices() doesn't have the same functionality as np.triu_indices() does to get the kth diagonals/triangles, another approach is to just use np.eye(n,k) to construct an nxn matrix with 1's on the kth diagonal, and then use np.where to extract a tuple of indices of where the 1's are located.

So we can do this with just:

rows, cols = np.where(np.eye(*a.shape, k=-1))

This is an extra allocation of the eye matrix that may be excessive in some cases, but its an easy one-liner.

Answer 5

There is another solution. Create a matrix E with np.eye. You can just change the main diagonal like the following. Then, create two more matrizes with the k parameter in eye. In the end just add all matrizes together.

E = np.eye(5)
E = E*2
F = -1*np.eye(len(E),k=1)
G = -1*np.eye(len(E),k=-1)

E = E+F+G
print(E)

Answer 6

Use numpy.diag(v, k=0)

Where k sets the diagonal location from center.

ie. {k=0: "default center", k=(-1): "1 row to the left of center", k=1: "1 row to the right of center}

Then perform the arithmetic as you would normally expect.

Check out the docs here: np.diag().

Examples:

In [3]: np.diag(np.arange(6), k=0)
Out[3]: 
array([[0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0],
       [0, 0, 0, 3, 0, 0],
       [0, 0, 0, 0, 4, 0],
       [0, 0, 0, 0, 0, 5]])

In [4]: np.diag(np.arange(6), k=1)
Out[4]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 1, 0, 0, 0, 0],
       [0, 0, 0, 2, 0, 0, 0],
       [0, 0, 0, 0, 3, 0, 0],
       [0, 0, 0, 0, 0, 4, 0],
       [0, 0, 0, 0, 0, 0, 5],
       [0, 0, 0, 0, 0, 0, 0]])

In [5]: np.diag(np.arange(6), k=-1)
Out[5]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0, 0],
       [0, 0, 0, 3, 0, 0, 0],
       [0, 0, 0, 0, 4, 0, 0],
       [0, 0, 0, 0, 0, 5, 0]])