CODEWITHSUNDEEP
python-3.xnumpyvectorization
user447457
user447457

Reputation: 71

Efficiently creating a tensor of diagonal matrices

Assume a map from scalars to matrices of a fixed dimension. How would one efficiently create a vectorized version of this map?

More specifically assume there is a constant vector lamb with n entries. Given a scalar t I'm interested in the diagonal matrix given by

np.diag(np.exp(lamb*t))

using numpy. This will be an n times n matrix. Now given a matrix T of size m_1 times m_2 I would like to calculate the tensor D of shape (m_1,m_2,n,n) given for 0 <= i < m_1, 0 <= j < m_2 by

D[i,j,:,:] = np.diag(np.exp(lamb*T[i,j]))

How would one efficiently get this tensor?

Upvotes: 2

Views: 435

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

One comparatively straight-forward way would be using einsum.

Example:

>>> T = np.array([[1,2,3], [4,6,7]])
>>> lam = np.array([1,2,5])
>>> D = np.zeros((*T.shape, n, n))
>>> np.einsum('ijkk->ijk', D)[...] = np.exp(np.multiply.outer(T, lam))
>>> D
array([[[[2.71828183e+00, 0.00000000e+00, 0.00000000e+00],
         [0.00000000e+00, 7.38905610e+00, 0.00000000e+00],
         [0.00000000e+00, 0.00000000e+00, 1.48413159e+02]],

        [[7.38905610e+00, 0.00000000e+00, 0.00000000e+00],
         [0.00000000e+00, 5.45981500e+01, 0.00000000e+00],
         [0.00000000e+00, 0.00000000e+00, 2.20264658e+04]],

        [[2.00855369e+01, 0.00000000e+00, 0.00000000e+00],
         [0.00000000e+00, 4.03428793e+02, 0.00000000e+00],
         [0.00000000e+00, 0.00000000e+00, 3.26901737e+06]]],


       [[[5.45981500e+01, 0.00000000e+00, 0.00000000e+00],
         [0.00000000e+00, 2.98095799e+03, 0.00000000e+00],
         [0.00000000e+00, 0.00000000e+00, 4.85165195e+08]],

        [[4.03428793e+02, 0.00000000e+00, 0.00000000e+00],
         [0.00000000e+00, 1.62754791e+05, 0.00000000e+00],
         [0.00000000e+00, 0.00000000e+00, 1.06864746e+13]],

        [[1.09663316e+03, 0.00000000e+00, 0.00000000e+00],
         [0.00000000e+00, 1.20260428e+06, 0.00000000e+00],
         [0.00000000e+00, 0.00000000e+00, 1.58601345e+15]]]])

You can speed this up a little bit using the out keyword to avoid one copy:

np.exp(np.multiply.outer(T, lam), out=np.einsum('ijkk->ijk', D))

Upvotes: 2

Related Questions