Каков наиболее независимый от платформы и Python-версии способ сделать быстрый цикл для использования в Python?

Question

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

В настоящее время он написан на Python с чертой NumPy, и я включил приведенный ниже код.

1. Есть ли решение, которое было бы достаточно быстрым и не требовало бы компиляции?Казалось бы, это самый простой способ сохранить независимость от платформы.
2. Если вы используете что-то вроде Pyrex, которое требует компиляции, есть простой способ объединить множество модулей и сделать выбор между ними в зависимости от обнаруженной ОС.а Python версия?Есть ли простой способ создать коллекцию модулей без необходимости доступа к каждой системе с каждой версией Python?
3. Подходит ли один метод особенно для многопроцессорной оптимизации?

(Если вам интересно, цикл должен вычислять магнитное поле в заданной точке внутри кристалла, складывая вклады большого количества соседних магнитных ионов, рассматриваемых как крошечные стержневые магниты. В основном, огромная сумма эти .)

# calculate_dipole
# -------------------------
# calculate_dipole works out the dipole field at a given point within the crystal unit cell
# ---
# INPUT
# mu = position at which to calculate the dipole field
# r_i = array of atomic positions
# mom_i = corresponding array of magnetic moments
# ---
# OUTPUT
# B = the B-field at this point

def calculate_dipole(mu, r_i, mom_i):
    relative = mu - r_i
    r_unit = unit_vectors(relative)
    #4pi / mu0 (at the front of the dipole eqn)
    A = 1e-7
    #initalise dipole field
    B = zeros(3,float)

    for i in range(len(relative)):
        #work out the dipole field and add it to the estimate so far
        B += A*(3*dot(mom_i[i],r_unit[i])*r_unit[i] - mom_i[i]) / sqrt(dot(relative[i],relative[i]))**3
    return B

Answer 1

Вы можете заставить это работать намного, намного быстрее, если вы исключите цикл и будете использовать векторизованные операции Numpy.Поместите ваши данные в массивы фигуры (3, N) и попробуйте следующее:

import numpy as np

N = 20000
mu = np.random.random((3,1))
r_i = np.random.random((3,N))
mom_i = np.random.random((3,N))

def unit_vectors(r):
     return r / np.sqrt((r*r).sum(0))

def calculate_dipole(mu, r_i, mom_i):
    relative = mu - r_i
    r_unit = unit_vectors(relative)
    A = 1e-7

    num = A*(3*np.sum(mom_i*r_unit, 0)*r_unit - mom_i)
    den = np.sqrt(np.sum(relative*relative, 0))**3
    B = np.sum(num/den, 1)
    return B

Для меня это работает примерно в 50 раз быстрее, чем при использовании цикла for.

Answer 2

Numpy использует некоторую встроенную оптимизацию для обработки массива. Вы можете использовать массив Numpy с Cython для ускорения.

Answer 3

Ваш код на python, возможно, можно было бы немного ускорить, заменив цикл на выражение генератора и удалив все запросы на поиск mom_i [i], относительного [i] и r_unit [i], выполнив итерации по всем трем последовательностям параллельно, используя itertools.izip.

т.е. замена

B = zeros(3,float)

for i in range(len(relative)):
    #work out the dipole field and add it to the estimate so far
    B += A*(3*dot(mom_i[i],r_unit[i])*r_unit[i] - mom_i[i]) / sqrt(dot(relative[i],relative[i]))**3
return B

с:

from itertools import izip
...
return sum((A*(3*dot(mom,ru)*ru - mom) / sqrt(dot(rel,rel))**3 
            for mom, ru, rel in izip(mom_i, r_unit, relative)),
           zeros(3,float)) 

Это также более читабельно ИМХО, так как основное уравнение не загромождено [i] повсюду ..

Однако я подозреваю, что это принесет вам лишь незначительные выгоды по сравнению с выполнением всей функции на скомпилированном языке, таком как Cython.

Answer 4

Одно простое, но существенное ускорение - вывести умножение на А за пределы вашей суммы. Вы можете просто умножить B с ним, как только вернете его:

for i in range(len(relative)):
    #work out the dipole field and add it to the estimate so far
    B += (3*dot(mom_i[i],r_unit[i])*r_unit[i] - mom_i[i]) / sqrt(dot(relative[i],relative[i]))**3

return A*B

Это дало примерно 8% ускорение при использовании 20 000 случайных диполей.

