Odległość między segmentami za pomocą numba jit, Python

Question

W zeszłym tygodniu zadawałem pytania pokrewne dotyczące tego stosu, próbując wyizolować rzeczy, których nie rozumiałem, używając dekoratora @jit z Numba w Pythonie. Uderzam jednak w ścianę, więc po prostu napiszę cały problem.

Problemem jest obliczenie minimalnej odległości między parami dużej liczby segmentów. Segmenty są reprezentowane przez ich początek i punkty końcowe w 3D. Matematycznie, każdy segment jest sparametryzowany jako [AB] = A + (B-A) * s, gdzie s w [0,1], i A i B są początkiem i punktami końcowymi segmentu. W przypadku dwóch takich segmentów można obliczyć minimalną odległość, a formułę podano here.

Już zdemaskowałem ten problem na innym thread, a otrzymana odpowiedź dotyczyła zastąpienia podwójnych pętli mojego kodu przez wektoryzację problemu, co jednak mogłoby spowodować problemy z pamięcią dla dużych zestawów segmentów. Dlatego postanowiłem trzymać się pętli i używać zamiast tego numby jit.

Ponieważ rozwiązanie tego problemu wymaga dużej ilości produktów dotowych, a produkt numpy's dotępny jest not supported by numba, zacząłem od wdrożenia mojego własnego produktu w postaci kropek 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. 

Funkcja obliczenie minimalnej odległości wszystkich par w segmentach N przyjmuje jako dane wejściowe tablicę NX6 współrzędnych (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) 

więc compute_stuff (Arg) wykonuje jako argument 2D np.array (double [:,:]), wykonuje kilka obsługiwanych przez numba (?) operacji i zwraca inne 2D np.array (double [:,:]). Jednak w przypadku każdej pętli otrzymuję 134 i 123 ms. Czy możesz rzucić trochę światła na to, dlaczego nie mogę przyspieszyć mojej funkcji? Wszelkie opinie będą mile widziane.

Answer 1

Oto moja wersja kodu, który jest znacznie szybsze:

@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 

A czasy:

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 

Moja podstawowa strategia dla przyspieszenia ten wynosił:

• Pozbądź wszystkich wyrażeń tablicowych i jawnie rozwinąć wektoryzację (zwłaszcza, że ​​twoje tablice są tak małe)
• Pozbądź się nadmiarowych obliczeń (odjęcie dwóch wektorów) w swoich wywołaniach do metody dot.
• Przenieś całe tworzenie tablicy na zewnątrz zagnieżdżonych pętli, aby numba mógł potencjalnie wykonać numer loop lifting. Pozwala to również uniknąć tworzenia wielu małych tablic, co jest kosztowne. Lepiej jest przydzielić raz i wykorzystać pamięć.

Inną rzeczą, aby pamiętać, że dla ostatnich wersjach Numba, po co nazywać autojit (tj wynajem Numba zrobić rodzaj wnioskowania na wejściach) i jest teraz tylko dekorator bez podpowiedzi typu jest zwykle tak szybko jak określanie typów wejść w moim doświadczeniu.

Również czasy zostały uruchomione przy użyciu numba 0.17.0 na OS X przy użyciu dystrybucji python Anaconda w Pythonie 2.7.9.