Применить точечный продукт ко всем столбцам в трехмерном массиве

Question

У меня есть 3D-массив, который мне нужно преобразовать в одном направлении.У меня есть матрица map_y, которая содержит преобразование в одном направлении, которое я применяю, используя np.dot.Можно ли избежать вложенного цикла for в приведенном ниже коде?

import numpy as np

nx, ny, nz = 64, 32, 24
nyc = 11

a = np.random.rand(nz, ny, nx)

ac = np.empty((nz, nyc, nx))

map_y = np.random.rand(nyc, ny)

# Can we do this in a single numpy function?
for k in range(nz):
    for i in range(nx):
        ac[k,:,i] = np.dot(map_y, a[k,:,i])

Answer 1

Использование np.einsum -

ac = np.einsum('ijk,lj->ilk',a,map_y)

Бит einsum объяснение проблемы:

Оси можно представить схематично, например, так:

a      : i      x j x k
map_y  :     l  x j
output : i x l      x k

Следовательно, j выровнен и также суммирован для выхода, в то время как остальные «выбиты», что дает нам желаемый результат.

Для получения дополнительной информациисм. связанные документы.

Кроме того, поэкспериментируйте с флагом optimize в np.einsum, установив его как True для использования BLAS.