Векторизация кумулятивной суммы 6 для цикла в python

Question

Математическая задача:

Выражение в суммах на самом деле намного сложнее, чем приведенное выше, но это для минимального рабочего примера, чтобы не слишком усложнять вещи. Я написал это на Python, используя 6 вложенных циклов for, и, как и ожидалось, он работает очень плохо (истинная форма работает плохо и требует оценки миллионы раз), даже с помощью Numba, Cython и его друзей. Здесь это написано с использованием вложенных циклов и совокупной суммы:

import numpy as np

def func1(a,b,c,d):
    '''
    Minimal working example of multiple summation
    '''
    B = 0
    for ai in range(0,a):
        for bi in range(0,b):
            for ci in range(0,c):
                for di in range(0,d):
                    for ei in range(0,ai+bi):
                        for fi in range(0,ci+di):
                            B += (2)**(ei-fi-ai-ci-di+1)*(ei**2-2*(ei*fi)-7*di)*np.math.factorial(ei)


    return a, b, c, d, B

Выражение управляется 4 цифрами в качестве ввода, а для func1(4,6,3,4) для B выводится 21769947.844726562.

Я искал помощи в этом и нашел несколько сообщений в стеке, с некоторыми примерами:

Внешний продукт в NumPy: векторизация шести вложенных циклов

Векторизация triple для цикла в Python / Numpy с различными формами массива

Векторизация Python, вложенная для циклов

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

Answer 1

Это только комментарий к ответу @jdehesa.

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

Код

import numpy as np
import numba as nb

@nb.njit()
def factorial(a):
  res=1.
  for i in range(1,a+1):
    res*=i
  return res

@nb.njit()
def func1(a, b, c, d):
    B = 0.

    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)

    fact_e=np.empty(a + b - 2)
    for i in range(a + b - 2):
      fact_e[i]=factorial(i)

    for ai in range(0, a):
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
    return B

Параллельная версия

@nb.njit(parallel=True)
def func_p(a_vec,b_vec,c_vec,d_vec):
  res=np.empty(a_vec.shape[0])
  for i in nb.prange(a_vec.shape[0]):
    res[i]=func1(a_vec[i], b_vec[i], c_vec[i], d_vec[i])
  return res

Пример

a_vec=np.random.randint(low=2,high=10,size=1000000)
b_vec=np.random.randint(low=2,high=10,size=1000000)
c_vec=np.random.randint(low=2,high=10,size=1000000)
d_vec=np.random.randint(low=2,high=10,size=1000000)

res_2=func_p(a_vec,b_vec,c_vec,d_vec)

Однопоточная версия приводит к 5,6 мкс (после первого запуска) в вашем примере.

Параллельверсия почти приведет к другому ускорению Number_of_Cores для расчета многих значений.Имейте в виду, что издержки компиляции больше для параллельной версии (более 0,5 с для первого вызова).

Answer 2

Используя эту функцию cartesian_product , вы можете преобразовать вложенный цикл в матрицы, а затем просто вычислить соответствующие вложенные сигмы в векторизованном виде:

In [37]: def nested_sig(args):
    ...:     base_prod = cartesian_product(*arrays)
    ...:     second_prod = cartesian_product(base_prod[:,:2].sum(1), base_prod[:,2:].sum(1))
    ...:     total = np.column_stack((base_prod, second_prod))
    ...:     # the items in each row denotes the following variables in order:
    ...:     # ai, bi, ci, di, ei, fi
    ...:     x = total[:, 4] - total[:, 5] - total[:, 0] - total[:, 2] - total[:, 3] + 1
    ...:     y = total[:, 4] - total[:, 5]
    ...:     result = np.power(2, x) * (np.power(total[:, 4], 2) - 2*y - 7*total[:, 3]) * np.math.factorial(total[:,4])
    ...:     return result

Answer 3

Я вижу три источника улучшения в вашем коде:

• range(0,a) is

• Вы выполняете большую работу во внутреннем цикле

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

Здесь версия (вероятно, еще не хорошая), которая пытается улучшить эти очки.

@numba.njit
def func1o(a,b,c,d):
    "2**(ei-fi-ai-ci-di+1)*(ei**2-2*(ei*fi)-7*di)*ei!"                    
    POW=2.;                 SUM=0.;              
    L=[]
    for ai in arange(0.,a+1):
        for bi in range(0,b+1):
            for ci in range(0,c+1):
                for di in range(0,d+1):
                    FACT=1.
                    for ei in arange(0,ai+bi+1):
                        for fi in range(0,ci+di+1):
                            L.append(POW*SUM*FACT)
                            POW /= 2
                            SUM -= 2*ei
                        POW *= 2    
                        SUM += 2*(ei-fi)+1
                        FACT *= ei+1
                    POW /=2
                    SUM -= 7*di
                POW /= 2
        POW /= 2
    A=np.array(L)
    I=np.abs(A).argsort()
    return A[I].sum()    

Answer 4

РЕДАКТИРОВАТЬ 3:

Окончательная (я думаю) версия, немного чище и быстрее, включающая идеи из max9111 ответа .

import numpy as np
from numba import as nb

@nb.njit()
def func1_jit(a, b, c, d):
    # Precompute
    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)
    fact_e = np.empty((a + b - 2))
    fact_e[0] = 1
    for ei in range(1, len(fact_e)):
        fact_e[ei] = ei * fact_e[ei - 1]
    # Loops
    B = 0
    for ai in range(0, a):
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
    return B

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