Помимо этого простого ускорения, я бы порекомендовал использовать Cython (который обычно рекомендуется вместо Pyrex) или Weave от Scipy. Взгляните на Performance Python для некоторых примеров и сравнений различных способов ускорения Numpy / Scipy.

Если вы хотите попробовать эту параллель, я бы порекомендовал взглянуть на Параллельное программирование Сципи , чтобы начать.

Приятно видеть другого физика на SO. Здесь их не очень много.

Edit:

Я решил принять это как вызов для развития некоторых навыков Cython и получил улучшение примерно в 10 раз по сравнению с оптимизированной для Psyco версией. Дайте мне знать, если вы хотите увидеть мой код.

Edit2:

Ладно, вернулся и обнаружил, что тормозит в моей версии Cython. Теперь ускорение превышает 100 раз. Если вы хотите или нуждаетесь в другом коэффициенте в 2 раза превышающем ускоренную версию Ray Numpy, дайте мне знать, и я опубликую свой код.

Исходный код Cython:

Вот код Cython, который я набрал:

import numpy as np
cimport numpy as np
cimport cython
cdef extern from "math.h":
    double sqrt(double theta)
ctypedef np.float64_t dtype_t

@cython.boundscheck(False)
@cython.wraparound(False)
def calculate_dipole_cython(np.ndarray[dtype_t,ndim=2,mode="c"] mu, 
                            np.ndarray[dtype_t,ndim=2,mode="c"] r_i, 
                            np.ndarray[dtype_t,ndim=2,mode="c"] mom_i):
    cdef Py_ssize_t i
    cdef np.ndarray[dtype_t,ndim=1,mode="c"] tmp = np.empty(3,np.float64)
    cdef np.ndarray[dtype_t,ndim=1,mode="c"] relative = np.empty(3,np.float64)
    cdef double A = 1e-7
    cdef double C, D, F
    cdef np.ndarray[dtype_t,ndim=1,mode="c"] B = np.zeros(3,np.float64)
    for i in xrange(r_i.shape[0]):
        relative[0] = mu[0,0] - r_i[i,0]
        relative[1] = mu[0,1] - r_i[i,1]
        relative[2] = mu[0,2] - r_i[i,2]
        C = relative[0]*relative[0] + relative[1]*relative[1] + relative[2]*relative[2]
        C = 1.0/sqrt(C)
        D = C**3
        tmp[0] = relative[0]*C
        F = mom_i[i,0]*tmp[0]
        tmp[1] = relative[1]*C
        F += mom_i[i,1]*tmp[1]
        tmp[2] = relative[2]*C
        F += mom_i[i,2]*tmp[2]
        F *= 3
        B[0] += (F*tmp[0] - mom_i[i,0])*D
        B[1] += (F*tmp[1] - mom_i[i,1])*D
        B[2] += (F*tmp[2] - mom_i[i,2])*D
    return A*B

Я немного оптимизировал это, я думаю, но может быть немного больше, что вы можете извлечь из этого. Вы все еще можете заменить np.zeros и np.empty прямыми вызовами из Numpy C API, но это не должно иметь большого значения. В нынешнем виде этот код дает в 2-3 раза больше, чем у оптимизированного кода Numpy. Тем не менее, вам нужно правильно ввести числа. Массивы должны быть в формате C (который используется по умолчанию для массивов Numpy, но в Numpy транспонирование массива в формате C представляет собой массив в формате Fortran).

Например, чтобы запустить код из вашего другого вопроса , вам нужно заменить np.random.random((3,N)) s на np.random.random((N,3)). Кроме того, `

r_test_fast = reshape_vector(r_test) 

необходимо изменить на

r_test_fast = np.array(np.matrix(r_test))

Эта последняя строка может быть сделана проще / быстрее, но, по моему мнению, это было бы преждевременной оптимизацией.

Если вы раньше не использовали Cython и не знаете, как его скомпилировать, сообщите мне, и я буду рад помочь.

Наконец, я бы порекомендовал посмотреть эту статью . Я использовал это в качестве руководства для моей оптимизации. Следующим шагом будет попытка использовать функции BLAS, которые используют набор инструкций SSE2, попытка использования API SSE или попытка использовать больше API Numpy C, который взаимодействует с SSE2. Также вы можете посмотреть на распараллеливание.

Answer 5

Python не предназначен для высокопроизводительных вычислений. Напишите основной цикл в C и вызовите его из Python.