L oop ускорение БПФ в python (с `np.einsum`)

Question

Проблема: Я хочу ускорить мой python l oop, содержащий множество продуктов и сумм с np.einsum, но я также открыт для любых других решений.

Моя функция принимает векторную конфигурацию S формы (n, n, 3) (в моем случае: n = 72) и выполняет Фурье-преобразование для корреляционной функции для N * N точек. Корреляционная функция определяется как произведение каждого вектора на каждый другой. Это умножается на функцию косинуса от положений векторов, умноженных на значения kx и ky. Каждая позиция i,j в итоге суммируется, чтобы получить одну точку в k-пространстве p,m:

def spin_spin(S,N):
    n= len(S)
    conf = np.reshape(S,(n**2,3))
    chi = np.zeros((N,N))
    kx = np.linspace(-5*np.pi/3,5*np.pi/3,N)
    ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N)

    x=np.reshape(triangular(n)[0],(n**2))
    y=np.reshape(triangular(n)[1],(n**2))
    for p in range(N):
        for m in range(N):
            for i in range(n**2):
                for j in range(n**2):        
                    chi[p,m] += 2/(n**2)*np.dot(conf[i],conf[j])*np.cos(kx[p]*(x[i]-x[j])+ ky[m]*(y[i]-y[j]))
    return(chi,kx,ky)

Моя проблема в том, что мне нужно примерно 100 * 100 точек, которые обозначаются через kx * ky и l oop нужно много часов, чтобы завершить эту работу для решетки с 72 * 72 векторами. Количество вычислений: 72 * 72 * 72 * 72 * 100 * 100 Я не могу использовать встроенное БПФ numpy из-за моей сетки tri angular, поэтому мне нужен какой-то другой вариант, чтобы уменьшить здесь вычислительные затраты.

Моя идея: Сначала я понял, что преобразование конфигурации в список векторов вместо матрицы снижает вычислительные затраты. Кроме того, я использовал пакет numba, который также снизил стоимость, но все еще слишком медленный. Я обнаружил, что хорошим способом вычисления таких объектов является функция np.einsum. Вычисление произведения каждого вектора с каждым вектором выполняется следующим образом:

np.einsum('ij,kj -> ik',np.reshape(S,(72**2,3)),np.reshape(S,(72**2,3)))

Сложной частью является вычисление члена внутри np.cos. Здесь я хочу рассчитать произведение между списком фигур (100,1) и позициями векторов (например, np.shape(x)=(72**2,1)). Особенно я действительно не знаю, как реализовать расстояние по оси X и Y с помощью np.einsum.

Чтобы воспроизвести код (возможно, вам это не понадобится): Сначала вам потребуется векторная конфигурация. Вы можете сделать это просто с помощью np.ones((72,72,3) или взять случайные векторы в качестве примера с:

def spherical_to_cartesian(r, theta, phi):
    '''Convert spherical coordinates (physics convention) to cartesian coordinates'''
    sin_theta = np.sin(theta)
    x = r * sin_theta * np.cos(phi)
    y = r * sin_theta * np.sin(phi)
    z = r * np.cos(theta)

    return x, y, z # return a tuple

def random_directions(n, r):
    '''Return ``n`` 3-vectors in random directions with radius ``r``'''
    out = np.empty(shape=(n,3), dtype=np.float64)

    for i in range(n):
        # Pick directions randomly in solid angle
        phi = random.uniform(0, 2*np.pi)
        theta = np.arccos(random.uniform(-1, 1))
        # unpack a tuple
        x, y, z = spherical_to_cartesian(r, theta, phi)
        out[i] = x, y, z

    return out
S = np.reshape(random_directions(72**2,1),(72,72,3))

(Изменение формы в этом примере необходимо для придания ей формы в функции spin_spin обратно к (72 *) * 2,3) shape.)

Для позиций векторов я использую сетку tri angular, определяемую

def triangular(nsize):
    '''Positional arguments of the spin configuration'''

    X=np.zeros((nsize,nsize))
    Y=np.zeros((nsize,nsize))
    for i in range(nsize):
        for j in range(nsize):
            X[i,j]+=1/2*j+i
            Y[i,j]+=np.sqrt(3)/2*j
    return(X,Y)

Answer 1

Оптимизированная реализация Numba

