How to create a diagonal multi-dimensional (ie greater than 2) in numpy

Question

Is there a higher (than two) dimensional equivalent of diag?

L = [...] # some arbitrary list.
A = ndarray.diag(L)

will create a diagonal 2-d matrix shape=(len(L), len(L)) with elements of L on the diagonal.

I'd like to do the equivalent of:

length = len(L)
A = np.zeros((length, length, length))
for i in range(length):
    A[i][i][i] = L[i]

Is there a slick way to do this?

Thanks!

Answer 1

You can use diag_indices to get the indices to be set. For example,

x = np.zeros((3,3,3))
L = np.arange(6,9)

x[np.diag_indices(3,ndim=3)] = L

gives

array([[[ 6.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  7.,  0.],
        [ 0.,  0.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  8.]]])

Under the hood diag_indices is just the code Jaime posted, so which to use depends on whether you want it spelled out in a numpy function, or DIY.