Using NumPy, create 3D array from function of 1D and 2D without loop

Question

Using python/numpy, I can create the 3D array (note the matrix exponential function) I want like so

import numpy as np
from scipy.linalg import expm

a = np.arange(3)
B = np.ones((2,2))
C = np.zeros((2,2,3))

for i in range(3):
    C[:,:,i] = expm(a[i]*B)

which yields for C, the 3D array

[[[  1.           4.19452805  27.79907502]
  [  0.           3.19452805  26.79907502]]

 [[  0.           3.19452805  26.79907502]
  [  1.           4.19452805  27.79907502]]]

But I'd like to eliminate the loop. Is there any way I can get rid of the for loop? Perhaps by NumPy broadcasting? I had thought of np.kron but can't seem to figure out a good way to reshape so that I can apply the expm function, which requires a square array as an argument.

expm needs a square array as input right, so broadcasting the arguments to it wont help — skmth
– skmth, Commented Mar 15, 2018 at 2:09
What is the actual sizes of A and B? You could always do it by eigndecomposition. — Daniel F
– Daniel F, Commented Mar 15, 2018 at 6:29
@Daniel F, I would say in practice A and B are no greater than dim(500). Eigendecomposition could work, but I can’t see how that would eliminate the loop/multiple calls to expm. — jjgoings
– jjgoings, Commented Mar 15, 2018 at 11:25
You wouldn't use expm in that case, you'd only need to exp the eigenvalues and then recompose by the eigenvectors. And np.linalg.eig can handle arrays of shape (*, N, N) all at once — Daniel F
– Daniel F, Commented Mar 15, 2018 at 15:25

Daniel F · Accepted Answer · 2018-03-16 07:33:38Z

1

Using Sylvester's Formula and a slightly different shape, you can do this vectorized:

C = a[:, None, None] * B

C.shape
 (3, 2, 2)

E, V = np.linalg.eig(C)
V.swapaxes(-1,-2) @ np.exp(E)[..., None] * V
array([[[  1.        ,   0.        ],
        [  0.        ,   1.        ]],

       [[  4.19452805,  -4.19452805],
        [ -3.19452805,  -3.19452805]],

       [[ 27.79907502, -27.79907502],
        [-26.79907502, -26.79907502]]])

This method can take any number of square matrces as inputs in a (*, N, N) array, but while it eliminates the for loop overhead it exchanges it for an np.linalg.eig call, which may be slow if N is large.

edited Mar 16, 2018 at 7:33

answered Mar 16, 2018 at 7:28

Daniel F

14.5k2 gold badges34 silver badges59 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

jjgoings Over a year ago

This looks pretty close to what I want, but the sign in the resulting array does not match? How does this change if B is antisymmetric?

mustafa-guler · Accepted Answer · 2018-03-15 02:02:35Z

0

I'm not sure if this is exactly what you'd like but it gives the same result without a for loop.

a = np.arange(3)
B = np.ones((2,2))
C = np.array([expm(a[i]*B) for i in range(3)]).T

Granted there's still a "for" in the list comprehension, not sure if that fits your purposes.

answered Mar 15, 2018 at 2:02

mustafa-guler

761 silver badge4 bronze badges

Comments

dROOOze · Accepted Answer · 2018-03-15 02:37:38Z

If you want to get the result without calling expm multiple times in a for loop, you can construct the input like this:

import numpy as np
from scipy.linalg import expm

n = 4
t = np.zeros([n*2, n*2])
for i in range(n):
    t[2*i:2*i+2, 2*i:2*i+2] = i

C0 = expm(t)

This results in a format

>>> t
array([[0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 1., 1., 0., 0., 0., 0.],
       [0., 0., 1., 1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 2., 2., 0., 0.],
       [0., 0., 0., 0., 2., 2., 0., 0.],
       [0., 0., 0., 0., 0., 0., 3., 3.],
       [0., 0., 0., 0., 0., 0., 3., 3.]])
>>> C0
array([[  1.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        ,   1.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        ,   0.        ,   4.19452805,   3.19452805,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        ,   0.        ,   3.19452805,   4.19452805,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        ,   0.        ,   0.        ,   0.        ,
         27.79907502,  26.79907502,   0.        ,   0.        ],
       [  0.        ,   0.        ,   0.        ,   0.        ,
         26.79907502,  27.79907502,   0.        ,   0.        ],
       [  0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        , 202.21439675, 201.21439675],
       [  0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        , 201.21439675, 202.21439675]])

If you insist on having the results in your example structure of C, I believe that you'll have to use some kind of loop (list comprehension, range, or otherwise) to achieve this, as I don't think that is a standard output structure for this kind of function.

i1 = range(0, 2*n, 2)
i2 = range(1, 2*n+1, 2)
i = (tuple(i1 + i2 + i1 + i2), tuple(i1 + i1 + i2 + i2))
C = C0[i].reshape(2,2,n)

>>> C
array([[[  1.        ,   4.19452805,  27.79907502, 202.21439675],
        [  0.        ,   3.19452805,  26.79907502, 201.21439675]],

       [[  0.        ,   3.19452805,  26.79907502, 201.21439675],
        [  1.        ,   4.19452805,  27.79907502, 202.21439675]]])

For larger inputs, expm handles sparse square arrays so you shouldn't have any issues with memory.

Collectives™ on Stack Overflow

Using NumPy, create 3D array from function of 1D and 2D without loop

3 Answers 3

1 Comment

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related