[FIXED] Making a Diagonal matrix with desired shape

Issue

How to make a random elements digonal matrix in a numpy array with the desired shape, for example I want to make F to be a numpy array with shape (3, 2, 4, 4), and the diagonal matrix dimensions 4*4 like this:

  F= np.array([[[[  a,  0,  0,  0],
                  [ 0,  b,  0,  0],
                  [ 0,  0,  c,  0],
                  [ 0,  0,  0,  d]],

                 [[e, 0, 0, 0],
                  [0, f, 0, 0],
                  [0, 0, g, 0],
                  [0, 0, 0, h]]],

                [[[i, 0, 0, 0],
                  [0, j, 0, 0],
                  [0, 0, k, 0],
                  [0, 0, 0, l]],

                 [[m, 0, 0, 0],
                  [0, n, 0, 0],
                  [0, 0, o, 0],
                  [0, 0, 0, p]]],

                [[[q, 0, 0, 0],
                  [0, r, 0, 0],
                  [0, 0, s, 0],
                  [0, 0, 0, t]],

                 [[ u,  0,  0, 0],
                  [ 0,  v,  0, 0],
                  [ 0,  0,  w, 0],
                  [ 0,  0,  0, x]]]]

is there any simple and handy command to do this instead of performing for loops, for example something like this simple command: F=(np.arange(96)).reshape(3, 2, 4, 4)

Solution

One way to do it is to use np.fromfunction. We can write a custom function which will be given input arrays x, y, z and w, all of which have the same shape as the target array, with the property that every element of x has value equal to its index in dimension 0, every element of y has value equal to its index in dimension 1, and so on.

Then the diagonal elements that you're interested in are just those for which z==w. We can use np.where to construct an array in which those elements are random, and the others zero.

def random_if_diagonal(x, y, z, w):
    return np.where(z==w,
                    np.random.default_rng(seed=42).uniform(size=w.shape),
                    np.zeros_like(w))

shape = (3, 2, 4, 4)
arr = np.fromfunction(random_if_diagonal, shape=shape)
print(arr)

The random values here are floats between 0 and 1. If you prefer, you can scale them by multiplying, and convert them to int.

Result:

[[[[0.77395605 0.         0.         0.        ]
   [0.         0.97562235 0.         0.        ]
   [0.         0.         0.37079802 0.        ]
   [0.         0.         0.         0.22723872]]

  [[0.55458479 0.         0.         0.        ]
   [0.         0.35452597 0.         0.        ]
   [0.         0.         0.466721   0.        ]
   [0.         0.         0.         0.96750973]]]


 [[[0.32582536 0.         0.         0.        ]
   [0.         0.47570493 0.         0.        ]
   [0.         0.         0.7002651  0.        ]
   [0.         0.         0.         0.2883281 ]]

  [[0.6824955  0.         0.         0.        ]
   [0.         0.66485086 0.         0.        ]
   [0.         0.         0.139797   0.        ]
   [0.         0.         0.         0.76499886]]]


 [[[0.63471832 0.         0.         0.        ]
   [0.         0.43671739 0.         0.        ]
   [0.         0.         0.05830274 0.        ]
   [0.         0.         0.         0.78389821]]

  [[0.66431354 0.         0.         0.        ]
   [0.         0.09004786 0.         0.        ]
   [0.         0.         0.1523121  0.        ]
   [0.         0.         0.         0.63028259]]]]

It's worth saying this isn't the most memory-optimal way of doing it, because it constructs an entire random matrix and an entire zero matrix of the same shape as the target. If the size of array in your example is a realistic use case though, that shouldn't be a big deal.

Answered By - slothrop