Create stack of arrays from diagonal values using numpy

Question

I'm trying to do some matrix calculations in python and came across a problem when I tried to speed up my code using stacked arrays instead of simple for loops. I need to create a 2D-array with values (given as 1D-array) on the diagonal, but could't figure out a smart way to do it with stacked arrays.

In the old (loop) version, I used the np.diag() method, which returns exactly what I need (a 2D-array in that case) if I give the values as 1D-array as input. However, when I switched to stacked arrays my input is not a 1D-array anymore, so that the np.diag() method returns a copy of the diagonal of my 2D-input instead.

Old version with 1D input:

import numpy as np
vals = np.array([1,2,3])
mat = np.diag(vals)
print(mat.shape)

Out: (3, 3)

New version with 2D input:

vals_stack = np.repeat(np.expand_dims(vals, axis=0), 5, axis=0) 
# btw: is there a better way to repeat/stack my array? 
mat_stack = np.diag(vals_stack)
print(mat_stack.shape)

Out: (3,)

So you can see that np.diag() returns a 1D-array (as expected from the documentation), but I actually need a stack of 2D-arrays. So the shape of the mat_stack must be (7,3,3) and not (3,). Is there any function for that in numpy? Or do I have to loop over that additional dimension like this:

def mydiag(stack):
    diag = np.zeros([stack.shape[0], stack.shape[1], stack.shape[1]])
    for i in np.arange(stack.shape[0]):
        diag[i,:,:] = np.diag([stack[i,:].ravel()])
    return diag

Simon Delecourt · Accepted Answer · 2019-07-09 15:34:37Z

2

In numpy you should use apply_along_axis. There is even an example at the end of the doc for your specific case (here). So the answer is :

np.apply_along_axis(np.diag, -1, vals_stack)

A more pythonic way would be something like this:

[np.diag(row) for row in vals_stack]

answered Jul 9, 2019 at 15:34

Simon Delecourt

1,60910 silver badges14 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

koxx Over a year ago

That is exactly what I was looking for and it works like a charm! Thank you!

hpaulj · Accepted Answer · 2019-07-09 17:34:06Z

Is this what you had in mind:

In [499]: x = np.arange(12).reshape(4,3)                                                                     
In [500]: X = np.zeros((4,3,3),int)                                                                          
In [501]: X[np.arange(4)[:,None],np.arange(3), np.arange(3)] = x                                             
In [502]: X                                                                                                  
Out[502]: 
array([[[ 0,  0,  0],
        [ 0,  1,  0],
        [ 0,  0,  2]],

       [[ 3,  0,  0],
        [ 0,  4,  0],
        [ 0,  0,  5]],

       [[ 6,  0,  0],
        [ 0,  7,  0],
        [ 0,  0,  8]],

       [[ 9,  0,  0],
        [ 0, 10,  0],
        [ 0,  0, 11]]])

X[0,np.arange(3), np.arange(3)] indexes the diagonal on the first plane. np.arange(4)[:,None] is a (4,1) array, which broadcasts with a (3,) to index a (4,3) block, matching the size of x.

Collectives™ on Stack Overflow

Create stack of arrays from diagonal values using numpy

2 Answers 2

1 Comment

Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Related