使用numba jit，Python的段间距离

Question

在过去的一周里，我一直在询问有关此堆栈的相关问题，以尝试找出我不了解的有关在 Python 中将 @jit 装饰器与 Numba 一起使用的问题。但是，我碰壁了，所以我会写整个问题。

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

我已经在另一个thread 上暴露了这个问题，并且给出的答案是通过向量化问题来替换我的代码的双循环，但这会导致大量段的内存问题。因此，我决定坚持使用循环，并改用 numba 的 jit。

由于解决问题需要很多点积，而numpy的点积是not supported by numba，所以我从实现自己的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 可能会做一些loop lifting。这也避免了创建许多成本高昂的小阵列。最好分配一次，重复使用内存。

要注意的另一件事是，对于最近版本的 numba，过去称为 autojit（即让 numba 对输入进行类型推断）现在只是没有类型提示的装饰器通常与指定一样快根据我的经验，您的输入类型。

还使用带有 Python 2.7.9 的 Anaconda python 发行版在 OS X 上使用 numba 0.17.0 运行计时。