CODEWITHSUNDEEP
pythonnumpyvectorizationnumpy-ndarraydiagonal
Wei
Wei

Reputation: 315

Create diagonal matrix from rows of a 2D matrix without for loop

Suppose that you have a mxn matrix A and want to create m diagonal matrices, each of which is from rows of A and thus has the shape nxn. The resulting matrix should have the shape mxnxn.

I know a typical solution would be:

result = numpy.stack([numpy.diag(A[i,:]) for i in range(A.shape[0])], axis=0)

I am wondering if it is possible we can get the same result without using a loop.

Any idea would be appreciated!

Upvotes: 0

Views: 281

Answers (4)

Shirai Yanagi
Shirai Yanagi

Reputation: 1

(Found my way here, because I happened to run into something requiring dealing with the Jacobian/partial derivatives of the Softmax activation function lately)

Inspired by the above answer by @Ananda :

Here's a way to do it using numpy.einsum in one line

import numpy as np
np.random.seed(0)       # set the seed for a reproducible result
A = np.random.rand(2,3) # MxN for M = 2, N =3
result = np.einsum('ij,jk->ijk', A, np.eye(A.shape[1]))

print(f'{A=}\n{result=}\n')
print('The outputs shape:', result.shape)

and the output will be

A=array([[0.5488135 , 0.71518937, 0.60276338],
   [0.54488318, 0.4236548 , 0.64589411]])
result=array([[[0.5488135 , 0.        , 0.        ],
    [0.        , 0.71518937, 0.        ],
    [0.        , 0.        , 0.60276338]],

   [[0.54488318, 0.        , 0.        ],
    [0.        , 0.4236548 , 0.        ],
    [0.        , 0.        , 0.64589411]]])

The output shape: (2, 3, 3)

Upvotes: 0

kmario23
kmario23

Reputation: 61335

You can try the following where we use broadcasting and numpy.identity() to obtain the desired result.

# `arr` is input of size (m, n); for instance,
arr = np.arange(1, 13).reshape(3, -1)   # (3, 4)

# get column size
In [160]: _ , n = arr.shape

# fill the rows from the input, after promoting it to 3D
In [161]: res = np.eye(n, dtype=arr.dtype) * arr[:,np.newaxis]

In [162]: res.shape
Out[162]: (3, 4, 4)

Note: since we'll have (n,n) as last two dimensions, np.identity() would also work.

In [171]: res = np.identity(n, dtype=arr.dtype) * arr[:,np.newaxis]

In [172]: res.shape
Out[172]: (3, 4, 4)

Another approach would be to use advanced indexing where first initialize the result array with zeros and slice out the indices of the diagonal elements and populate them with the input array.

Upvotes: 0

fountainhead
fountainhead

Reputation: 3722

This seems to work:

result = np.zeros((A.shape[0], A.shape[1], A.shape[1]), dtype=A.dtype)
result[:, range(A.shape[1]), range(A.shape[1])] = A

Test input:

A = np.arange(24).reshape(4,6)

Output of print(result):

[[[ 0  0  0  0  0  0]
  [ 0  1  0  0  0  0]
  [ 0  0  2  0  0  0]
  [ 0  0  0  3  0  0]
  [ 0  0  0  0  4  0]
  [ 0  0  0  0  0  5]]

 [[ 6  0  0  0  0  0]
  [ 0  7  0  0  0  0]
  [ 0  0  8  0  0  0]
  [ 0  0  0  9  0  0]
  [ 0  0  0  0 10  0]
  [ 0  0  0  0  0 11]]

 [[12  0  0  0  0  0]
  [ 0 13  0  0  0  0]
  [ 0  0 14  0  0  0]
  [ 0  0  0 15  0  0]
  [ 0  0  0  0 16  0]
  [ 0  0  0  0  0 17]]

 [[18  0  0  0  0  0]
  [ 0 19  0  0  0  0]
  [ 0  0 20  0  0  0]
  [ 0  0  0 21  0  0]
  [ 0  0  0  0 22  0]
  [ 0  0  0  0  0 23]]]

Upvotes: 1

Ananda
Ananda

Reputation: 3272

Not going to be the most efficient way of doing it but something like this might work

import numpy as np

A = np.random.rand(10, 5)

S = np.einsum("ai,ij->aij", A, np.ones((5, 5)))
M = np.eye(5).astype(np.bool)
M = np.repeat(M[None, ...], 10, axis=0)
S[~M]=0
print(S.shape)

Upvotes: 0

Related Questions