Векторизованная матрица-Векторное умножение вдоль заданной оси многомерного массива

#arrays #numpy #multidimensional-array #scipy #vectorization

Вопрос:

Предположим,у меня есть массив A размера (M,N,D,D) и вектор x размера (M,N, D). Я хочу создать массив b размера (M,N,D), где b[m, n,:] = np.matmul(A[m, n,:,:], x[m,n,:]). Таким образом, в принципе,матрица D по D, содержащаяся в[m,n,:], и вектор размера D, содержащийся в x[m, n,:], должны быть умножены вместе после стандартного умножения матрицы на вектор. Наивной реализацией цикла было бы

 for m in range(M):  for n in range(N):  b[m,n,:] = np.matmul(A[m,n,:],x[m,n,:])

Очевидно, что это неэффективно для больших N и M (что касается меня; N составляет порядка 100 000, а M-порядка 1000). Мой вопрос: как векторизовать этот цикл? Я пробовал использовать np.tensordot, но я продолжаю получать несоответствия формы, независимо от того, как я пытаюсь его структурировать, предположительно потому, что у A 4 оси, в то время как у x 3 (Обратите внимание, что, насколько я знаю, np.tensordot—правильная функция для использования, я просто неправильно ее использую)

РЕДАКТИРОВАТЬ Попытался использовать np.einsum, но я не получаю ожидаемого результата (который должен быть [5,11]:

1. np.matmul(A[m,n,:,:],x[m,n,:]) вычисляет сумму продуктов в последнем измерении обоих массивов, где x[m,n,:] 1d, и обрабатывается (согласно документам) как x[m,n,:,None] (2-е по последнему). Ключевая matmul концепция — «Последний из A со 2-го по последнее из B». Ведущими измерениями здесь m,n являются «пакетные». ( tensordot не обрабатывает «пакеты»).

Ответ №1:

Вот где np.einsum сияет:

 np.einsum('mnij,mnj-gt;mni', A,B)

Уточните у

 np.allclose(np.einsum('mnij,mnj-gt;mni', A,x), b)

ВОЗВРАТ True

1. Я попробовал это; хотя конечный результат имеет правильную форму, он не дает правильного результата. Рассмотрим случай A[0,0,:] = np.массив([[1,2],[3,4]]) и x[0,0,:] = np.массив([1,2]). В этом случае я бы ожидал, что b[0,0,:] = np.массив([5,11]), так как это матрично-векторное произведение. Однако выполнение b = np.einsum(‘mndd, mnd-gt;mnd’, A,x) показывает,что b[0,0,:] = np.массив([1,8]).

2. @torola вам следует дважды проверить код. Я вошел [5, 11] в свою систему.

3. Я отредактировал исходное сообщение, чтобы показать свой вывод — может ли ошибка быть связана с тем, как я присваиваю значения A и/или x?

4. @torola моя формула-это 'mnij,mnj-gt;mni' , а не mndd,mnd-gt;mnd то , что является формулой трассировки.

5. Ааааа, теперь это имеет смысл (я думал, что i и j, которые вы использовали, были просто заполнителями и что я должен заменить их на d). Теперь это работает отлично, большое спасибо!

Ответ №2:

Того же результата можно достичь с помощью вещания:

 x1 = x.reshape(M,N,d,1) b1 = np.matmul(A,x1) b2 = b1.reshape(M,N,d)