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?

Answer 1

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]]])