Create Diagonal Matrix from column of a multi shape array

user_12216747 0
Viewed 425

Suppose I have an array like

mat = np.array([[1,1],[2,2],[3,3]])

which has shape (3,2). I want to create A new array with the shape (3,2,2) by taking the last axis of the array and create diagonal matrices out of it. I can do it with a for loop like

mat2 = np.zeros((3,2,2))
for i in np.arange(0,3):
    mat2[i] = np.diag(mat[i])

which yields the desired output 

[[[1. 0.]
  [0. 1.]]

 [[2. 0.]
  [0. 2.]]

 [[3. 0.]
  [0. 3.]]]

But is there a way to do it in a direct vectorized (faster?!) version? In my real problem I have a large highly dimensional array of shape (...,n) and need to convert the last axis into a diagonal matrix with shape (...,n,n) in the end.

numpy arrays python
4 Answers

Here a possible implementation: use np.diagonal to take a view of the relevant diagonals and force the view to be writeable with setflags, and write to the view:

expanded = np.zeros(mat.shape + mat.shape[-1:], dtype=mat.dtype)
diagonals = np.diagonal(expanded, axis1=-2, axis2=-1)
diagonals.setflags(write=True)

diagonals[:] = mat

expanded
array([[[1, 0],
        [0, 1]],

       [[2, 0],
        [0, 2]],

       [[3, 0],
        [0, 3]]])
user_124115360

There is a great library called numba that allows to just in time compile python functions to get a speed comparable to C/Fortran when dealing with numpy arrays. If you are at a point where you cannot express your problem vectorized this can save your day.

from numba import jit

@jit(nopython=True)
def get_diag_mat(mat):
    mat2 = np.zeros((mat.shape[0], mat.shape[1], mat.shape[1]))
    for i in range(mat2.shape[0]):
        mat2[i] = np.diag(mat[i])
    return mat2

Even if you can vectorize your problem with numpy, you can sometimes save some memory, when you switch to "jitted" loops.

PS: The first execution of the function has an overhead because of compilation.

PPS: Note that you rarely want to use np.arange. Especially not for looping. python3's range is lazy and does not construct an array of integers in your memory.

user_53126060

You can use the np.eye function and multiply. 

>>> mat = np.array([[1,1],[2,2],[3,3]])
>>> np.multiply(np.eye(2),mat[:,np.newaxis])
array([[[ 1.,  0.],
        [ 0.,  1.]],

       [[ 2.,  0.],
        [ 0.,  2.]],

       [[ 3.,  0.],
        [ 0.,  3.]]])
user_106423640 
In [219]: m = mat.shape[-1]                                                                    
In [220]: mat2 = np.zeros(mat.shape+(m,),mat.dtype)                                            
In [221]: idx = np.arange(m)                                                                   
In [222]: mat2[...,idx,idx] = mat[...,idx]                                                     
In [223]: mat2                                                                                 
Out[223]: 
array([[[1, 0],
        [0, 1]],

       [[2, 0],
        [0, 2]],

       [[3, 0],
        [0, 3]]])

testing on a larger mat:

In [224]: mat = np.array([[1,2,3],[2,3,4],[3,4,5]])                                            
In [225]: m = mat.shape[-1]                                                                    
In [226]: mat2 = np.zeros(mat.shape+(m,),mat.dtype)                                            
In [227]: idx = np.arange(m)                                                                   
In [228]: mat2[...,idx,idx] = mat[...,idx]                                                     
In [230]: mat2                                                                                 
Out[230]: 
array([[[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3]],

       [[2, 0, 0],
        [0, 3, 0],
        [0, 0, 4]],

       [[3, 0, 0],
        [0, 4, 0],
        [0, 0, 5]]])
user_9019250
Related