Indexing numpy arrays only if within bounds
Question:
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.
Answers:
This problem smells like finite differences. Your best bet is to write a (fast, possibly recursive) loop for the inner points, and then loop over the boundary points separately, imposing the desired boundary conditions there. Obviously, the other way around also works: start by assigning boundary points, then loop over inner points.
That said, if you are having issues with speed (probably because your grid is gigantic), you may want to do a few optimizations, as 2d arrays in python are S L O W:
-
try reversing the order of looping: in python (NumPy, in case you are using that), 2d arrays are traversed by rows first. You may want to experiment with that at least.
-
try allocating your 2d thing as a big 1d chunk where its unique index is n = i + nx * j, with i,j your original 2d indices. Again, experiment with running the new index n along rows vs columns first.
These two suggestions combined should give you a massive speedup.
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.
This problem smells like finite differences. Your best bet is to write a (fast, possibly recursive) loop for the inner points, and then loop over the boundary points separately, imposing the desired boundary conditions there. Obviously, the other way around also works: start by assigning boundary points, then loop over inner points.
That said, if you are having issues with speed (probably because your grid is gigantic), you may want to do a few optimizations, as 2d arrays in python are S L O W:
-
try reversing the order of looping: in python (NumPy, in case you are using that), 2d arrays are traversed by rows first. You may want to experiment with that at least.
-
try allocating your 2d thing as a big 1d chunk where its unique index is
n = i + nx * j, withi,jyour original 2d indices. Again, experiment with running the new index n along rows vs columns first.
These two suggestions combined should give you a massive speedup.
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.