根据单元格索引填充numpy数组

Question

我正在尝试创建一个3d数组，其单元格条目将根据单元格索引计算。具体来说，我想要(i,j,k) = sqrt(i+j+k)。

这很容易通过以下循环来完成：

N=10
A=np.zeros((N,N,N))
for i in range(N):
for j in range(N):
for k in range(N):
A[i][j][k] = np.sqrt(i+j+k)

我想知道numpy是否有使这些嵌套for循环变得多余的内置函数。

Answer 1

最简单、最高效的方法是使用np.ogrid打开网格，然后执行相关操作-

I,J,K = np.ogrid[:N,:N,:N]
A = np.sqrt(I+J+K)

或者用np.sum对一行的开放网格的广播求和

A = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))

相关：General workflow on vectorizing loops involving range iterators。

Answer 2

您可以使用np.arange，然后使用np.newaxis来创建不同的维度。用简单的sum和np.sqrt做作业后：

arr = np.arange(N)
A = np.sqrt(arr + arr[:,np.newaxis]+ arr[:,np.newaxis,np.newaxis])

你会得到同样的结果：

N = 10
arr = np.arange(N)
A = np.sqrt(arr + arr[:,np.newaxis]+ arr[:,np.newaxis,np.newaxis])
B = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))
print ((A==B).all())
#True

这种方法比使用np.ogrid:快一点

N = 10
%timeit arr = np.arange(N); A = np.sqrt(arr + arr[:,np.newaxis]+ arr[:,np.newaxis,np.newaxis])
#18.6 µs ± 3.29 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit A = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))
#58.5 µs ± 8.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 3

这对于大型N来说更快，但可能被视为作弊；-(

它充分利用了高度规则和重复的模式，节省了大量的平方根求值。

def cheat(N):
values = np.sqrt(np.arange(3*N-2))
result = np.lib.stride_tricks.as_strided(values, (N, N, N), 3*values.strides)
return np.ascontiguousarray(result)

如果你可以使用非连续的只读(通过所有实际手段(视图，这可能会更快：

def cheat_nc_view(N):
values = np.sqrt(np.arange(3*N-2))
return np.lib.stride_tricks.as_strided(values, (N, N, N), 3*values.strides)

供参考：

def cheek(N):
arr = np.arange(N)
return np.sqrt(arr + arr[:,np.newaxis] + arr[:,np.newaxis,np.newaxis])
>>> np.all(cheek(20) == cheat(20))
True
>>> np.all(cheek(200) == cheat_nc_view(200))
True

时间安排：

>>> timeit(lambda: cheek(20), number=1000)
0.05387042500660755
>>> timeit(lambda: cheat(20), number=1000)
0.020798540994292125
>>> timeit(lambda: cheat_nc_view(20), number=1000)
0.010791150998556986
>>> timeit(lambda: cheek(200), number=100)
6.823299437994137
>>> timeit(lambda: cheat(200), number=100)
2.0583883369981777
>>> timeit(lambda: cheat_nc_view(200), number=100)
0.0014881940151099116