将NumPy数组上的for循环转换为切片

Question

我有这个 for 循环来计算 3D numpy 数组的元素：

A = np.zeros((N - 2, N - 2, N - 2))
for i in range(1, N - 1):
    for j in range(1, N - 1):
        for k in range(1, N - 1):
            A[i - 1, j - 1, k - 1] = (B[i,j,k] * dy[j] * dz[k] + B[i,j-1,k] * dy[j-1] * dz[k]
                                        + B[i,j,k-1] * dy[j] * dz[k-1] + B[i,j-1,k-1] * dy[j-1] * dz[k-1]) / (4 * dx[i])

其中 B 是形状 (N, N, N) numpy 数组，dx、dy 和 dz 是长度 N-1 numpy arrays。 对于大 N 这真的很慢所以我尝试了这样的事情：

A = (B[1:-1, 1:-1, 1:-1] * dy[1:] * dz[1:] + B[1:-1, :-2, 1:-1] * dy[:-1] * dz[1:]
                                        + B[1:-1, 1:-1, :-2] * dy[1:] * dz[:-1] + B[1:-1, :-2, :-2] * dy[:-1] * dz[:-1]) / (4 * dx[1:])

除非 dx、dy 和 dz 是常数，否则这不起作用。 我也试过：

dX, dY, dZ = np.meshgrid(dx, dy, dz)
A = (B[1:-1, 1:-1, 1:-1] *  dY[1:, 1:, 1:] * dZ[1:, 1:, 1:] + B[1:-1, :-2, 1:-1] * dY[:-1, :-1, :-1] * dZ[1:, 1:, 1:]
                            + B[1:-1, 1:-1, :-2] * dY[1:, 1:, 1:] * dZ[:-1, :-1, :-1] + B[1:-1, :-2, :-2] * dY[:-1, :-1, :-1] * dZ[:-1, :-1, :-1]) / (4 * dX[:-1, :-1, :-1])

但这也行不通。

关于如何做到这一点的任何想法？

Answer 1

A[i - 1, j - 1, k - 1] = (B[i,j,k] * dy[j] * dz[k] + B[i,j-1,k] * dy[j-1] * dz[k]
                                    + B[i,j,k-1] * dy[j] * dz[k-1] + B[i,j-1,k-1] * dy[j-1] * dz[k-1]) / (4 * dx[i])

让我们将迭代更改为 (0,N-2)

A[i, j, k] = (B[i+1,j+1,k+1] * dy[j+1] * dz[k+1] + B[i+1,j,k+1] * dy[j] * dz[k+1]
                                    + B[i+1,j+1,k] * dy[j+1] * dz[k] + B[i+1,j,k] * dy[j] * dz[k]) / (4 * dx[i+1])

并分别处理 4 个术语

B[1:,1:,1:]*dy[None,1:,None]*dz[None,None,1:]

B[1:,:-1,1:]*dy[None,:-1,None]*dz[None,None,1:]

etc

我没有测试就写了这些，所以可能会有错误，但希望这足以让你开始。