Умножение путем трансляции строк 2D-массива на каждый фрагмент 3D-массива с использованием np.einsum - PullRequest
Новая
       52

Умножение путем трансляции строк 2D-массива на каждый фрагмент 3D-массива с использованием np.einsum

0 голосов
/ 02 мая 2020

У меня есть массивы A и B: 

>>> import numpy as np

>>> A = np.ones((3,3,2))

>>> B = np.array([
    [0,0],
    [1,1],
    [2,2],
])

Я хочу умножить каждую строку B на каждый фрагмент A, чтобы каждая строка B получает широковещательную рассылку по каждому срезу A, то есть: 

>>> np.array([A_slice*B_row for A_slice, B_row in zip(A, B)])
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

Мне нужен наиболее эффективный метод для этого, который, я считаю, может использовать np.einsum (однако, если вы считаете, что использование другого метода быстрее, например, тот, который я упомянул ниже, скажите мне).

Я попытался выполнить следующее: 

>>> np.einsum('ijk,lk->ijk', A, B)
[[[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]]

Как вы можете видеть, это явно не тот же вывод, что и выше.

Другое решение, о котором я могу подумать: 

>>> A*B[:,np.newaxis,:].repeat(3, axis=1)
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

, которое дает правильный вывод, но я все еще очень хочу знать, как это сделать с np.einsum

РЕДАКТИРОВАТЬ: Уоррен Векессер отметил в комментариях, что решение выше можно упростить до A*B[:,np.newaxis,:], который является самым чистым решением из всех, что я видел, без использования np.einsum.

1 Ответ

0 голосов
/ 02 мая 2020

numpy.einsum soution: 

C = np.einsum('ijk,jk->jik', A, B)

Использование elipsis: 

C = np.einsum('ij...,j...->ji...', A, B)

Вывод 

[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]
...