Эффективный расчет поэлементного матричного произведения с индексно-зависимой функцией в Python

Question

У меня есть интересная проблема, когда я хочу вычислить сумму поэлементного произведения трех матриц

При вычислении \ mathbf { p} _ {ijk} и c_ {ijk} можно сделать априори, у меня проблема с f_ {ijk} (x, y, z). Элементы этой матрицы являются многомерными многочленами, которые зависят от индексов матрицы, поэтому numpy .vectorize не может быть тривиально применен. Лучше всего при решении этой проблемы рассматривать (i, j, k) как дополнительные переменные, так что numpy .vectorize затем применяется к 6-мерному, а не к 3-мерному входу. Однако я не уверен, существуют ли более эффективные или альтернативные способы.

Answer 1

Это простой способ эффективно реализовать эту формулу:

import numpy as np

np.random.seed(0)
l, m, n = 4, 5, 6
x, y, z = np.random.rand(3)
p = np.random.rand(l, m, n)
c = np.random.rand(l, m, n)
i, j, k = map(np.arange, (l, m, n))
xi = (x ** (l - i)) * (x ** l)
yj = (y ** (m - j)) * (y ** m)
zk = (z ** (n - k)) * (z ** n)
res = np.einsum('ijk,ijk,i,j,k->', p, c, xi, yj, zk)
print(res)
# 0.0007208482648476157

Или даже немного более компактный:

import numpy as np

np.random.seed(0)
l, m, n = 4, 5, 6
x, y, z = np.random.rand(3)
p = np.random.rand(l, m, n)
c = np.random.rand(l, m, n)
t = map(lambda v, s: (v ** (s - np.arange(s))) * (v ** s), (x, y, z), (l, m, n))
res = np.einsum('ijk,ijk,i,j,k->', p, c, *t)
print(res)
# 0.0007208482648476157

Используя np.einsum, вы минимизируете необходимость в промежуточных массивах, поэтому должно быть быстрее, чем сначала создать f (что вы можете получить, например, как f = np.einsum('i,j,k->ijk', xi, yj, zk)), умножить p, c и f, а затем суммировать результат.