# [SciPy-user] assign to diagonal values?

Anand Patil anand.prabhakar.patil@gmail....
Wed Jan 30 23:28:22 CST 2008

On Jan 30, 2008 7:44 PM, David Warde-Farley <dwf@cs.toronto.edu> wrote:
> On 30-Jan-08, at 10:08 PM, Alan G Isaac wrote:
>
> > #diagrv: insert v as diagonal of matrix x (2D only!)
> > def diagrv(x,v,copy=True):
> >       assert(len(x.shape)==2), "For 2-d arrays only."
> >       x = numpy.matrix( x, copy=copy )
> >       stride = 1 + x.shape[1]
> >       x.flat[ slice(0,None,stride) ] = v
> >       return x
>
> Ooh that is clever.  This should really go on the Cookbook page.
>
> Perhaps more generally, for in place modification:
>
> def setdiag(m, d):
>         assert(len(x.shape) == 2)
>         stride = 1 + x.shape[1]
>         m.flat[slice(0,None,stride)] = d

Yeah, that's awesome. I have so many for-loops littering my code for
setting diagonals. Here's an nd-version:

def setdiag(a, d):
assert(all([s == len(d) for s in a.shape]))
stride = 1+sum(cumprod(a.shape[:-1]))
a.flat[::stride] = d

Cheers,
Anand