CODEWITHSUNDEEP
pythonnumpymatrixscipydiagonal
kinder chen
kinder chen

Reputation: 1461

How to assign elements into the diagonal of a 3d matrix efficiently?

a=np.zeros((3,3,3))
b=np.arange(3)
c=np.arange(9).reshape(3,3)

I wanna put the elements of the array b or c along the diagonal (or above/below the diagonal) of the 3d matrix (tensor) a with respect to a specific axis.

I tired numpy.diagflat, but it only works for 2d matrix.

For instance, how to make the following matrix?

array([[[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

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

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

Upvotes: 1

Views: 2096

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53029

For the main diagonals you can use np.einsum. For example:

>> np.einsum('iii->i', a)[...] = b
>>> a
array([[[ 0.,  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.,  2.]]])

Or:

>>> np.einsum('iji->ji', a)[...] = c
>>> a
array([[[ 0.,  0.,  0.],
        [ 3.,  0.,  0.],
        [ 6.,  0.,  0.]],

       [[ 0.,  1.,  0.],
        [ 0.,  4.,  0.],
        [ 0.,  7.,  0.]],

       [[ 0.,  0.,  2.],
        [ 0.,  0.,  5.],
        [ 0.,  0.,  8.]]])

Edit: Broadcasting works normally:

>>> np.einsum('ijj->ij', a)[...] = b
>>> a
array([[[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

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

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

Subdiagonals also work but are more tricky as some manual slicing is required. For example:

>>> a=np.zeros((3,3,3))
>>> np.einsum('iij->ij', a[:2,1:])[...] = c[1:]
>>> a
array([[[ 0.,  0.,  0.],
        [ 3.,  4.,  5.],
        [ 0.,  0.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 6.,  7.,  8.]],

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

Upvotes: 6

ChootsMagoots
ChootsMagoots

Reputation: 670

If a sparse matrix will serve your needs, here are some resources:

http://www.janeriksolem.net/sparray-sparse-n-dimensional-arrays-in.html

sparse 3d matrix/array in Python?

Upvotes: -1

Related Questions