Skip to main content
 首页 » 编程设计

python之大型数组上的函数应用和归约

2025年05月04日10sky-heaven

我有两个 numpy 数组,X 和 Y,它们的形状分别是 X.shape == (m,d) 和 Y.shape == (n,d),其中 m、n 和 d 是非平凡的大小。我需要创建第三个数组 Z,其形状为 Z.shape == (m,n)。

元素 Z[i,j] 是 f(X[i,k],Y[j,k]) for k in range(d)< 的结果 然后对所有 k 求和,对于一些非线性 f。

最明显的方法是这样做:

Z = numpy.zeros((m,n), dtype = numpy.float64) 
for i in range(m): 
    for j in range(n): 
        Z[i,j] += (f(X[i,:],Y[j,:])).sum() # I can compose f from ufuncs 

但我真正想问的是,是否有某种巧妙的广播技巧可用于计算 Z,这将:

  1. 尽可能利用 numpy 的优化
  2. 执行此操作时不要将形状为 (n,m,d) 的数组放入内存中(n*m double 将适合内存,但 n* m*d double 不会)

有谁知道这样做的方法吗?提前致谢。

请您参考如下方法:

这是您不想要的解决方案,我将其包括在内是因为我相信这是针对您的问题的“规范”解决方案。

# A simple function of x, y 
def f(x, y): 
    return 2*x + 3*y**2 
 
x = x.reshape((m, 1, d)) 
y = y.reshape((1, n, d)) 
temp = f(x, y) 
Z = temp.sum(2) 

如果您想避免创建临时数组 temp,它非常大,您可以尝试遍历 d 维度。在某些情况下,以下循环的开销会非常小,您将获得几乎相同的性能,但内存使用量要少得多。

Z = np.zeros((m, n)) 
for i in range(d): 
    Z += f(x[:, :, i], y[:, :, i]) 

如果有帮助,请告诉我。