Optimierung und Beschleunigung einer mathematischen Funktion in Python

Question

Der Zweck dieser mathematischen Funktion ist ein Abstand zwischen zwei (oder mehr) Proteinstrukturen mit Diederwinkeln zu berechnen:

Es ist sehr nützlich in Struktur Biologie zum Beispiel. Und ich schreibe diese Funktion bereits in Python mit numpy, aber das Ziel ist eine schnellere Implementierung. Als Berechnungszeitreferenz verwende ich die euklidische Abstandsfunktion, die im scikit-learn-Paket verfügbar ist.

Hier ist der Code, den ich für den Moment haben:

import numpy as np 
import numexpr as ne 
from sklearn.metrics.pairwise import euclidean_distances 

# We have 10000 structures with 100 dihedral angles 
n = 10000 
m = 100 

# Generate some random data 
c = np.random.rand(n,m) 
# Generate random int number 
x = np.random.randint(c.shape[0]) 

print c.shape, x 

# First version with numpy of the dihedral_distances function 
def dihedral_distances(a, b): 
    l = 1./a.shape[0] 
    return np.sqrt(l* np.sum((0.5)*(1. - np.cos(a-b)), axis=1)) 

# Accelerated version with numexpr 
def dihedral_distances_ne(a, b): 
    l = 1./a.shape[0] 
    tmp = ne.evaluate('sum((0.5)*(1. - cos(a-b)), axis=1)') 
    return ne.evaluate('sqrt(l* tmp)') 

# The function of reference I try to be close as possible 
# in term of computation time 
%timeit euclidean_distances(c[x,:], c)[0] 
1000 loops, best of 3: 1.07 ms per loop 

# Computation time of the first version of the dihedral_distances function 
# We choose randomly 1 structure among the 10000 structures. 
# And we compute the dihedral distance between this one and the others 
%timeit dihedral_distances(c[x,:], c) 
10 loops, best of 3: 21.5 ms per loop 

# Computation time of the accelerated function with numexpr 
%timeit dihedral_distances_ne(c[x,:], c) 
100 loops, best of 3: 9.44 ms per loop 

9,44 ms es sehr schnell ist, aber es ist sehr langsam, wenn Sie es eine Million Mal ausgeführt werden müssen. Jetzt ist die Frage, wie man das macht? Was ist der nächste Schritt? Cython? PyOpenCL? Ich habe einige Erfahrung mit PyOpenCL, aber ich kodiere nie etwas so ausführlich wie dieses. Ich weiß nicht, ob es möglich ist, auf einem GPU die Zweiflächendistanzen in einem Schritt zu berechnen, wie ich es mit numpy mache und wie es weitergeht.

Vielen Dank für Ihre Hilfe!

EDIT: Danke Jungs! Ich arbeite gerade an der vollständigen Lösung und sobald es fertig ist, werde ich den Code hier einfügen.

Cython VERSION:

%load_ext cython 
import numpy as np 

np.random.seed(1234) 

n = 10000 
m = 100 

c = np.random.rand(n,m) 
x = np.random.randint(c.shape[0]) 

print c.shape, x 

%%cython --compile-args=-fopenmp --link-args=-fopenmp --force 

import numpy as np 
cimport numpy as np 
from libc.math cimport sqrt, cos 
cimport cython 
from cython.parallel cimport parallel, prange 

# Define a function pointer to a metric 
ctypedef double (*metric)(double[: ,::1], np.intp_t, np.intp_t) 

cdef extern from "math.h" nogil: 
    double cos(double x) 
    double sqrt(double x) 

@cython.boundscheck(False) 
@cython.wraparound(False) 
@cython.cdivision(True) 
cdef double dihedral_distances(double[:, ::1] a, np.intp_t i1, np.intp_t i2): 
    cdef double res 
    cdef int m 
    cdef int j 

    res = 0. 
    m = a.shape[1] 

    for j in range(m): 
     res += 1. - cos(a[i1, j] - a[i2, j]) 

    res /= 2.*m 

    return sqrt(res) 

@cython.boundscheck(False) 
@cython.wraparound(False) 
@cython.cdivision(True) 
cdef double dihedral_distances_p(double[:, ::1] a, np.intp_t i1, np.intp_t i2): 
    cdef double res 
    cdef int m 
    cdef int j 

    res = 0. 
    m = a.shape[1] 

    with nogil, parallel(num_threads=2): 
     for j in prange(m, schedule='dynamic'): 
      res += 1. - cos(a[i1, j] - a[i2, j]) 

    res /= 2.*m 

    return sqrt(res) 

