Question

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

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)

有更快的方法吗？我知道它可以与向量化函数一起使用，但是我无法解决它。

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

Answer 1

向量化循环，并使用broadcasting：

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

编辑：

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

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

Answer 2

您可以使用numpy.fromfunction()进行此操作：

import numpy as np

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

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

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

updatedArray = np.fromfunction(calcP, (4,4,3), a=a,b=b,ps=ps, dtype=int)
print (updatedArray)

请注意，我对您的功能calcP做了一些修改，以适应实际情况。另外，我添加了0.0，以确保输出数组将是float s而不是int s。

还要注意，fromfunction()的第二个参数仅指定网格的形状，在该网格上将调用功能calcP()。

输出（由于randint的不同，每次都会有所不同）：

[[[  8.   5.   3.]
  [  9.   6.  12.]
  [ 10.   7.  21.]
  [ 11.   8.  30.]]

 [[ 12.   9.  39.]
  [ 13.  10.  48.]
  [ 14.  11.  57.]
  [ 15.  12.  66.]]

 [[ 16.  13.  75.]
  [ 17.  14.  84.]
  [ 18.  15.  93.]
  [ 19.  16. 102.]]

 [[ 20.  17. 111.]
  [ 21.  18. 120.]
  [ 22.  19. 129.]
  [ 23.  20. 138.]]]

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

2 个答案: