Can you use 3 seperate 1D numpy arrays to manipulate a 3D array using vectorization?

Rob Lasch

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. ]]]

user7138814

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]

Collected from the Internet

Please contact [email protected] to delete if infringement.