python - 优化成对减法

Question

我必须计算大约 1e6 个物体的粒子速度的成对差异，每个物体的速度约为。 1e4 粒子。现在我正在使用 itertools.combinations 来循环粒子，但仅对于一个对象我的代码就已经花费了 30 多分钟。我想知道我还能做些什么来将其加速到可行的速度，因为并行化似乎并没有在 python 中增加太多。 cython 是正确的选择吗？

这是我仅针对其中一个对象的代码:

def pairwisevel(hist,velj,velk, xj, xk):
    vlos = (velj - velk) 
    if( (xj - xk) < 0.):
        vlos = - vlos
    hist.add_from_value(vlos)


for i in itertools.combinations(np.arange(0,int(particles_per_group[0]),1),2):
    pairwisevel(hist, pvel[i[0]], pvel[i[1]],\
                pcoords[i[0]], pcoords[i[1]])

Answer 1

我希望我能理解你的问题。在此示例中，我计算了一个粒子对象的直方图。但如果您不想对所有 1e6 组 (1e4*1e4*1e6=1e14) 进行比较，这仍然需要几天时间。在此示例中，我使用 Numba完成任务。

代码

import numpy as np
import numba as nb
import time

#From Numba source
#Copyright (c) 2012, Anaconda, Inc.
#All rights reserved.

@nb.njit(fastmath=True)
def digitize(x, bins, right=False):
    # bins are monotonically-increasing
    n = len(bins)
    lo = 0
    hi = n

    if right:
        if np.isnan(x):
            # Find the first nan (i.e. the last from the end of bins,
            # since there shouldn't be many of them in practice)
            for i in range(n, 0, -1):
                if not np.isnan(bins[i - 1]):
                    return i
            return 0
        while hi > lo:
            mid = (lo + hi) >> 1
            if bins[mid] < x:
                # mid is too low => narrow to upper bins
                lo = mid + 1
            else:
                # mid is too high, or is a NaN => narrow to lower bins
                hi = mid
    else:
        if np.isnan(x):
            # NaNs end up in the last bin
            return n
        while hi > lo:
            mid = (lo + hi) >> 1
            if bins[mid] <= x:
                # mid is too low => narrow to upper bins
                lo = mid + 1
            else:
                # mid is too high, or is a NaN => narrow to lower bins
                hi = mid

    return lo

#Variant_1
@nb.njit(fastmath=True,parallel=True)
def bincount_comb_1(pvel,pcoords,bins):
  vlos_binned=np.zeros(bins.shape[0]+1,dtype=np.uint64)
  for i in nb.prange(pvel.shape[0]):
    for j in range(pvel.shape[0]):
      if( (pcoords[i] - pcoords[j]) < 0.):
        vlos = 0.
      else:
        vlos = (pvel[i] - pvel[j])

      dig_vlos=digitize(vlos, bins, right=False)
      vlos_binned[dig_vlos]+=1
  return vlos_binned

#Variant_2
#Is this also working?
@nb.njit(fastmath=True,parallel=True)
def bincount_comb_2(pvel,pcoords,bins):
  vlos_binned=np.zeros(bins.shape[0]+1,dtype=np.uint64)
  for i in nb.prange(pvel.shape[0]):
    for j in range(pvel.shape[0]):
      #only particles which fulfill this condition are counted?
      if( (pcoords[i] - pcoords[j]) < 0.):
        vlos = (pvel[i] - pvel[j])
        dig_vlos=digitize(vlos, bins, right=False)
        vlos_binned[dig_vlos]+=1
  return vlos_binned

#Variant_3
#Only counting once
@nb.njit(fastmath=True,parallel=True)
def bincount_comb_3(pvel,pcoords,bins):
  vlos_binned=np.zeros(bins.shape[0]+1,dtype=np.uint64)
  for i in nb.prange(pvel.shape[0]):
    for j in range(i,pvel.shape[0]):
      #only particles, where this condition is met are counted?
      if( (pcoords[i] - pcoords[j]) < 0.):
        vlos = (pvel[i] - pvel[j])
        dig_vlos=digitize(vlos, bins, right=False)
        vlos_binned[dig_vlos]+=1
  return vlos_binned


#Create some data to test
bins=np.arange(2,32)
pvel=np.random.rand(10_000)*35
pcoords=np.random.rand(10_000)*35

#first call has compilation overhead, we don't measure this
res_1=bincount_comb_1(pvel,pcoords,bins)
res_2=bincount_comb_2(pvel,pcoords,bins)

t1=time.time()
res=bincount_comb_1(pvel,pcoords,bins)
print(time.time()-t1)
t1=time.time()
res=bincount_comb_2(pvel,pcoords,bins)
print(time.time()-t1)
t1=time.time()
res=bincount_comb_3(pvel,pcoords,bins)
print(time.time()-t1)

性能

#Variant_1: 0.5s 5.78d for 1e6 groups of points
#Variant_2: 0.3s 3.25d for 1e6 groups of points
#Variant_3: 0.22s 2.54d for 1e6 groups of points