将索引的函数应用于numpy数组的所有元素

Question

我正在寻找一种方法，将函数应用于numpy数组的所有元素。函数不仅接收值，还接收元素的索引作为参数。目标是在大的2或3码上表现出色。阵列。

(我知道对于只接收元素值的函数，有几种方法可以做到这一点(

要替换的代码是

def foo(x, i, j)
return (i*x)**2 - (j*x)**3  # or some other fancy stuff
...
arr = np.zeros((nx, ny))
...
# nested loop to be replaced, e.g. via vectorization
for i in range(nx):
for j in range(ny):
arr[i,j] = foo(arr[i,j], i, j)               

Answer 1

您可以使用简单的广播规则来实现这一点，方法是使用具有适当形状的适当生成的索引，使标准广播规则与输入的形状相匹配。

这可以手动生成(例如，使用np.arange()和np.reshape()的组合(，也可以更简洁地使用np.ogrid()生成。

import numpy as np

np.random.seed(0)

def foo(x):
n, m = arr.shape
i, j = np.ogrid[:n, :m]
return (i * x) ** 2 - (j * x) ** 3

n, m = 2, 3
arr = np.random.random((n, m))

foo(arr)
# array([[ 0.        , -0.36581638, -1.7519857 ],
#        [ 0.29689768,  0.10344439, -1.73844954]])

这种方法可能需要用于中间结果的大型临时存储器阵列。

---

通过使用像numba:这样的JIT编译器保持显式循环加速，可以获得更有效的方法

import numba as nb

@nb.njit
def foo_nb(arr):
n, m = arr.shape
out = np.empty((n, m), dtype=arr.dtype)
for i in range(n):
for j in range(m):
x = arr[i, j]
out[i, j] = (i * x) ** 2 - (j * x) ** 3
return out

foo_nb(arr)
# array([[ 0.        , -0.36581638, -1.7519857 ],
#        [ 0.29689768,  0.10344439, -1.73844954]])