python - 使用 Parakeet 优化 Python 函数

Question

我需要优化此函数，因为我正试图让我的 OpenGL 模拟运行得更快。我想用 Parakeet ，但我不太明白我需要以何种方式修改下面的代码才能这样做。你能看到我应该做什么吗？

def distanceMatrix(self,x,y,z):
    " ""Computes distances between all particles and places the result in a matrix such that the ij th matrix entry corresponds to the distance between particle i and j"" "
    xtemp = tile(x,(self.N,1))
    dx = xtemp - xtemp.T
    ytemp = tile(y,(self.N,1))
    dy = ytemp - ytemp.T
    ztemp = tile(z,(self.N,1))
    dz = ztemp - ztemp.T

    # Particles 'feel' each other across the periodic boundaries
    if self.periodicX:
        dx[dx>self.L/2]=dx[dx > self.L/2]-self.L
        dx[dx<-self.L/2]=dx[dx < -self.L/2]+self.L
    if self.periodicY:
        dy[dy>self.L/2]=dy[dy>self.L/2]-self.L
        dy[dy<-self.L/2]=dy[dy<-self.L/2]+self.L
    if self.periodicZ:
        dz[dz>self.L/2]=dz[dz>self.L/2]-self.L
        dz[dz<-self.L/2]=dz[dz<-self.L/2]+self.L

    # Total Distances
    d = sqrt(dx**2+dy**2+dz**2)

    # Mark zero entries with negative 1 to avoid divergences
    d[d==0] = -1

    return d, dx, dy, dz

据我所知，Parakeet 应该能够在不进行修改的情况下使用上述函数——它只使用 Numpy 和数学。但是，从 Parakeet jit 包装器调用函数时，我总是会收到以下错误:

AssertionError: Unsupported function: <bound method Particles.distanceMatrix of <particles.Particles instance at 0x04CD8E90>>

Answer 1

Parakeet 还很年轻，它对 NumPy 的支持还不完整，而且您的代码涉及到几个尚不可用的功能。

1) 您正在包装一个方法，而 Parakeet 到目前为止只知道如何处理函数。常见的解决方法是制作一个 @jit 包装的辅助函数，并让您的方法使用所有必需的成员数据调用它。方法不起作用的原因是将有意义的类型分配给“self”并非易事。这并非不可能，但足够棘手，以至于在摘下悬垂的果实之前，方法不会进入长尾小鹦鹉。说到唾手可得的果实...

2) bool 索引。尚未实现，但将在下一个版本中实现。

3) np.tile:也不起作用，也可能会出现在下一个版本中。如果您想查看哪些内置函数和 NumPy 库函数可以工作，请查看 Parakeet 的 mappings。模块。

我重写了您的代码，以便对 Parakeet 更友好一些:

@jit 
def parakeet_dist(x, y, z, L, periodicX, periodicY, periodicZ):
  # perform all-pairs computations more explicitly 
  # instead of tile + broadcasting
  def periodic_diff(x1, x2, periodic):
    diff = x1 - x2 
    if periodic:
      if diff > (L / 2): diff -= L
      if diff < (-L/2): diff += L
    return diff
  dx = np.array([[periodic_diff(x1, x2, periodicX) for x1 in x] for x2 in x])
  dy = np.array([[periodic_diff(y1, y2, periodicY) for y1 in y] for y2 in y])
  dz = np.array([[periodic_diff(z1, z2, periodicZ) for z1 in z] for z2 in z])
  d= np.sqrt(dx**2 + dy**2 + dz**2)

  # since we can't yet use boolean indexing for masking out zero distances
  # have to fall back on explicit loops instead 
  for i in xrange(len(x)):
    for j in xrange(len(x)):
      if d[i,j] == 0: d[i,j] = -1 
  return d, dx, dy, dz 

在我的机器上，当 N = 2000 时，它的运行速度仅比 NumPy 快约 3 倍(NumPy 为 0.39 秒，Parakeet 为 0.14 秒)。如果我重写数组遍历以更明确地使用循环，那么性能将比 NumPy 快约 6 倍(Parakeet 在约 0.06 秒内运行):

@jit 
def loopy_dist(x, y, z, L, periodicX, periodicY, periodicZ):
  N = len(x)
  dx = np.zeros((N,N))
  dy = np.zeros( (N,N) )
  dz = np.zeros( (N,N) )
  d = np.zeros( (N,N) )

  def periodic_diff(x1, x2, periodic):
    diff = x1 - x2 
    if periodic:
      if diff > (L / 2): diff -= L
      if diff < (-L/2): diff += L
    return diff

  for i in xrange(N):
    for j in xrange(N):
      dx[i,j] = periodic_diff(x[j], x[i], periodicX)
      dy[i,j] = periodic_diff(y[j], y[i], periodicY)
      dz[i,j] = periodic_diff(z[j], z[i], periodicZ)
      d[i,j] = dx[i,j] ** 2 + dy[i,j] ** 2 + dz[i,j] ** 2 
      if d[i,j] == 0: d[i,j] = -1
      else: d[i,j] = np.sqrt(d[i,j])
  return d, dx, dy, dz 

通过一些创造性的重写，您还可以让上述代码在 Numba 中运行，但它只比 NumPy 快 1.5 倍(0.25 秒)。编译时间是 Parakeet w/理解:1 秒，Parakeet w/循环:0.5 秒，Numba w/循环:0.9 秒。

希望接下来的几个版本能够更加地道地使用 NumPy 库函数，但目前理解或循环通常是可行的方法。