Python 沿两个不同大小数组的共享维度进行元素相乘的优雅方法？

假设我有两个数组，比如形状（X，Y）的数组A和形状（X，Y，Z）的数组B。我的目标是获得由A和B沿Z方向的元素相乘组成的数组C。我可以使用for循环来实现这一点，例如：

import numpy as np
C = np.zeros(B.shape)
for z in range(Z):
    C[:,:,z] = A*B[:,:,z]

---

然而，在理想情况下，我希望避免使用for循环。您知道更优雅的方法吗？

您可以使用形状数组

（x，y）

和

（z，x，y）

进行元素相乘

例如：

a = np.arange(12).reshape((4, 3))
b = np.ones((4, 3, 2))

b[:, :, 0] = 10

x = a * b # :( error

然而，如果我们这样做

c = np.ones((2, 4, 3))
c[0, :, :] = 10
y = a * c

这很有效，我们得到了

array([[[  0.,  10.,  20.],
        [ 30.,  40.,  50.],
        [ 60.,  70.,  80.],
        [ 90., 100., 110.]],

       [[  0.,   1.,   2.],
        [  3.,   4.,   5.],
        [  6.,   7.,   8.],
        [  9.,  10.,  11.]]])

现在，要移动阵列的轴，可以使用函数

b_moved = np.moveaxis(b, -1, 0) # move last axis of b to first position

x = a * b_moved # In the shape of b_moved

y = np.moveaxis(x, 0, -1) # Move first axis of x to the last position

这就给了

x: shape 2x4x3
array([[[  0.,  10.,  20.],
        [ 30.,  40.,  50.],
        [ 60.,  70.,  80.],
        [ 90., 100., 110.]],

       [[  0.,   1.,   2.],
        [  3.,   4.,   5.],
        [  6.,   7.,   8.],
        [  9.,  10.,  11.]]])

y: shape 4x3x2
array([[[  0.,   0.],
        [ 10.,   1.],
        [ 20.,   2.]],

       [[ 30.,   3.],
        [ 40.,   4.],
        [ 50.,   5.]],

       [[ 60.,   6.],
        [ 70.,   7.],
        [ 80.,   8.]],

       [[ 90.,   9.],
        [100.,  10.],
        [110.,  11.]]])

总之，您可以编写一个函数来完成以下任务：

def emult3(arr2d, arr3d):
    # make sure shapes are compatible
    assert (arr2d.ndim, arr3d.ndim) == (2, 3)
    assert arr2d.shape == arr3d.shape[0:2] 

    # do the calculations
    x = arr2d * np.moveaxis(arr3d, -1, 0)
    return np.moveaxis(x, 0, -1)

---

太好了，这正是我想要的！乘法运算的方式真奇怪。我甚至可以用np.moveaxis（a*np.moveaxis（B，-1,0），0，-1）将其简化为一行。多谢各位！