Эффективное создание тензора диагональных матриц

Question

Предположим карту от скаляров до матриц фиксированного измерения. Как эффективно создать векторизованную версию этой карты?

Более конкретно, предположим, что существует постоянный вектор ягненка с n записями. Учитывая скаляр t, меня интересует диагональная матрица, заданная

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

используя NumPy. Это будет n раз n матрицы. Теперь, учитывая матрицу T размером m_1, умноженной на m_2, я бы хотел вычислить тензор D формы (m_1, m_2, n, n), заданный для 0 <= i <m_1, 0 <= j <m_2, </p>

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

Как эффективно получить этот тензор?

Answer 1

Одним из сравнительно простых путей будет использование einsum.

Пример:

>>> 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]]]])

Вы можете немного ускорить это, используя ключевое слово out, чтобы избежать одногокопия:

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