Search code examples
pythonarraysnumpyarray-broadcasting

Broadcasting 2D array to 4D array in numpy

I have a 2D array x of shape (48, 7), and a 4D array T of shape (48, 7, 48, 7). When I multiply x * T, python broadcasts the dimensions, but not in the way I expected (actually, I don´t understand how it is broadcasting). The following loop would achieve what I want:

for i in range(48):
    for j in range(7):
        Tx[i, j, :, :] = x[i, j] * T[i, j, :, :]

Where Tx is an array of shape (48, 7, 48, 7). My question is, is there a way to achieve the same result using broadcasting?

Solution

  • Broadcasting aligns trailing dimensions. In other words, x * Tx is doing this:

    for i in range(48):
    for j in range(7):
        Tx[:, :, i, j] = x[i, j] * T[:, :, i, j]

    To get the leading dimensions to line up, add unit dimensions to x:

    Tx = x[..., None, None] * T

    Alternatively, you can use np.einsum to specify the dimensions explicitly:

    Tx = np.einsum('ij,ij...->ij...', x, T)