Turn for loop over NumPy array into slice-CodePudding

I have this for loop to calculate the elements of a 3D numpy array:

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])

Where B is a shape (N, N, N) numpy array and dx, dy and dz are length N-1 numpy arrays. For large N this is really slow so I tried something like this:

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:])

This does not work unless dx, dy and dz are constant. I also tried:

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])

But this also does not work.

Any ideas on how to do this?

CodePudding user response：

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])

lets change the iterations to (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])

and tackle the 4 terms separately

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

I wrote these without testing, so there might be errors, but hopefully that's enough to get you started.