Distanza segmento inter utilizzando numba jit, Python

Question

Ho fatto domande correlate su questo stack per la settimana passata per cercare di isolare le cose che non ho capito sull'uso del decoratore @jit con Numba in Python. Comunque, sto sbattendo contro il muro, quindi scriverò solo l'intero problema.

Il problema è quello di calcolare la distanza minima tra le coppie di un numero di segmenti grande. I segmenti sono rappresentati dal loro inizio e dai loro endpoint in 3D. Matematicamente, ogni segmento è parametrizzato come [AB] = A + (B-A) * s, dove s in [0,1], e A e B sono l'inizio e il punto finale del segmento. Per due di questi segmenti, la distanza minima può essere calcolata e la formula è data here.

Ho già esposto questo problema su un altro thread e la risposta fornita riguardava la sostituzione di doppi loop del mio codice mediante la vettorizzazione del problema, che tuttavia avrebbe colpito problemi di memoria per insiemi di segmenti di grandi dimensioni. Pertanto, ho deciso di restare con i loop e utilizzare invece lo jolly di numba.

Poiché la soluzione del problema richiede molti prodotti dot e il prodotto con punti di Numpy è not supported by numba, ho iniziato implementando il mio prodotto 3D a punti.

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. 

La funzione di calcolo della distanza minima di tutte le coppie di segmenti N prende come input un array NX6 (6) coordinate

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) 

Quindi, compute_stuff (arg) prende come argomento un np.array 2D (double [:,:]), esegue una serie di operazioni supportate da numba (?) e restituisce un altro np.array 2D (double [:,:]). Tuttavia,

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

Ottengo 134 e 123 ms per ciclo. Puoi far luce sul motivo per cui non riesco ad accelerare la mia funzione? Qualsiasi feedback sarebbe molto apprezzato.

Answer 1

Ecco la mia versione del codice che è significativamente più veloce:

@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 

E i tempi:

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 

La mia strategia di base per accelerare questo up è stata:

• Sbarazzarsi di tutte le espressioni di matrice e srotolare esplicitamente la vettorizzazione (soprattutto dal momento che gli array sono così piccoli)
• Elimina i calcoli ridondanti (sottraendo due vettori) nelle chiamate al metodo dot.
• Sposta tutta la creazione di matrice all'esterno dei cicli nidificati in modo che numba possa potenzialmente eseguire loop lifting. Questo evita anche la creazione di molti piccoli array che è costoso. È meglio allocare una volta e riutilizzare la memoria.

Un'altra cosa da notare è che per le versioni recenti di numba, a che serve per essere chiamato autojit (cioè locazione numba fare inferenza di tipo sugli ingressi) ed è ora solo il decoratore, senza digitare i suggerimenti di solito è solo il più velocemente specificando i tipi di input nella mia esperienza.

Anche le temporizzazioni sono state eseguite utilizzando numba 0.17.0 su OS X utilizzando la distribuzione Python Anaconda con Python 2.7.9.