CODEWITHSUNDEEP
numpy
River Satya
River Satya

Reputation: 1099

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

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!

Upvotes: 6

Views: 3393

Answers (2)

tom10
tom10

Reputation: 69192

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.

Upvotes: 5

Jaime
Jaime

Reputation: 67437

You can use fancy indexing:

In [2]: a = np.zeros((3,3,3))

In [3]: idx = np.arange(3)

In [4]: a[[idx]*3] = 1

In [5]: a
Out[5]:
array([[[ 1.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

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

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  1.]]])

For a more general approach, you could set the diagonal of an arbitrarily sized array doing something like:

def set_diag(arr, values):
    idx = np.arange(np.min(arr.shape))
    arr[[idx]*arr.ndim] = values

Upvotes: 1

Related Questions