Python 计算三维和二维NumPy阵列的加权和

假设我有一个形状为（n，3）的2D NumPy数组a和形状为（x，y，n）的3D数组B。我想创建一个形状为（x，Y，3）的三维阵列Y，这是一个RGB图像

B的第三维包含概率[0，1]，其总和为1。数组A是RGB颜色的列表（每种概率一种颜色）。 对于B的第三维中的每个元素C，我想计算A*C的和，从而得到一个RGB颜色向量

示例：

给定一个颜色映射作为

A = [(255   0   0),
     (  0 255   0),
     (  0   0 255)]

以及点（x，y）处B的第三维中的元素C，如

我想计算

C'=sum（A*C）

as

  (255   0   0) * 0.2
+ (  0 255   0) * 0.6
+ (  0   0 255) * 0.2
---------------------
  ( 51 153  51)

分配

Y[x, y, :] = C'

我知道我可以用for循环在x和y上迭代，一次计算每个元素，但我想知道这是否可以用矢量化的方式来完成，而不必自己迭代数组（主要是出于性能原因）

该方法对此很方便：

np.einsum('ij,kli->klj', A, B)

这个符号表示：将A[i，j]乘以B[k，l，i]，然后在

i

上求和；将结果放入单元格

[k，l，j]

例如：

A = np.array([(255, 0, 0), (0, 255, 0), (0, 0, 255)])
B = np.array([[[0.2, 0.6, 0.20], [0.2, 0.2, 0.60]], [[0.4, 0.4, 0.2], [0.3, 0.3, 0.4]]])
Y = np.einsum('ij,kli->klj', A, B)

然后是

Y

array([[[  51. ,  153. ,   51. ],
        [  51. ,   51. ,  153. ]],

       [[ 102. ,  102. ,   51. ],
        [  76.5,   76.5,  102. ]]])

---

您正在将

A

中的第一个轴与

B

中的第三个轴进行求和，而其余轴则展开。这是一个完美的设置，可以利用基于BLAS的张量矩阵乘法，就像这样-

C = np.tensordot(B,A,axes=((2),(0)))

B.reshape(-1,n).dot(A).reshape(x,y,3)

我们还可以手动将形状重塑为

2D

，并将矩阵乘法用于

2D

：

np.dot

，如下所示-

C = np.tensordot(B,A,axes=((2),(0)))

B.reshape(-1,n).dot(A).reshape(x,y,3)

请注意，

B.dot（A）

也可以工作，但速度会慢一些，很可能是因为它会在

B

的第一个轴上循环，同时针对

A

对每个2D切片执行

2D

矩阵乘法

运行时测试-

In [180]: np.random.seed(0)
     ...: x,y,n = 100,100,100
     ...: A = np.random.rand(n,3)
     ...: B = np.random.rand(x,y,n)

# @Crazy Ivan's soln
In [181]: %timeit np.einsum('ij,kli->klj', A, B)
100 loops, best of 3: 4.21 ms per loop

In [182]: %timeit np.tensordot(B,A,axes=((2),(0)))
1000 loops, best of 3: 1.72 ms per loop

In [183]: np.random.seed(0)
     ...: x,y,n = 200,200,200
     ...: A = np.random.rand(n,3)
     ...: B = np.random.rand(x,y,n)

# @Crazy Ivan's soln
In [184]: %timeit np.einsum('ij,kli->klj', A, B)
10 loops, best of 3: 33.2 ms per loop

In [185]: %timeit np.tensordot(B,A,axes=((2),(0)))
100 loops, best of 3: 15.3 ms per loop

---

工作起来很有魅力。我接受这个答案，因为它提出了更快的解决方案。