Transform matrix columns into diagonal matrices and wrap them up without for loops

Question

I have M vectors with N entries each (MxN matrix, if you will):

A = [A1, A2, ... , AN]
B = [B1, B2, ... , BN]
...
Z = [Z1, Z2, ... , ZN]

I want to transform these so that the output is:

[[[A1,  0, ... ,  0],
  [0,  B1, ... ,  0],
   ...
  [0,   0, ... , Z1]],
  
 [[A2,  0, ... ,  0],
  [0,  B2, ... ,  0],
   ...
  [0,   0, ... , Z2]],
  
  ...
  
 [[AN,  0, ... ,  0],
  [0,  BN, ... ,  0],
   ...
  [0,   0, ... , ZN]]]

The goal is not to use for loops at all and achieve this solely with numpy operations. Ideas?

Hey, Ivan. Yeah, I have tried `np.diag`, but it only works for 1D or 2D arrays and not in the way I want. Now I am trying to implement it with matrix multiplication. — giochanturia, Dec 05 '20 at 13:02
This is quite vague question. There multiple ways to do this. Did you try np.fill_diagonal https://numpy.org/doc/stable/reference/generated/numpy.fill_diagonal.html — Valentyn, Dec 05 '20 at 17:10

Ivan · Accepted Answer · 2021-08-27T21:21:11.070

I came upon an example here which uses np.apply_along_axis to fill multiple diagonals.

A = [11, 12, 13, 14]
B = [21, 22, 23, 24]
C = [31, 32, 33, 34]
D = [41, 42, 43, 44]
E = [51, 52, 53, 54]

Z = np.array([A, B, C, D, E])

Having constructed Z, you take its transposed and fill an empty diagonal array with its values:

>>> np.apply_along_axis(np.diag, -1, Z.T)
array([[[11,  0,  0,  0,  0],
        [ 0, 21,  0,  0,  0],
        [ 0,  0, 31,  0,  0],
        [ 0,  0,  0, 41,  0],
        [ 0,  0,  0,  0, 51]],

       [[12,  0,  0,  0,  0],
        [ 0, 22,  0,  0,  0],
        [ 0,  0, 32,  0,  0],
        [ 0,  0,  0, 42,  0],
        [ 0,  0,  0,  0, 52]],

       [[13,  0,  0,  0,  0],
        [ 0, 23,  0,  0,  0],
        [ 0,  0, 33,  0,  0],
        [ 0,  0,  0, 43,  0],
        [ 0,  0,  0,  0, 53]],

       [[14,  0,  0,  0,  0],
        [ 0, 24,  0,  0,  0],
        [ 0,  0, 34,  0,  0],
        [ 0,  0,  0, 44,  0],
        [ 0,  0,  0,  0, 54]]])

hpaulj · Answer 2 · 2020-12-05T17:04:53.403

While your question is clear enough, we prefer to see a [MCVE] (that tag expands in a comment, see SO intro). And when asking for no-loop solutions, it's polite to show a working loop solution - that gives us something test against. I at least like to prove that my answer works.

So here my minimal working example - with a loop solution:

In [308]: M = np.stack([np.arange(i,j) for i,j in zip([1,11,21],[5,15,25])])
In [309]: M.shape
Out[309]: (3, 4)
In [310]: M
Out[310]: 
array([[ 1,  2,  3,  4],
       [11, 12, 13, 14],
       [21, 22, 23, 24]])
In [311]: np.stack([np.diag(M[:,i]) for i in range(4)])
Out[311]: 
array([[[ 1,  0,  0],
        [ 0, 11,  0],
        [ 0,  0, 21]],

       [[ 2,  0,  0],
        [ 0, 12,  0],
        [ 0,  0, 22]],

       [[ 3,  0,  0],
        [ 0, 13,  0],
        [ 0,  0, 23]],

       [[ 4,  0,  0],
        [ 0, 14,  0],
        [ 0,  0, 24]]])

np.diagonal lets us specify axes. So for example we can use it to extract the starting array from Out[311]:

In [318]: np.diagonal(Out[311],axis1=1, axis2=2)
Out[318]: 
array([[ 1, 11, 21],
       [ 2, 12, 22],
       [ 3, 13, 23],
       [ 4, 14, 24]])

I don't (off hand) see a multidimensional building. We could dig in docs some more, or look at the existing code to construct an equivalent. Or just accept the time penalty of that loop :)

The diag(onal) functions use a flat indexing for speed, but it's easier to use multidimensional indexing. We can access those same values with

In [319]: Out[311][:,np.arange(3),np.arange(3)]
Out[319]: 
array([[ 1, 11, 21],
       [ 2, 12, 22],
       [ 3, 13, 23],
       [ 4, 14, 24]])

We can use the same indexing to assign values.

In [320]: res = np.zeros((4,3,3),int)
In [321]: res[:,np.arange(3), np.arange(3)] = M.T
In [322]: res
Out[322]: 
array([[[ 1,  0,  0],
        [ 0, 11,  0],
        [ 0,  0, 21]],

       [[ 2,  0,  0],
        [ 0, 12,  0],
        [ 0,  0, 22]],

       [[ 3,  0,  0],
        [ 0, 13,  0],
        [ 0,  0, 23]],

       [[ 4,  0,  0],
        [ 0, 14,  0],
        [ 0,  0, 24]]])

If these last steps are confusing, I'd suggest experimenting with creating your own 2d diagonal array. Start small, and build on that knowledge.

score 0 · Answer 3 · answered Dec 05 '20 at 13:07

You can use the diag numpy function, e.g.

import numpy as np

A = [0, 1, 2]
B = [0, 1, 2]
C = [0, 1, 2]

# Merge lists into matrix
X = np.array([A, B, C])

# Index columns and merge into list of diagonal matrices
print([np.diag(X[:, 0]), np.diag(X[:, 1]), np.diag(X[:, 2])])

To automate the final step (above) you may need to use a for-loop however e.g.

diag_list = []
for n in N:
   diag_list.append(np.diag(X[:, n])

Transform matrix columns into diagonal matrices and wrap them up without for loops

3 Answers3