Indexing numpy arrays only if within bounds
I've realized that using numpy array operations is definitely the way to go to make the code faster. Pairing this with np.pad to add zeros to the edges of a matrix makes this fairly simple.
I have a 2d numpy array psi with shape (nx,ny). I want to create a new array phi of the same shape where for each element phi[i][j] I need to evaluate an expression containing psi[i][j] and neighboring elements psi[i-1][j],psi[i+1][j],psi[i][j+1] and psi[i][j-1],except for edge cases where any of these neighbors are not in the bounds of psi, treat that element as 0 in the expression.
I can implement this using nested for loops and checking for boundary conditions, but I would like to perform this operation as time efficient as possible. I've tried by assigning
phi[1:-1,1:-1] = f(psi[1:-1,1:-1], psi[0:-2,1:-1], psi[2:,1:-1], psi[1:-1,0:-2], psi[1:-1,2:])
but this does not cover edge cases which get messy, so if there were some conditional way to only reference when within bounds else just be 0 it might work. Or, of course, if there is an even more time efficient way that would be better.