Me gustaría hacer aritmética con la diagonal K-th de un numpy.array. Necesito esos índices. Por ejemplo, algo como:

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

Desafortunadamente, diag_indices solo devuelve los índices que comprenden la diagonal principal, por lo que en este momento estoy haciendo:

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

Pero eso no parece tan bonito ni robusto. :-)

¿Hay alguna manera, en números, de obtener índices para otra cosa que no sea la diagonal principal?