将依赖于其输入值位置的函数映射到 numpy 数组答案

【问题标题】：Mapping a function that depends on the location of its input value to a numpy array将依赖于其输入值位置的函数映射到 numpy 数组
【发布时间】：2018-08-27 00:21:03
【问题描述】：

假设我们有一个数组，

arr = np.random.rand(3,3)

通常在将函数映射到数组时，我们只对数组元素的值感兴趣，即

f = lambda val : val**2
arr_squared = f(arr)

但是如果我们函数的输出取决于输入值在数组中的位置，即

f = lambda x,y,val : x*y*val

现在我正在使用 meshgrids 和 ravel。

X, Y = np.arange(arr.shape[0]), np.arange(arr.shape[1])
X, Y = np.meshgrid(X,Y)

result = np.zeros(arr.shape)

for x,y in zip(np.ravel(X), np.ravel(Y)):
    result[x,y] = f(x,y,arr[x,y])

这可行，但速度很慢。我很难弄清楚是否有更好/更快的方法来做到这一点，并且在线搜索没有产生太多有用的信息。

【问题讨论】：

标签： python arrays numpy dictionary mapping

【解决方案1】：

原则上，索引的行为与任何其他参数一样。

这一切都归结为您的函数是否矢量化。

如果是，如您的示例所示：

>>> Y, X = np.ogrid[(*map(slice, arr.shape),)]
>>> def f(X, Y, val): return X*Y*val
... 
>>> f(X, Y, arr)
array([[0.        , 0.        , 0.        ],
       [0.        , 0.92796409, 0.20353397],
       [0.        , 1.01294541, 1.30677315]])

如果不是，例如

>>> def g(X, Y, val): return X+Y if val>0.5 else X-Y
... 
>>> g(X, Y, arr)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in g
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

快速修复是 np.vectorize，但这并不快：

>>> np.vectorize(g)(X, Y, arr)
array([[0, 1, 2],
       [1, 2, 1],
       [2, 3, 0]])

如果可能，最好手动矢量化：

>>> def gv(X, Y, val): return X + (2*(val>0.5)-1) * Y
... 
>>> gv(X, Y, arr)
array([[0, 1, 2],
       [1, 2, 1],
       [2, 3, 0]])

关于你对网格网格的使用的旁注：对于索引，你可以使用它来缩写

Y, X = np.indices(arr.shape)

在上面的示例中，我使用了利用广播来节省内存的开放网格。

【讨论】：