Оптимизация органайзера функций системы уравнений с узким местом в памяти в python

Код

Итак, у меня есть следующая функция, которая просто собирает систему уравнений, чтобы ее можно было ввести в стандартные решатели scipy.

def EDP7 (v):
    
    T_N = np.reshape(v, (N_points_1, N_points_2)) #Variable declaration
    
                 
    #Model
    sys = np.empty_like(v)
    l=0
    ##Boundary Condition x = 0
    for i in range(N_points_1):
        BC1 = T_N[i,0] - T_x_0
        sys[l] = BC1
        l+=1
    
    ##Boundary Condition x = L
    for i in range(N_points_1):
        BC2 = T_N[i,-1] - T_x_L
        sys[l] = BC2
        l+=1
    
    ##Initial Condition
    for j in range(1,N_points_2-1):
        IC1 = T_N[0,j] - T_t_0
        sys[l] = IC1
        l+=1
    
    ## Energy Balance 
    for i in range(N_points_1-1):
        for j in range (1,N_points_2 -1):
            EB = alpha*(T_N[i,j+1]-2*T_N[i,j]+T_N[i,j-1])/dx**2 - (T_N[i+1,j]-T_N[i,j])/dt  
            sys[l] = EB
            l+=1
    
    return sys

Параметры здесь:

t_max = 25
x_max = 1
N_points_1 = 1000000
N_points_2 = 11
dt = t_max/(N_points_1-1) 
dx = x_max/(N_points_2-1)
N_variables = N_points_1 * N_points_2
alpha = 0.01
T_x_0 = 60 
T_x_L = 25
T_t_0 = 25

initial_guess = T_t_0*np.ones((N_points_1,N_points_2))
initial_guess[:,0] = T_x_0
initial_guess[:,-1] = T_x_L
initial_guess = initial_guess.flatten()  

и я обнаружил, что для задач с большим количеством переменных (что меня интересует), таких как N = 1e6-1e7, значительное время тратится на простое создание системы массива уравнений

sys = np.empty_like(v)

На это уходит примерно половина времени функции (тесты ниже).

Что я пробовал

Ну, я понял, что могу просто объявить вектор снаружи и передать его в функцию в качестве параметра.

syst = np.empty_like(initial_guess)

@njit(fastmath=True)
def EDP8 (v,system):
    sys = system
    T_N = np.reshape(v, (N_points_1, N_points_2))
...
32.3 ms ± 245 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) -Declaration inside

21.9 ms ± 91.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) - Declaration outside

Улучшение около 33%

Однако при передаче в решатели scipy функция отказывается покидать исходную позицию предположения.

Solution_T7 = optimize.least_squares(EDP8,initial_guess, args = [syst], verbose = 2)
Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality   
       0              1         3.0625e+04                                    0.00e+00    
`gtol` termination condition is satisfied.
Function evaluations 1, initial cost 3.0625e+04, final cost 3.0625e+04, first-order optimality 0.00e+00.

print(Solution_7.x)
[60. 25. 25. 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25. 25. 25. 25. 25.
 25. 25. 25. 25. 60. 25. 25. 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25.
 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25. 25. 25. 25. 25. 25. 25. 25.
 25. 60. 25. 25. 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25. 25. 25. 25.
...]

Это массив initial_guess, который я объявил. Я предполагаю, что что-то в махинациях scipy (я тестировал его со многими другими решателями с теми же результатами) препятствует итерациям с переменной syst, объявленной снаружи, может быть, потому, что sys становится глобальной переменной? Кроме того, это не вызвано декоратором njit, то же самое происходит и без него.

В любом случае, есть ли способ устранить узкое место объявления массива исходной функции?

Правки

Это исходный полный код

import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from scipy.optimize import fsolve
from scipy import optimize
import pandas as pd
from numba import njit
import time

t_max = 25
x_max = 1
N_points_1 = 1000000
N_points_2 = 11
dt = t_max/(N_points_1-1) 
dx = x_max/(N_points_2-1)
N_variables = N_points_1 * N_points_2
alpha = 0.01
T_x_0 = 60 
T_x_L = 25
T_t_0 = 25

initial_guess = T_t_0*np.ones((N_points_1,N_points_2))
initial_guess[:,0] = T_x_0
initial_guess[:,-1] = T_x_L
initial_guess = initial_guess.flatten()

@njit(fastmath=True)
def EDP7 (v):
    
    T_N = np.reshape(v, (N_points_1, N_points_2)) #Variable declaration
    
                 
    #Model
    sys = np.empty_like(v)
    l=0
    ##Boundary Condition x = 0
    for i in range(N_points_1):
        BC1 = T_N[i,0] - T_x_0
        sys[l] = BC1
        l+=1

    ##Boundary Condition x = L
    for i in range(N_points_1):
        BC2 = T_N[i,-1] - T_x_L
        sys[l] = BC2
        l+=1

    ##Initial Condition
    for j in range(1,N_points_2-1):
        IC1 = T_N[0,j] - T_t_0
        sys[l] = IC1
        l+=1

    ## Energy Balance 
    for i in range(N_points_1-1):
        for j in range (1,N_points_2 -1):
            EB = alpha*(T_N[i,j+1]-2*T_N[i,j]+T_N[i,j-1])/dx**2 - (T_N[i+1,j]-T_N[i,j])/dt  #partial(T_N,t_1,1)[i,j] - alpha*partial_2(T_N,x,2)[i,j] 
            sys[l] = EB
            l+=1
    
    
    
    return sys