@cython.boundscheck(False) 
@cython.wraparound(False) 
def pairwise(double[: ,::1] c not None, np.intp_t x, p = True): 
    cdef metric dist_func 
    if p: 
     dist_func = &dihedral_distances_p 
    else: 
     dist_func = &dihedral_distances 

    cdef np.intp_t i, n_structures 
    n_samples = c.shape[0] 

    cdef double[::1] res = np.empty(n_samples) 

    for i in range(n_samples): 
     res[i] = dist_func(c, x, i) 

    return res 

%timeit pairwise(c, x, False) 
100 loops, best of 3: 17 ms per loop  

# Parallel version 
%timeit pairwise(c, x, True) 
10 loops, best of 3: 37.1 ms per loop 

So folge ich Ihrer link die cython Version der V-Form Entfernungen Funktion zu erstellen. Wir gewinnen etwas an Geschwindigkeit, nicht so sehr, aber es ist immer noch langsamer als die NUMEXPR-Version (17ms vs 9,44ms). Also habe ich versucht, die Funktion mit prange zu parallelisieren und es ist schlimmer (37.1ms vs 17ms vs 9.4ms)!

Vermisse ich etwas?

Answer 1

Wenn Sie bereit sind http://pythran.readthedocs.io/ zu verwenden, können Sie auf der numpy Implementierung nutzen und eine bessere Leistung als cython für diesen Fall erhalten:

#pythran export np_cos_norm(float[], float[]) 
import numpy as np 
def np_cos_norm(a, b): 
    val = np.sum(1. - np.cos(a-b)) 
    return np.sqrt(val/2./a.shape[0]) 

und kompilieren Sie es mit:

pythran fast.py 

zu bekomme einen Durchschnitt von x2 über die Cython-Version.

Bei Verwendung:

pythran fast.py -march=native -DUSE_BOOST_SIMD -fopenmp 

Sie erhalten eine vektorisiert, parallele Version erhalten, die etwas schneller läuft:

100000 loops, best of 3: 2.54 µs per loop 
1000000 loops, best of 3: 674 ns per loop 

100000 loops, best of 3: 16.9 µs per loop 
100000 loops, best of 3: 4.31 µs per loop 

10000 loops, best of 3: 176 µs per loop 
10000 loops, best of 3: 42.9 µs per loop 

(die gleiche Testbed wie ev-br verwenden)

Answer 2

Hier ist ein schneller und unsauberer Versuch mit cython, für nur ein Paar von 1D-Arrays:

(in einem IPython Notebook)

%%cython 

cimport cython 
cimport numpy as np 

cdef extern from "math.h": 
    double cos(double x) nogil 
    double sqrt(double x) nogil 

def cos_norm(a, b): 
    return cos_norm_impl(a, b) 

@cython.boundscheck(False) 
@cython.wraparound(False) 
@cython.cdivision(True) 
cdef double cos_norm_impl(double[::1] a, double[::1] b) nogil: 
    cdef double res = 0., val 
    cdef int m = a.shape[0] 
    # XXX: shape of b not checked 
    cdef int j 

    for j in range(m): 
     val = a[j] - b[j] 
     res += 1. - cos(val) 
    res /= 2.*m 

    return sqrt(res) 

mit einer einfachen numpy Implementierung Vergleich,

def np_cos_norm(a, b): 
    val = np.add.reduce(1. - np.cos(a-b)) 
    return np.sqrt(val/2./a.shape[0]) 

ich

np.random.seed(1234) 

for n in [100, 1000, 10000]: 
    x = np.random.random(n) 
    y = np.random.random(n) 
    %timeit cos_norm(x, y) 
    %timeit np_cos_norm(x, y) 
    print '\n' 

100000 loops, best of 3: 3.04 µs per loop 
100000 loops, best of 3: 12.4 µs per loop 

100000 loops, best of 3: 18.8 µs per loop 
10000 loops, best of 3: 30.8 µs per loop 

1000 loops, best of 3: 196 µs per loop 
1000 loops, best of 3: 223 µs per loop 

Also, abhängig von der Dimensionalität Ihrer Vektoren, können Sie von einem Faktor von 4 bis Null einer Beschleunigung erhalten.

Für die Berechnung von paarweisen Abständen können Sie wahrscheinlich viel besser tun, wie in this blog post gezeigt, aber natürlich YMMV.