向量化numpy分配

Question

我通过在其他 numpy 数组中查找值来为 numpy 数组赋值。这些数组可能具有不同的索引。这是一个例子：

import numpy as np
A=1; B=2; C=3; D=4; E=5
X = np.random.normal(0,1,(A,B,C,E))
Y = np.random.normal(0,1,(A,B,D))
Z = np.random.normal(0,1,(A,C))
Result = np.zeros((A,B,C,D,E))
for a in range(A):
    for b in range(B):
        for c in range(C):
            for d in range(D):
                for e in range(E):
                    Result[a,b,c,d,e] = Z[a,c] + Y[a,b,d] + X[a,b,c,e]

优化此代码的最佳方法是什么？我可以使用 Result[a,b,c,d,:] = Z[a,c] + Y[a,b,d] + X[a,b,c,:] 删除 E for 循环。但是那么如何删除其余的循环呢？我还认为我可以在分配之前操纵 X、Y、Z，以便它可以轻松地与 Result 的维度合并。必须有更优雅的方式。感谢您的提示。

Answer 1

这是一种方法：

Result = Z[:,None,:,None,None] + Y[:,:,None,:,None] + X[:,:,:,None,:]

为了生成这个矢量化版本，我所做的只是将各种索引替换为 X、Y 和 Z，并使用完整的 a,b,c,d,e 样式索引，在找到缺失索引的地方插入 None。比如Y[a,b,d]变成Y[a,b,None,d,None]，向量化成Y[:,:,None,:,None]。

在 numpy 中，None 的索引告诉数组假装它有一个额外的轴。这不会改变数组的大小，但会改变操作的广播方式，这正是我们需要的。查看numpy broadcasting docs 了解更多信息。