compile_1 = EDP7(initial_guess)

И модифицированный

syst = np.empty_like(initial_guess)

@njit(fastmath=True)
def EDP8 (v,system):
    sys = system
    T_N = np.reshape(v, (N_points_1, N_points_2))
    
                 
    #Model
    
    l=0
    ##Boundary Condition x = 0
    for i in range(N_points_1):
        BC1 = T_N[i,0] - T_x_0
        sys[l] = BC1
        l+=1

    ##Boundary Condition x = L
    for i in range(N_points_1):
        BC2 = T_N[i,-1] - T_x_L
        sys[l] = BC2
        l+=1

    ##Initial Condition
    for j in range(1,N_points_2-1):
        IC1 = T_N[0,j] - T_t_0
        sys[l] = IC1
        l+=1

    ## Energy Balance 
    for i in range(N_points_1-1):
        for j in range (1,N_points_2 -1):
            EB = alpha*(T_N[i,j+1]-2*T_N[i,j]+T_N[i,j-1])/dx**2 - (T_N[i+1,j]-T_N[i,j])/dt  #partial(T_N,t_1,1)[i,j] - alpha*partial_2(T_N,x,2)[i,j] 
            sys[l] = EB
            l+=1
    
    
    
    return sys

compile_2 = EDP8(initial_guess,syst)

И тесты, которые я только что сделал.

%timeit EDP7(initial_guess)
%timeit EDP8(initial_guess,syst)

47.2 ms ± 1.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
23.3 ms ± 633 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Это уже большая разница, чем ОП. Поскольку единственная разница между этими кодами — это внешнее объявление, я полагаю, это из-за него…

1 ответ
1

В исходном посте утверждается, что

    sys = np.empty_like(v)

занимает значительное время.

Я не нахожу доказательств этого ни в ОП, ни в этих тестах:
(См. новые результаты синхронизации ниже.)

$ python -m timeit \
   -s 'import numpy as np; shp = (1000, 1000, 1000); v = np.ones(shp)' \
      'assert shp == np.empty_like(v).shape'
20000 loops, best of 5: 19.2 usec per loop
$
$ python -m timeit \
   -s 'import numpy as np; shp = (10,) * 9; v = np.ones(shp)' \
      'assert shp == np.empty_like(v).shape'
10000 loops, best of 5: 20.4 usec per loop
$
$ python -m timeit \
   -s 'import numpy as np; shp = (2,) * 29; v = np.ones(shp)' \
      'assert shp == np.empty_like(v).shape'
20000 loops, best of 5: 11 usec per loop

Поднимите этот последний до 2 ** 30 приводит нас к «22 мкс на цикл», что связано с началом страницы в резервном хранилище на этом ноутбуке с 8 ГБ, а не с отражением подпрограмм BLAS / numpy.

.empty_like(...) операция занимает незначительное время, даже при работе с миллиардом записей. То есть эмпирически это занимает O(1) времени, независимо от размера v.


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


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

Основываясь на cProfile, описывающем вложенные циклы как использующие большую часть циклов, я сразу понял это:

@njit(fastmath=True)
def energy_balance(t_n, sys, l):
    ...

Я хотел использовать timeit.timeit(), но, увы, это не совместимо с JIT. Итак, у меня это наверху:

@njit(fastmath=True)
def el(v):
    return np.empty_like(v)


# @njit(fastmath=True)
def edp7(v):
    t_n = np.reshape(v, (n_points_1, n_points_2))  # Variable declaration

    # Model
    sys = el(v)
    assert len(sys) == len(v) == 11_000_000
    print("")
    number = 10_000
    elapsed = timeit(lambda: np.empty_like(v), number=number)
    print(f"elapsed: {elapsed:.3f} s, {elapsed/number:.6f} s/loop")

Если вы отрегулируете lambdaвы увидите, что прошедшее время увеличивается в 100 раз, когда мы вызываем тривиальное el обертка вместо empty_like напрямую. Но вызов оболочки является «безопасным» (быстрым), если он не JIT.

Кроме того, вот вторая подсказка. Я закончил с этим внизу:

if __name__ == "__main__":
    edp7(initial_guess.copy())

    t0 = time.time()
    compile_1 = edp7(initial_guess)
    print(f"total:  {time.time() - t0}:.3f s")

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

elapsed: 0.018 s, 0.000002 s/loop

elapsed: 42.895 s, 0.004290 s/loop
total:  43.702 s

timeit фигура просто развалилась во второй раз.

OTOH, если оба звонки делают initial_guess.copy()
тогда все идет гладко:

elapsed: 0.020 s, 0.000002 s/loop

elapsed: 0.019 s, 0.000002 s/loop
total:  0.823 s

Во время прогона на ноутбуке много свободной оперативной памяти. Я не могу объяснить такое поведение. Но это является можно обойти его край и избежать его.


Переназначить не помешает v так

    v = np.float64(v)
    sys = np.empty_like(v)

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

Если вам это сойдет с рук, назначьте np.float32(v) для экономии памяти, что подразумевает экономию пропускной способности памяти. Меньшее количество перемещенных байтов означает меньшее затраченное время.

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *