I’m trying to cast a m x n matrix with a p x n x q matrix in a particular way. In effect, I want to replace each element of the p x n x q array with the m x n array added to the element it replaced.
Formally, given mn matrix X and pnq matrix Y, I want to fill a qnmp (not pmnq) numpy such that,
How would I go about doing this in numpy without creating a separate array and filling it out with the above relation?
>Solution :
IIUC, align the axes for broadcasting, and then use transposing to move the axes into the desired output shape.
import numpy as np
M, N, P, Q = 2, 3, 4, 5
X = np.arange(M * N).reshape(M, N)
Y = np.arange(P * N * Q).reshape(P, N, Q)
# X: (1, M, N, 1)
# Y: (P, 1, N, Q)
# =============== Broadcast
# Z: (P, M, N, Q)
Z = X[None, :, :, None] + Y[:, None, :, :]
# Transpose the output array to achieve
# Z: (Q, N, M, P)
Z = Z.T
# Or alternatively do the transposing of X and Y first, then align axes.