Оптимизация и ускорение математической функции в Python

Question

Цели данной математической функции для вычисления расстояния между двумя (или более) белковыми структурами с использованием двугранных углов:

Это очень полезно в структурно биологии, например. И я уже кодирую эту функцию в python, используя numpy, но цель состоит в том, чтобы иметь более быструю реализацию. В качестве ссылки на вычисление времени я использую функцию евклидовой дистанции, доступную в пакете scikit-learn.

Вот код у меня есть на данный момент:

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 мс это очень быстро, но это очень медленно, если вам нужно запустить его в миллион раз. Теперь вопрос в том, как это сделать? Каким будет следующий шаг? Cython? PyOpenCL? У меня есть опыт работы с PyOpenCL, но я никогда не кодирую что-то столь же сложное, как это. Я не знаю, можно ли вычислить двугранные расстояния за один шаг на GPU, как я делаю с numpy и как продолжить.

Спасибо, что помогли мне!

EDIT: Спасибо, ребята! В настоящее время я работаю над полным решением, и как только он будет закончен, я поставлю код здесь.

Cython версия:

%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 

Так я следую вашему link, чтобы создать версию Cython функции двугранного расстояния. Мы получаем некоторую скорость, не так много, но она все еще медленнее, чем версия numexpr (17 мс против 9,44 мс). Поэтому я попытался распараллелить функцию, используя prange, и это хуже (37.1ms против 17ms против 9.4ms)!

Пропустить что-нибудь?

Answer 1

Если вы готовы использовать http://pythran.readthedocs.io/, вы можете использовать на реализацию Numpy и получить более высокую производительность, чем Cython для этого случая:

#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]) 

И скомпилировать его с:

pythran fast.py 

Для получить среднее значение x2 над версией cython.

При использовании:

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

Вы получите Векторизованные, параллельную версию, которая работает немного быстрее:

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 

(используя тот же тестовый стенд как эй-BR)

Answer 2

Вот быстрый и грязный попробовать с Cython, для всего пары 1D массивов:

(в записной книжке IPython)

%%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) 

По сравнению с непосредственным осуществлением Numpy,

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

я

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 

Итак, в зависимости от размерности ваших векторов вы можете получить коэффициент 4 с нуля ускорения.

Для вычисления попарных расстояний вы, вероятно, можете сделать намного лучше, как показано в this blog post, но, конечно, YMMV.