How can I access the neighboring elements of the matrix using numpy?

Question

I have made a code that calculates the dissolution of fluids, the problem is that the code is very poor, so I have been looking at that with numpy I can optimize it but I have been stuck without knowing how to do the following code using numpy and the roll function. Basically I have a matrix that the index and cannot be more than 1024, for this I use% to calculate what index it is. But this takes a long time.

I tried using numpy, using roll, rotating the matrix and then I don't have to calculate the module. But I don't know how to take the values of the neighbors.

def evolve(grid, dt, D=1.0):
  xmax, ymax = grid_shape
  new_grid = [[0.0,] * ymax for x in range(xmax)]
  for i in range(xmax):
    for j in range(ymax):
      grid_xx = grid[(i+1)%xmax][j] + grid[(i-1)%xmax][j] - 2.0 * grid[i][j]
      grid_yy = grid[i][(j+1)%ymax] + grid[i][(j-1)%ymax] - 2.0 * grid[i][j]
      new_grid[i][j] = grid[i][j] + D * (grid_xx + grid_yy) * dt
  return new_grid 

Answer 1

You have to rewrite the evolve function from (almost) zero using numpy .

Here the guidelines:

• First, grid must be a 2D numpy array, not a list of lists.
• Your teacher suggested the roll function: look at its docs and try to understand how it works. roll will solve the problem of finding neighbour entries in the matrix by shifting (or rolling) the matrix over one of the axis. You can then create shifted versions of grid in the four directions and use them, instead of searching for neighbours.
• Once you have the shifted grid s, you'll see that you will not need the for loops to calculate each cell of new_grid : you can use vectorized calculation, which is faster.

So the code will look like this:

def evolve(grid, dt, D=1.0):
    if not isinstance(grid, np.ndarray): #ensuring that is a numpy array.
        grid = np.array(grid)
    u_grid = np.roll(grid, 1, axis=0)
    d_grid = np.roll(grid, -1, axis=0)
    r_grid = np.roll(grid, 1, axis=1)
    l_grid = np.roll(grid, -1, axis=1)
    new_grid = grid + D * (u_grid + d_grid + r_grid + l_grid - 4.0*grid) * dt
    return new_grid

With a 1024 x 1024 matrix, each numpy evolve takes (on my machine) ~0.15 seconds to return the new_grid . Your evolve with the for loops takes ~3.85 seconds.