使用numba jit,Python的段间距离

Question

在过去的一周里,我一直在问这个堆栈上的相关问题,试图找出我不理解在Python中使用Numba的@jit装饰器的事情.但是,我正在撞墙,所以我只会写出整个问题.

手头的问题是计算大量段之间的最小距离.这些段由它们在3D中的起点和终点表示.在数学上,每个段被参数化为[AB] = A +(BA)*s,其中s在[0,1]中,A和B是段的起点和终点.对于两个这样的段,可以计算最小距离并且在此给出公式.

我已经在另一个线程上暴露了这个问题,并且给出的答案是通过向量化问题替换我的代码的双循环,然而这会对大量段的内存问题造成影响.因此,我决定坚持使用循环,而是使用numba的jit代替.

由于问题的解决方案需要大量的点积,并且numba不支持 numpy的点积,所以我开始实现自己的3D点积.

import numpy as np
from numba import jit, autojit, double, float64, float32, void, int32

def my_dot(a,b):
    res = a[0]*b[0]+a[1]*b[1]+a[2]*b[2]
    return res

dot_jit = jit(double(double[:], double[:]))(my_dot)    #I know, it's not of much use here.

计算N个段中所有对的最小距离的函数将Nx6数组作为输入(6个坐标)

def compute_stuff(array_to_compute):
    N = len(array_to_compute)
    con_mat = np.zeros((N,N))
    for i in range(N):
        for j in range(i+1,N):

            p0 = array_to_compute[i,0:3]
            p1 = array_to_compute[i,3:6]
            q0 = array_to_compute[j,0:3]
            q1 = array_to_compute[j,3:6]

            s = ( dot_jit((p1-p0),(q1-q0))*dot_jit((q1-q0),(p0-q0)) - dot_jit((q1-q0),(q1-q0))*dot_jit((p1-p0),(p0-q0)))/( dot_jit((p1-p0),(p1-p0))*dot_jit((q1-q0),(q1-q0)) - dot_jit((p1-p0),(q1-q0))**2 )
            t = ( dot_jit((p1-p0),(p1-p0))*dot_jit((q1-q0),(p0-q0)) -dot_jit((p1-p0),(q1-q0))*dot_jit((p1-p0),(p0-q0)))/( dot_jit((p1-p0),(p1-p0))*dot_jit((q1-q0),(q1-q0)) - dot_jit((p1-p0),(q1-q0))**2 )

            con_mat[i,j] = np.sum( (p0+(p1-p0)*s-(q0+(q1-q0)*t))**2 ) 

return con_mat

fast_compute_stuff = jit(double[:,:](double[:,:]))(compute_stuff)

因此,compute_stuff(arg)将2D np.array(double [:,]]作为参数,执行一堆numba支持的(？)操作,并返回另一个2D np.array(double [:,:]) .然而,

v = np.random.random( (100,6) )
%timeit compute_stuff(v)
%timeit fast_compute_stuff(v)

我得到每循环134和123毫秒.你能说清楚为什么我没能加快我的功能吗？任何反馈都将非常感激.

Answer 1

这是我的代码版本,速度明显更快:

@jit(nopython=True)
def dot(a,b):
    res = a[0]*b[0]+a[1]*b[1]+a[2]*b[2]
    return res

@jit
def compute_stuff2(array_to_compute):
    N = array_to_compute.shape[0]
    con_mat = np.zeros((N,N))

    p0 = np.zeros(3)
    p1 = np.zeros(3)
    q0 = np.zeros(3)
    q1 = np.zeros(3)

    p0m1 = np.zeros(3)
    p1m0 = np.zeros(3)
    q0m1 = np.zeros(3)
    q1m0 = np.zeros(3)
    p0mq0 = np.zeros(3)

    for i in range(N):
        for j in range(i+1,N):

            for k in xrange(3):
                p0[k] = array_to_compute[i,k]
                p1[k] = array_to_compute[i,k+3]
                q0[k] = array_to_compute[j,k]
                q1[k] = array_to_compute[j,k+3]

                p0m1[k] = p0[k] - p1[k]
                p1m0[k] = -p0m1[k]

                q0m1[k] = q0[k] - q1[k]
                q1m0[k] = -q0m1[k]

                p0mq0[k] = p0[k] - q0[k]

            s = ( dot(p1m0, q1m0)*dot(q1m0, p0mq0) - dot(q1m0, q1m0)*dot(p1m0, p0mq0))/( dot(p1m0, p1m0)*dot(q1m0, q1m0) - dot(p1m0, q1m0)**2 )
            t = ( dot(p1m0, p1m0)*dot(q1m0, p0mq0) - dot(p1m0, q1m0)*dot(p1m0, p0mq0))/( dot(p1m0, p1m0)*dot(q1m0, q1m0) - dot(p1m0, q1m0)**2 )


            for k in xrange(3):
                con_mat[i,j] += (p0[k]+(p1[k]-p0[k])*s-(q0[k]+(q1[k]-q0[k])*t))**2 

    return con_mat

时间安排:

In [38]:

v = np.random.random( (100,6) )
%timeit compute_stuff(v)
%timeit fast_compute_stuff(v)
%timeit compute_stuff2(v)

np.allclose(compute_stuff2(v), compute_stuff(v))

#10 loops, best of 3: 107 ms per loop
#10 loops, best of 3: 108 ms per loop
#10000 loops, best of 3: 114 µs per loop
#True

我加快这项工作的基本策略是:

• 摆脱所有数组表达式并显式展开矢量化(特别是因为你的数组很小)
• 在对dot方法的调用中去掉冗余计算(减去两个向量).
• 将所有数组创建移动到嵌套for循环之外,以便numba可能会执行一些循环提升.这也避免了创建许多昂贵的小阵列.最好分配一次并重用内存.

另外需要注意的是,对于numba的最新版本,有什么用途要调用autojit(即让numba对输入进行类型推断),现在只是没有类型提示的装饰器通常和我的经验中指定输入类型一样快.

另外,使用带有Python 2.7.9的Anaconda python发行版在OS X上使用numba 0.17.0运行时序.