Основная проблема в вашем коде - многократный вызов внешней функции BLAS np.dot с чрезвычайно маленькими данными. В этом коде имеет смысл вычислять их только один раз, но если вам нужно выполнить эти вычисления в al oop, напишите реализацию Numba. Пример

Оптимизированная функция (перебор)

import numpy as np
import numba as nb

@nb.njit(fastmath=True,error_model="numpy",parallel=True)
def spin_spin(S,N):
    n= len(S)
    conf = np.reshape(S,(n**2,3))
    chi = np.zeros((N,N))
    kx = np.linspace(-5*np.pi/3,5*np.pi/3,N).astype(np.float32)
    ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N).astype(np.float32)

    x=np.reshape(triangular(n)[0],(n**2)).astype(np.float32)
    y=np.reshape(triangular(n)[1],(n**2)).astype(np.float32)

    #precalc some values
    fact=nb.float32(2/(n**2))
    conf_dot=np.dot(conf,conf.T).astype(np.float32)

    for p in nb.prange(N):
        for m in range(N):
            #accumulating on a scalar is often beneficial
            acc=nb.float32(0)
            for i in range(n**2):
                for j in range(n**2):        
                    acc+= conf_dot[i,j]*np.cos(kx[p]*(x[i]-x[j])+ ky[m]*(y[i]-y[j]))
            chi[p,m]=fact*acc

    return(chi,kx,ky)

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

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

@nb.njit()
def precalc(S):
    #There may not be all redundancies removed
    n= len(S)
    conf = np.reshape(S,(n**2,3))
    conf_dot=np.dot(conf,conf.T)
    x=np.reshape(triangular(n)[0],(n**2))
    y=np.reshape(triangular(n)[1],(n**2))

    x_s=set()
    y_s=set()
    for i in range(n**2):
        for j in range(n**2):
            x_s.add((x[i]-x[j]))
            y_s.add((y[i]-y[j]))

    x_arr=np.sort(np.array(list(x_s)))
    y_arr=np.sort(np.array(list(y_s)))


    conf_dot_sel=np.zeros((x_arr.shape[0],y_arr.shape[0]))
    for i in range(n**2):
        for j in range(n**2):
            ii=np.searchsorted(x_arr,x[i]-x[j])
            jj=np.searchsorted(y_arr,y[i]-y[j])
            conf_dot_sel[ii,jj]+=conf_dot[i,j]

    return x_arr,y_arr,conf_dot_sel

@nb.njit(fastmath=True,error_model="numpy",parallel=True)
def spin_spin_opt_2(S,N):
    chi = np.empty((N,N))
    n= len(S)

    kx = np.linspace(-5*np.pi/3,5*np.pi/3,N)
    ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N)

    x_arr,y_arr,conf_dot_sel=precalc(S)
    fact=2/(n**2)
    for p in nb.prange(N):
        for m in range(N):
            acc=nb.float32(0)
            for i in range(x_arr.shape[0]):
                for j in range(y_arr.shape[0]):        
                    acc+= fact*conf_dot_sel[i,j]*np.cos(kx[p]*x_arr[i]+ ky[m]*y_arr[j])
            chi[p,m]=acc

    return(chi,kx,ky)

@nb.njit()
def precalc(S):
    #There may not be all redundancies removed
    n= len(S)
    conf = np.reshape(S,(n**2,3))
    conf_dot=np.dot(conf,conf.T)
    x=np.reshape(triangular(n)[0],(n**2))
    y=np.reshape(triangular(n)[1],(n**2))

    x_s=set()
    y_s=set()
    for i in range(n**2):
        for j in range(n**2):
            x_s.add((x[i]-x[j]))
            y_s.add((y[i]-y[j]))

    x_arr=np.sort(np.array(list(x_s)))
    y_arr=np.sort(np.array(list(y_s)))


    conf_dot_sel=np.zeros((x_arr.shape[0],y_arr.shape[0]))
    for i in range(n**2):
        for j in range(n**2):
            ii=np.searchsorted(x_arr,x[i]-x[j])
            jj=np.searchsorted(y_arr,y[i]-y[j])
            conf_dot_sel[ii,jj]+=conf_dot[i,j]

    return x_arr,y_arr,conf_dot_sel

