用P优化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

据我所知，长尾鹦鹉应该可以使用上面的函数而不做任何修改——它只使用Numpy和math。但是，当我从Paraket jit包装器调用函数时，总是出现以下错误：

^{pr2}$

Answer 1

鹦鹉还很年轻，它对NumPy的支持还不完整，并且您的代码涉及到一些还不起作用的特性。在

1）你在包装一个方法，而鹦鹉到目前为止只知道如何处理函数。常见的解决方法是生成一个@jit包装的helper函数，并使用所有必需的成员数据调用该函数。方法不起作用的原因是将有意义的类型赋给“self”是非常重要的。这不是不可能的，但足够棘手的是，在摘下悬挂较低的水果之前，这些方法是不会进入长尾鹦鹉的。说到低垂的水果。。。在

2）布尔索引。尚未实现，但将在下一个版本中提供。在

3）np.瓷砖：同样不起作用，也可能在下一个版本中。如果您想知道哪些内置函数和NumPy库函数可以工作，请看一下鹦鹉的mappings模块。在

我重写了你的代码，让它对鹦鹉更友好一点：

@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.39s，而Parakeet为0.14s）。如果我重写数组遍历以更明确地使用循环，那么性能将比NumPy快约6倍（鹦鹉在~0.06s内运行）：

^{pr2}$

通过一些创造性的重写，您也可以让上面的代码在Numba中运行，但它只比NumPy快1.5倍（0.25秒）。编译时间是带理解的鹦鹉：1秒，带循环的鹦鹉：.5秒，带循环的Numba:0.9秒。在

希望接下来的几个发行版将支持对NumPy库函数的更多惯用用法，但目前理解或循环通常是最好的方法。在