Numpy вычесть два массива одинаковой ndim, но разной формы

Question

Итак, у меня есть два ndarrays:

A с формой (N, a, a), стек из N массивов формы (a, a) в основном

B с формой (8, M, a, a), матрица из 8 x M массивов формы (a, a)

Мне нужно вычесть B из A (AB) так, чтобы полученный массив имел форму (8, M* Н, а, а).Более подробно каждый (M всего) из 8 массивов B необходимо вычесть из каждого массива в A, что приведет к 8 * M * N вычитаниям между (a, a) массивами форм.

Как я могу это сделатьэто в векторизации без петель?Этот поток делает что-то похожее, но в меньших размерах, и я не могу понять, как его расширить.

Answer 1

A = np.arange(8).reshape(2,2,2)
B = np.ones(shape=(8,4,2,2))

Общее вещание работает, если измерения одинаковы или если одно измерение равно 1, поэтому мы делаем это;

a = A[np.newaxis, :, np.newaxis, :, :]
b = B[:, np.newaxis, :, :, :]

a.shape  # <- (1,2,1,2,2)
b.shape  # <- (8,1,4,2,2)

Теперь, когда вы можете делать вещание

c = a - b
c.shape  # <- (8,2,4,2,2)

И когда вы изменяете форму, компоненты (2x4 = 8) выравниваются.

c.reshape(8,-1,2,2) 

Порядок расположения новых осей требует изменения формы, поэтому будьте осторожны с этим.