CODEWITHSUNDEEP
pythonarraysnumpyvectorization
Rob Lasch
Rob Lasch

Reputation: 107

Can you use a 1D array to index a 3D array and apply multiplication across a single axis?

I am trying to vectorize this process. I have a 4x4x4 matrix. I am trying to multiply the value in [i,0,num[i]] by multiplier i where i is the range(data.shape[0]) and put the result in [i,0,num[i]]. I am trying to figure out how to vectorize that. I can only do it in a loop. The code will hopefully make more sense. The code is not correct, but I think I might be close.

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

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


data[:, 1, num] = multipler * data[:, 0, num]

data[:, 2, num] = data[:, 0, num] - data[:, 1, num]

print(data) # This is the goal output


[[[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.6 0.]
  [0.  0.  2.4 0.]
  [0.  0.  0.  0. ]]

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

Upvotes: 2

Views: 48

Answers (1)

fountainhead
fountainhead

Reputation: 3722

This works:

my_range = range(data.shape[0])
data[my_range, 1, num] = data[my_range, 0, num] * multipler
data[my_range, 2, num] = data[my_range, 0, num] - data[my_range, 1, num]

Output:

print (data)

[[[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.6 0. ]
  [0.  0.  2.4 0. ]
  [0.  0.  0.  0. ]]

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

Upvotes: 1

Related Questions