Eu gostaria de fazer aritmética com a k-ésima diagonal de um numpy.array. Eu preciso desses índices. Por exemplo, algo como:

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

Infelizmente, diag_indices só retorna os índices que compõem a diagonal principal, então no momento eu estou fazendo:

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

Mas isso não parece tão bom ou robusto. :-)

Existe uma maneira em numpy para obter índices para outros que não a diagonal principal?