Broadcasting 2D array to 4D array in numpy

user_17965789 0
Viewed 138

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?

array-broadcasting numpy arrays python
2 Answers

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)
user_29887300

I found the solution. Python broadcasts from the rightmost dimension and works its way to the left (source). By transposing the first two dimensions and the last two dimensions:

T = np.transpose(T, (2,3,0,1))

It will then broadcast the way I expected. After that, the resulting array can be transposed again to recover the original shape:

Tx = x*T
Tx = np.transpose(Tx, (2,3,0,1))
user_179657890
Related