@nb.njit(fastmath=True,error_model="numpy",parallel=True)
def spin_spin_opt_2(S,N):
    chi = np.empty((N,N))
    n= len(S)

    kx = np.linspace(-5*np.pi/3,5*np.pi/3,N)
    ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N)

    x_arr,y_arr,conf_dot_sel=precalc(S)
    fact=2/(n**2)
    for p in nb.prange(N):
        for m in range(N):
            acc=nb.float32(0)
            for i in range(x_arr.shape[0]):
                for j in range(y_arr.shape[0]):        
                    acc+= fact*conf_dot_sel[i,j]*np.cos(kx[p]*x_arr[i]+ ky[m]*y_arr[j])
            chi[p,m]=acc

    return(chi,kx,ky)

Время

#brute-force
%timeit res=spin_spin(S,100)
#48 s ± 671 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

#new version
%timeit res_2=spin_spin_opt_2(S,100)
#5.33 s ± 59.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit res_2=spin_spin_opt_2(S,1000)
#1min 23s ± 2.43 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

Редактировать (проверка SVML)

import numba as nb
import numpy as np

@nb.njit(fastmath=True)
def foo(n):
    x   = np.empty(n*8, dtype=np.float64)
    ret = np.empty_like(x)
    for i in range(ret.size):
            ret[i] += np.cos(x[i])
    return ret

foo(1000)

if 'intel_svmlcc' in foo.inspect_llvm(foo.signatures[0]):
    print("found")
else:
    print("not found")

#found

Если есть not found, прочитайте эту ссылку. Она должна работать на Linux и Windows, но я не проверял ее на macOS.

Answer 2

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

Главное, что замедляло ваш код, это избыточные пересчеты одного и того же. Вложенный l oop здесь является преступником:

for p in range(N):
        for m in range(N):
            for i in range(n**2):
                for j in range(n**2):        
                    chi[p,m] += 2/(n**2)*np.dot(conf[i],conf[j])*np.cos(kx[p]*(x[i]-x[j])+ ky[m]*(y[i]-y[j]))

Он содержит много избыточности, многократно пересчитывая векторные операции.

Рассмотрим np.dot (...) : этот расчет полностью не зависит от точек kx и ky. Но только точки kx и ky требовали индексации по m и n. Таким образом, вы можете запустить точечные произведения для всех i и j только один раз и сохранить результат, в отличие от пересчета для каждого m, n (что будет 10000 раз!).

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

Таким образом, зафиксировав циклы и использовав словари с индексами (i , j) в качестве ключей для хранения всех значений вы можете просто найти соответствующее значение в течение l oop над i, j. Вот мой код:

def spin_spin(S, N):
    n = len(S)
    conf = np.reshape(S,(n**2, 3))

    chi = np.zeros((N, N))
    kx = np.linspace(-5*np.pi/3, 5*np.pi/3, N)
    ky = np.linspace(-3*np.pi/np.sqrt(3), 3*np.pi/np.sqrt(3), N)

    # Minor point; no need to use triangular twice
    x, y = triangular(n)
    x, y = np.reshape(x,(n**2)), np.reshape(y,(n**2))

    # Build a look-up for all the dot products to save calculating them many times
    dot_prods = dict()
    x_diffs, y_diffs = dict(), dict()
    for i, j in itertools.product(range(n**2), range(n**2)):
        dot_prods[(i, j)] = np.dot(conf[i], conf[j])
        x_diffs[(i, j)], y_diffs[(i, j)] = x[i] - x[j], y[i] - y[j]    

    # Minor point; improve syntax by converting nested for loops to one line
    for p, m in itertools.product(range(N), range(N)):
        for i, j in itertools.product(range(n**2), range(n**2)):
            # All vector operations are replaced by look ups to the dictionaries defined above
            chi[p, m] += 2/(n**2)*dot_prods[(i, j)]*np.cos(kx[p]*(x_diffs[(i, j)]) + ky[m]*(y_diffs[(i, j)]))
    return(chi, kx, ky)

В данный момент я выполняю это с указанными вами размерами на приличной машине, и l oop над i, j заканчивается через две минуты. Это должно произойти только один раз; тогда это всего лишь oop над m, n. Каждый из них занимает около 90 секунд, так что время работы 2-3 часа. Я приветствую любые предложения о том, как оптимизировать этот расчет, чтобы ускорить его!

Я столкнулся с низко висящим плодом оптимизации, но чтобы дать ощущение скорости, l oop из i, j занимает 2 минут, и таким образом он работает на 9 999 раз меньше!