python - Python：将函数应用于numpy 3d数组中的每个条目

Question

我想在我的 3d numpy 数组上应用一个（更复杂的？）函数，形状为 x,y,z = (4,4,3)。假设我有以下数组：

array = np.arange(48)
array = array.reshape([4,4,3])

现在我想在数组的每个点上调用以下函数：

p(x,y,z) = a(z) + b(z)*ps(x,y)

假设 a 和 b 是以下 1d 数组，分别是 ps 一个 2d 数组。

a = np.random.randint(1,10, size=3)
b = np.random.randint(1,10, size=3)
ps = np.arrange(16)
ps = ps.reshape([4,4])

我直观的方法是遍历我的数组并在每个点上调用函数。它可以工作，但当然它太慢了：

def calcP(a,b,ps,x,y,z):
    p = a[z]+b[z]*ps[x,y]
    return p

def stupidLoop(array, a, b, ps, x, y, z):
    dummy = array
    for z in range (0, 3):
        for x in range (0, 4):
            for y in range (0, 4):
                dummy[x,y,z]=calcP(a,b,ps,x,y,z)
    return dummy

updatedArray=stupidLoop(array,a, b, ps, x, y, z)

有更快的方法吗？我知道它适用于矢量化函数，但我无法用我的方法弄清楚。

我实际上并没有尝试使用这些数字。这只是为了举例说明我的问题。它来自气象学世界，有点复杂。

标签： pythonarraysnumpy

Answer 1

向量化循环，并使用广播：

a.reshape([1,1,-1]) + b.reshape([1,1,-1]) * ps.reshape([4,4,1])

编辑：

感谢@NilsWerner 提供了一种更常见的评论方式：

a + b * ps[:, :, None]