正确的numpy矢量化 [英] Proper numpy vectorization

问题描述

我正在尝试向量化我的代码.我要转换吗

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]

对此

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]

正确吗?

我知道我也可以将其表示为矩阵乘法.但是将其转换为矩阵运算似乎很麻烦.通过将for循环转换为"nx"和"ny"上的隐式迭代，是否可以获得最佳的加速效果?

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"?

推荐答案

似乎您正在尝试进行卷积.您可以使用numpy完全做到这一点，类似于您所显示的.另外，我们可以使用内置的卷积功能:

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

---

仅使用numpy即可做到:

---

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]

这篇关于正确的numpy矢量化的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！