Proper numpy vectorization

Question

I am trying to vectorize my code. Am I converting

def acceleration_1(grid):

  nx = grid.shape[0]
  ny = grid.shape[1]
  acc = np.zeros((nx,ny,2))
  for i in range(1,nx-1):
    for j in range(1,ny-1):
      acc[i,j,0] = grid[i+1,j,0] + grid[i-1,j,0] - 2*grid[i,j,0]
      acc[i,j,1] = grid[i,j+1,1] + grid[i,j-1,1] - 2*grid[i,j,1]

to this

def acceleration_2(grid):

  nx = np.arange(1,grid.shape[0]-1)
  ny = np.arange(1,grid.shape[1]-1)
  acc = np.zeros((grid.shape[0],grid.shape[0],2))

  acc[nx,ny,0] = grid[nx+1,ny,0] + grid[nx-1,ny,0] - 2*grid[nx,ny,0]
  acc[nx,ny,1] = grid[nx,ny+1,1] + grid[nx,ny-1,1] - 2*grid[nx,ny,1]

properly?

I know I could represent this as matrix multiplication as well. But it just seems cumbersome to have to convert it to matrix operations. Do I get the best speedup by converting the for loops to an implicit iteration over "nx" and "ny"?

The first example doesn't work. Why not [i,j,0] and [i,j,1]? Your array shape is only 2D, as well. — Mateen Ulhaq
– Mateen Ulhaq, Commented Feb 19, 2018 at 17:40

Mateen Ulhaq · Accepted Answer · 2018-02-19 19:17:26Z

1

It seems like you're trying to do convolution. You could do this purely with numpy, similar to what you've shown. Alternatively, we can use built-in convolution functionality:

from scipy.signal import convolve2d

k_y = np.array([[1, -2, 1]]).T
k_x = np.array([[1, -2, 1]])

acc = np.zeros_like(grid)
acc[:, :, 0] = convolve2d(grid[:, :, 0], k_y, mode='same')
acc[:, :, 1] = convolve2d(grid[:, :, 1], k_x, mode='same')

To do it purely with numpy:

pad_y = np.pad(grid[:, :, 0], ((1, 1), (0, 0)), mode='constant')
pad_x = np.pad(grid[:, :, 1], ((0, 0), (1, 1)), mode='constant')

up    = pad_y[:-2,  1:-1]
down  = pad_y[2:,   1:-1]
left  = pad_x[1:-1, :-2]
right = pad_x[1:-1, 2:]

acc = np.zeros_like(grid)
acc[:, :, 0] = up   + down  - 2 * grid[:, :, 0]
acc[:, :, 1] = left + right - 2 * grid[:, :, 1]

edited Feb 19, 2018 at 19:17

answered Feb 19, 2018 at 17:53

Mateen Ulhaq

27.9k21 gold badges121 silver badges155 bronze badges

Sign up to request clarification or add additional context in comments.

7 Comments

user18764 Over a year ago

It seems that you cannot iterate over two separate arrays. So acc_2 would not work.

Mateen Ulhaq Over a year ago

What do you mean by acc_2?

user18764 Over a year ago

Sorry, the function acceleration_2 does not work properly. You cannot have arrays of different shapes as arguments to iterate over.

Mateen Ulhaq Over a year ago

Ah. You can make them the correct shape by padding (adding zeros to the boundaries).

user18764 Over a year ago

No, I meant in my function. Yours is fine. You just cant make 2 arrays, nx and ny, and assume that numpy will iterate over all pairs of (nx,ny).

|

Collectives™ on Stack Overflow

Proper numpy vectorization

1 Answer 1

7 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

7 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related