Любой способ ускорить этот код Python?

Question

Я написал код Python для выполнения некоторых работ по обработке изображений, но для запуска требуется огромное количество времени. Я потратил последние несколько часов, пытаясь его оптимизировать, но я думаю, что дошел до конца своих способностей.

Глядя на выходы из профилировщика, функция ниже занимает большую часть общего времени моего кода. Есть ли способ ускорить его?

def make_ellipse(x, x0, y, y0, theta, a, b): 
    c = np.cos(theta) 
    s = np.sin(theta) 
    a2 = a**2 
    b2 = b**2 
    xnew = x - x0 
    ynew = y - y0 
    ellipse = (xnew * c + ynew * s)**2/a2 + (xnew * s - ynew * c)**2/b2 <= 1 

    return ellipse 

Чтобы дать контекст, он вызывается с x и y в качестве выхода из np.meshgrid с достаточно большим размером сетки, а все остальные параметры, как простых целочисленных значений.

Хотя эта функция, по-видимому, занимает много времени, возможно, есть способы ускорить и остальную часть кода. Я оставил код в this gist.

Любые идеи были бы с благодарностью получены. Я пробовал использовать numba и autojit основных функций, но это мало помогает.

Answer 1

Давайте попробуем оптимизировать make_ellipse в сочетании с его вызывающим абонентом.

Во-первых, обратите внимание, что a и b одинаковы во многих вызовах. Так как make_ellipse каждый раз квадратизирует их, просто вместо этого вызывающий абонент делает это.

Во-вторых, обратите внимание, что np.cos(np.arctan(theta)) - 1/np.sqrt(1 + theta**2), который кажется немного быстрее в моей системе. Аналогичный трюк можно использовать для вычисления синуса, как из тета, так и из cos (тета) (или наоборот).

В-третьих, и менее конкретно, подумайте о коротком замыкании некоторых окончательных оценок формулы эллипса. Например, везде, где (xnew * c + ynew * s)**2/a2 больше 1, значение эллипса должно быть False. Если это случается часто, вы можете «замаскировать» вторую половину (дорогого) вычисления эллипса в этих местах. Я не планировал этого полностью, но смотрю numpy.ma на некоторые возможные выводы.

Answer 2

Вы можете значительно ускорить его, используя Cython. Существует очень хорошая документация о том, как это сделать.

Answer 3

Это не ускорит работу для всех случаев, но если ваши эллипсы не занимают все изображение, вы должны ограничить поиск точек внутри эллипса своим ограничивающим прямоугольником. Я ленив с математикой, так что я googled it и повторно @JohnZwinck аккуратный косинус арктангенса трюк, чтобы придумать с этой функцией:

def ellipse_bounding_box(x0, y0, theta, a, b): 
    x_tan_t = -b * np.tan(theta)/a 
    if np.isinf(x_tan_t) : 
     x_cos_t = 0 
     x_sin_t = np.sign(x_tan_t) 
    else : 
     x_cos_t = 1/np.sqrt(1 + x_tan_t*x_tan_t) 
     x_sin_t = x_tan_t * x_cos_t 
    x = x0 + a*x_cos_t*np.cos(theta) - b*x_sin_t*np.sin(theta) 

    y_tan_t = b/np.tan(theta)/a 
    if np.isinf(y_tan_t): 
     y_cos_t = 0 
     y_sin_t = np.sign(y_tan_t) 
    else: 
     y_cos_t = 1/np.sqrt(1 + y_tan_t*y_tan_t) 
     y_sin_t = y_tan_t * y_cos_t 
    y = y0 + b*y_sin_t*np.cos(theta) + a*y_cos_t*np.sin(theta) 

    return np.sort([-x, x]), np.sort([-y, y]) 

Теперь Вы можете изменить оригинальную функцию, чтобы что-то вроде этого:

def make_ellipse(x, x0, y, y0, theta, a, b): 
    c = np.cos(theta) 
    s = np.sin(theta) 
    a2 = a**2 
    b2 = b**2 
    x_box, y_box = ellipse_bounding_box(x0, y0, theta, a, b) 
    indices = ((x >= x_box[0]) & (x <= x_box[1]) & 
       (y >= y_box[0]) & (y <= y_box[1])) 
    xnew = x[indices] - x0 
    ynew = y[indices] - y0 
    ellipse = np.zeros_like(x, dtype=np.bool) 
    ellipse[indices] = ((xnew * c + ynew * s)**2/a2 + 
         (xnew * s - ynew * c)**2/b2 <= 1) 
    return ellipse 

Answer 4

Поскольку все, кроме x и y, являются целыми числами, вы можете попытаться свести к минимуму количество вычислений массива. Я полагаю, большую часть времени проводит в этом заявлении:

ellipse = (xnew * c + ynew * s)**2/a2 + (xnew * s - ynew * c)**2/b2 <= 1 

Простой рерайтинг как и должно сократить число операций с массивами:

a = float(a) 
b = float(b) 
ellipse = (xnew * (c/a) + ynew * (s/a))**2 + (xnew * (s/b) - ynew * (c/b))**2 <= 1 

Что было 12 операций массива в настоящее время 10 (плюс 4 скалярные операторы). Я не уверен, что джитт из numba попытался бы это сделать.Он может просто сделать все вещание первым, а затем выполнить соответствующие операции. В этом случае однократное переупорядочение таких общих операций должно помочь.

Содействуя вместе, вы можете переписать это снова, как

ellipse = ((xnew + ynew * (s/c)) * (c/a))**2 + ((xnew * (s/c) - ynew) * (c/b))**2 <= 1 

Или

t = numpy.tan(theta) 
ellipse = ((xnew + ynew * t) * (b/a))**2 + (xnew * t - ynew)**2 <= (b/c)**2 

Заменяя еще одну операцию массива с скаляр и устраняет другие скалярные опа, чтобы получить 9 операций массивов и 2 скаляр ОПС.

Как всегда, имейте в виду, что диапазон входов должен избегать ошибок округления.

К сожалению, нет способа сделать прогонную сумму и залог раньше, если любой из двух слагаемых больше, чем правая сторона сравнения. Это было бы очевидным ускорением, но вам понадобится cython (или c/C++) для кода.