How to multiply a four dimensional python numpy array by a vector indexed by second dimension

user_4798217 0
Viewed 388

I have a four-dimensional array T (shape=(361,30,100,257)) that I need to scale by a 30-element vector, P. I need the resulting vector to have the same shape as the original T array, and I need the scaling to be indexed by the second element of T. I have tried the obvious:

for i in range(len(PressFrac)):
  theta = T[:,i,:,:]*(PressFrac[i])

but the resulting theta loses the second dimension.

What I need is the first, third, and fourth dimensions multiplied by PressFrac[i] and the second dimension (which is an array called "lev") preserved.

I have tried using np.expand_dims and np.insert below the loop to add back in the second dimension but I get various broadcast errors. I have tried the following and combinations thereof:

theta = np.expand_dims(theta, axis = 1)
np.insert(theta,1,lev,axis=1)
theta = np.stack(theta, lev, axis =1)
theta = np.array(theta)[:,np.newaxis,:,:]

Any ideas would be great!

numpy arrays python
2 Answers

You can just write

T*P[:, np.newaxis, np.newaxis]

to get back a version of T with the second dimension scaled by P. NumPy's broadcasting rules line up the dimensions of arrays on the right, so you need to add two singleton dimensions to the end of P with np.newaxis like above. Then, we're multiplying T with something that has shape (30,1,1), which is allowed.

user_11030440

@macroeconomist's answer is totally correct, and has a nice explanation of the broadcasting rules to boot. The only thing I have to add is that you can make the syntax ever-so-slightly more concise by using None instead of np.newaxis. I wanted to see the broadcasting in action, so I wrote a short demonstration script:

import numpy as np

T = np.arange(4*5*6*7).reshape(4,5,6,7)
w = np.random.randint(1, 10, size=5)

res = T*w[:, None, None]

print('res shape\n%s\n' % (res.shape,))

print('T[:, 2, :, 0]\n%s\n' % T[:, 2, :, 0])
print('T[:, 2, :, 0] * w[2]\n%s\n' % (T[:, 2, :, 0]*w[2]))
print('res[:, 2, :, 0]\n%s\n' % res[:, 2, :, 0])

Output:

res shape
(4, 5, 6, 7)

T[:, 2, :, 0]
[[ 84  91  98 105 112 119]
 [294 301 308 315 322 329]
 [504 511 518 525 532 539]
 [714 721 728 735 742 749]]

T[:, 2, :, 0] * w[2]
[[ 588  637  686  735  784  833]
 [2058 2107 2156 2205 2254 2303]
 [3528 3577 3626 3675 3724 3773]
 [4998 5047 5096 5145 5194 5243]]

res[:, 2, :, 0]
[[ 588  637  686  735  784  833]
 [2058 2107 2156 2205 2254 2303]
 [3528 3577 3626 3675 3724 3773]
 [4998 5047 5096 5145 5194 5243]]
user_4254580
Related