您可以使用 1D 数组来索引 3D 数组并在单个轴上应用乘法吗?

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

我正在尝试矢量化这个过程。我有一个 4x4x4 矩阵。我正在尝试将 [i,0,num[i]] 中的值乘以乘数 i,其中 i 是范围 (data.shape[0]) 并将结果放入 [i,0,num[i] ].我想弄清楚如何对其进行矢量化。我只能循环进行。该代码有望更有意义。代码不正确,但我想我可能很接近。

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

这个有效:

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]

输出:

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