Как сделать np.einsum быстрее, когда на входе много одинаковых массивов? (Или любой другой более быстрый метод)

Question

У меня есть фрагмент кода типа:

nnt = np.real(np.einsum('xa,xb,yc,yd,abcde->exy',evec,evec,evec,evec,quartic))

, где evec (скажем) массив L x L np.float32, а quartic - это L x L x L x L x T np.complex64 array.

Я обнаружил, что эта процедура довольно медленная.

Я подумал, что, поскольку все evec идентичны, может быть более быстрый способ это делаете?

Заранее спасибо.

Answer 1

Для начала вы можете повторно использовать первый расчет:

evec2 = np.real(np.einsum('xa,xb->xab',evec,evec))
nnt = np.real(np.einsum('xab,ycd,abcde->exy',evec2,evec2,quartic))

А если вам не важна память, а нужна только производительность:

evec2 = np.real(np.einsum('xa,xb->xab',evec,evec))
nnt = np.real(np.einsum('xab,ycd,abcde->exy',evec2,evec2,quartic,optimize=True))