Turn for loop over NumPy array into slice

Question

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?

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

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.