import numpy as np
from numba import as nb

@nb.njit(parallel=True)
def func1_par(a, b, c, d):
    # Precompute
    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)
    fact_e = np.empty((a + b - 2))
    fact_e[0] = 1
    for ei in range(1, len(fact_e)):
        fact_e[ei] = ei * fact_e[ei - 1]
    # Loops
    B = np.empty((a,))
    for ai in nb.prange(0, a):
        Bi = 0
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            Bi += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
        B[ai] = Bi
    return np.sum(B)

Или, если у вас есть много точек, где вы хотите оценить функцию, вы также можете распараллелить на этом уровне. Здесь a_arr, b_arr, c_arr и d_arr - векторы значений, где должна оцениваться функция:

from numba import as nb

@nb.njit(parallel=True)
def func1_arr(a_arr, b_arr, c_arr, d_arr):
    B_arr = np.empty((len(a_arr),))
    for i in nb.prange(len(B_arr)):
        B_arr[i] = func1_jit(a_arr[i], b_arr[i], c_arr[i], d_arr[i])
    return B_arr

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

---

РЕДАКТИРОВАТЬ 2:

На самом деле, забудь, что я говорил раньше. Лучше всего JIT-компилировать алгоритм, но более эффективно. Сначала вычислите дорогие части (я взял экспоненту и факториал), а затем передайте их скомпилированной циклической функции:

import numpy as np
from numba import njit

def func1(a, b, c, d):
    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)
    ee = np.arange(a + b - 2)
    fact_e = scipy.special.factorial(ee)
    return func1_inner(a, b, c, d, exp_min, exp, fact_e)

@njit()
def func1_inner(a, b, c, d, exp_min, exp, fact_e):
    B = 0
    for ai in range(0, a):
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
    return B

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

a, b, c, d = 4, 6, 3, 4
# The original function
%timeit func1_orig(a, b, c, d)
# 2.07 ms ± 33.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# The grid-evaluated function
%timeit func1_grid(a, b, c, d)
# 256 µs ± 25 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
# The precompuation + JIT-compiled function
%timeit func1_jit(a, b, c, d)
# 19.6 µs ± 3.25 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)

---

Ну, всегда есть возможность оценить всю сетку целиком:

import numpy as np
import scipy.special

def func1(a, b, c, d):
    ai, bi, ci, di, ei, fi = np.ogrid[:a, :b, :c, :d, :a + b - 2, :c + d - 2]
    # Compute
    B = (2.) ** (ei - fi - ai - ci - di + 1) * (ei ** 2 - 2 * (ei * fi) - 7 * di) * scipy.special.factorial(ei)
    # Mask out of range elements for last two inner loops
    m = (ei < ai + bi) & (fi < ci + di)
    return np.sum(B * m)

print(func1(4, 6, 3, 4))
# 21769947.844726562

Я использовал scipy.special.factorial, потому что, по-видимому, np.factorial по некоторым причинам не работает с массивами.

Очевидно, что стоимость памяти будет расти очень быстро при увеличении параметров. Код на самом деле выполняет больше вычислений, чем необходимо, потому что два внутренних цикла имеют различное количество итераций, поэтому (в этом методе) вам нужно использовать наибольшее, а затем удалить то, что вам не нужно. Надеюсь, что векторизация восполнит это. Небольшой тест IPython:

a, b, c, d = 4, 6, 3, 4
# func1_orig is the original loop-based version
%timeit func1_orig(a, b, c, d)
# 2.9 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# func1 here is the vectorized version
%timeit func1(a, b, c, d)
# 210 µs ± 6.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

---

EDIT:

Обратите внимание, что предыдущий подход также не является "все или ничего". Вы можете выбрать для оценки сетки только некоторые из циклов. Например, два самых внутренних цикла могут быть векторизованы так:

def func1(a, b, c, d):
    B = 0
    e = np.arange(a + b - 2).reshape((-1, 1))
    f = np.arange(c + d - 2)
    for ai in range(0, a):
        for bi in range(0, b):
            ei = e[:ai + bi]
            for ci in range(0, c):
                for di in range(0, d):
                    fi = f[:ci + di]
                    B += np.sum((2.) ** (ei - fi - ai - ci - di + 1) * (ei ** 2 - 2 * (ei * fi) - 7 * di) * scipy.special.factorial(ei))
    return B

Это все еще имеет циклы, но оно избегает дополнительных вычислений, и требования к памяти намного ниже. Какой из них лучше, зависит от размеров входов, я думаю. В моих тестах с исходными значениями (4, 6, 3, 4) это даже медленнее, чем исходная функция; Кроме того, в этом случае кажется, что создание новых массивов для ei и fi в каждом цикле происходит быстрее, чем работа с фрагментом предварительно созданного. Однако, если вы умножите входное значение на 4 (14, 24, 12, 16), то это будет намного быстрее, чем оригинал (около х5), хотя все же медленнее, чем полностью векторизованный (около х3). С другой стороны, я мог бы вычислить значение для входа, масштабированное до десяти (40, 60, 30, 40) с этим (за ~ 5 минут), но не с предыдущим из-за памяти (я не проверял, как долго бы с оригинальной функцией). Использование @numba.jit немного помогает, хотя и не сильно (не может использовать nopython из-за функции факториала). Вы можете поэкспериментировать с векторизацией более или менее циклов в зависимости от размера ваших входных данных.