您可以使用 3 个单独的 1D numpy arrays 来使用矢量化操作 3D 数组吗？

Question

I am trying multiply a specific location of an array by certain value, where the location is determined by the index and the value of the num array. The certain value comes from the same index position of the multiplier array. We only want to apply this multiplier if the needs_multiplier is value at that index position is true. I think the code will do a better job explaining this. I am trying to vectorize this and avoid the for loop.

import numpy as np

data = np.array([[[ 2.,  2.,  2.,  2.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.]],

                 [[ 1.,  1.,  1.,  1.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.]],

                 [[ 3.,  3.,  3.,  3.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.]],

                 [[ 5.,  5.,  5.,  5.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.],
                  [ 0.,  0.,  0.,  0.]]])

needs_multiplier = np.array([True, True, False, True])
num = np.array([1, 2, 2, 3])
multipler = np.array([0.5, 0.6, 0.2, 0.3])


for i, cfn in enumerate(num):
    if needs_multiplier[i]:
        data[i, 1, cfn] = multipler[i] * data[i, 0, cfn]
        data[i, 2, cfn] = data[i, 0, cfn]-data[i, 1, cfn]

print(data) # this is the result I am looking for

[[[2.  2.  2.  2. ]
  [0.  1.  0.  0. ]
  [0.  1.  0.  0. ]
  [0.  0.  0.  0. ]]

 [[1.  1.  1.  1. ]
  [0.  0.  0.6 0. ]
  [0.  0.  0.4 0. ]
  [0.  0.  0.  0. ]]

 [[3.  3.  3.  3. ]
  [0.  0.  0.  0. ]
  [0.  0.  0.  0. ]
  [0.  0.  0.  0. ]]

 [[5.  5.  5.  5. ]
  [0.  0.  0.  1.5]
  [0.  0.  0.  3.5]
  [0.  0.  0.  0. ]]]

Answer 1

num can be used as index array after selecting "active" values with num[needs_multiplier]

Then vectorizing the expressions is pretty straight forward:

b = needs_multiplier
num_b = num[needs_multiplier]

data[b, 1, num_b] = multipler[b] * data[b, 0, num_b]
data[b, 2, num_b] = data[b, 0, num_b] - data[b, 1, num_b]