Оптимизация кода для дифференциальных уравнений
|
См. эти ответы на вопросы для получения информации о типичных сложностях и способах повышения производительности. |
Оптимизация кода в Julia
Прежде чем приступить к работе с данным руководством, рекомендуется ознакомиться с одним из многочисленных руководств по оптимизации кода Julia. Вот неполный список:
Оптимизации на стороне пользователя имеют большое значение, так как для достаточно сложных задач значительная часть времени будет проводиться внутри функции f — функции, которую вы пытаетесь решить. Эффективные интеграторы позволяют сократить необходимое количество вызовов функции f для получения допустимой погрешности. Основные идеи по оптимизации кода DiffEq или любой функции Julia заключаются в следующем.
-
Сделать его невыделяемым
-
Использовать StaticArrays для небольших массивов
-
Использовать слияние трансляции
-
Сделать его стабильным по типу
-
Сократить избыточные вычисления
-
Использовать вызовы BLAS
-
Оптимизировать выбор алгоритма
Мы обсудим эти стратегии в контексте дифференциальных уравнений. Начнем с небольших систем.
Пример ускорения решения нежесткого уравнения. Уравнение Лоренца
Возьмем классическую систему Лоренца. Сначала просто напишем систему в ее форме не на месте:
function lorenz(u, p, t)
dx = 10.0 * (u[2] - u[1])
dy = u[1] * (28.0 - u[3]) - u[2]
dz = u[1] * u[2] - (8 / 3) * u[3]
[dx, dy, dz]
endlorenz (generic function with 1 method)Здесь lorenz возвращает объект [dx,dy,dz], который создается в теле lorenz.
Это распространенный шаблон кода из языков высокого уровня, таких как MATLAB, SciPy или deSolve в R. Однако проблема этой формы заключается в том, что на каждом шаге выделяется вектор [dx,dy,dz]. Проведем сравнительный анализ процесса решения при таком выборе функции:
using DifferentialEquations, BenchmarkTools
u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 100.0)
prob = ODEProblem(lorenz, u0, tspan)
@btime solve(prob, Tsit5()); 3.811 ms (101102 allocations: 7.82 MiB)@benchmark пакета BenchmarkTools.jl выполняет код несколько раз, чтобы получить точное измерение. Минимальное время — это время, которое требуется, когда ОС и другие фоновые процессы не мешают работе. Обратите внимание, что в данном случае решение занимает около 5 мс и выделяет около 11,11 МиБ. Однако если использовать это в реальном пользовательском коде, мы увидим, что много времени тратится на сборку мусора (GC) для очистки всех созданных массивов. Даже если отключить сохранение, эти выделения все равно останутся.
@btime solve(prob, Tsit5(), save_everystep = false); 3.098 ms (89529 allocations: 6.83 MiB)Проблема заключается в том, что массивы создаются каждый раз, когда вызывается производная функция. Эта функция вызывается несколько раз за шаг и, таким образом, является основным источником потребления памяти. Для исправления этого момента можно использовать форму на месте, чтобы сделать код не поддерживающим выделение:
function lorenz!(du, u, p, t)
du[1] = 10.0 * (u[2] - u[1])
du[2] = u[1] * (28.0 - u[3]) - u[2]
du[3] = u[1] * u[2] - (8 / 3) * u[3]
nothing
endlorenz! (generic function with 1 method)Здесь вместо того, чтобы каждый раз создавать массив, мы использовали массив кэша du. Когда действует форма на месте, DifferentialEquations.jl использует другой внутренний маршрут, который также минимизирует внутренние выделения.
|
Обратите внимание, что ничего не возвращается. В форме на месте решатель ODE игнорирует возврат. Убедитесь, что исходный массив |
При поведении анализа этой функции мы увидим существенную разницу.
u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 100.0)
prob = ODEProblem(lorenz!, u0, tspan)
@btime solve(prob, Tsit5()); 806.701 μs (11415 allocations: 996.33 KiB)@btime solve(prob, Tsit5(), save_everystep = false); 368.701 μs (49 allocations: 4.00 KiB)Только из-за этого изменения достигнута 16-кратная разница во времени. Однако некоторые выделения все же есть, и это связано с построением кэша интеграции. Но они не масштабируются с размером задачи:
tspan = (0.0, 500.0) # в 5 раз дольше, чем раньше
prob = ODEProblem(lorenz!, u0, tspan)
@btime solve(prob, Tsit5(), save_everystep = false); 1.708 ms (49 allocations: 4.00 KiB)Поскольку это все выделения при настройке, оптимизация на стороне пользователя завершена.
Дальнейшая оптимизация небольших нежестких обыкновенных дифференциальных уравнений с помощью StaticArrays
Выделения являются дорогими только в том случае, если они являются выделениями памяти в куче. Более подробное определение выделений памяти в куче можно найти во многих источниках в Интернете. Но хорошим рабочим определением является то, что выделения памяти в куче — это участки памяти переменного размера, на которые необходимо указывать, и эта косвенная адресация указателей стоит времени. Кроме того, кучей необходимо управлять, контроллеры мусора должны активно следить за тем, что находится в куче.
Однако существует альтернатива выделению памяти в куче, известная как выделение памяти в стеке. Стек имеет статический размер (известный во время компиляции), поэтому доступ к нему осуществляется быстро. Кроме того, компилятору заранее известен точный блок памяти, и поэтому повторное использование памяти обходится дешево. Это означает, что выделение в стеке практически ничего не стоит.
Массивы должны выделяться в куче, поскольку их размер (а значит, и объем занимаемой ими памяти) определяется во время выполнения программы. Но в Julia есть структуры с выделением памяти в стеке. Например, структуры (struct) — это «значения-типы» с выделением памяти в стеке. Кортежи (Tuple) — это коллекция с выделением памяти в стеке. Однако наиболее полезной структурой данных для DiffEq является структура StaticArray из пакета StaticArrays.jl. Длина этих массивов определяется во время компиляции. Они создаются с помощью макросов, присоединенных к обычным выражениям массивов, например:
using StaticArrays
A = SA[2.0, 3.0, 5.0]
typeof(A) # SVector{3, Float64} (alias for SArray{Tuple{3}, Float64, 1, 3})SVector{3, Float64} (alias for SArray{Tuple{3}, Float64, 1, 3})Обратите внимание, что 3 после SVector дает размер SVector. Его нельзя изменить. Кроме того, SVector являются неизменяемыми, поэтому для изменения значений необходимо создать новый SVector. Но помните, что беспокоиться о выделении не нужно, поскольку эта структура данных выделяется в стеке. SArray имеют множество дополнительных оптимизаций: быстрое умножение матриц, быстрые QR-разложения и т. д., которые напрямую используют информацию о размере массива. Поэтому, когда это возможно, следует применять именно их.
К сожалению, статические массивы можно использовать только для достаточно небольших массивов данных. После определенного размера для них требуется выделять память в куче после некоторых инструкций, и их время компиляции возрастает. Таким образом, статические массивы не следует использовать, если в системе более ~20 переменных. Кроме того, в полной мере использовать статические массивы могут только собственные алгоритмы Julia.
Давайте оптимизируем lorenz с использованием статических массивов. Обратите внимание, что в данном случае мы хотим использовать форму выделения не на месте, но на этот раз нам нужно вывести статический массив:
function lorenz_static(u, p, t)
dx = 10.0 * (u[2] - u[1])
dy = u[1] * (28.0 - u[3]) - u[2]
dz = u[1] * u[2] - (8 / 3) * u[3]
SA[dx, dy, dz]
endlorenz_static (generic function with 1 method)Чтобы решатель внутренним образом использовал статические массивы, достаточно задать ему статический массив в качестве начального условия:
u0 = SA[1.0, 0.0, 0.0]
tspan = (0.0, 100.0)
prob = ODEProblem(lorenz_static, u0, tspan)
@btime solve(prob, Tsit5()); 326.822 μs (1293 allocations: 387.39 KiB)@btime solve(prob, Tsit5(), save_everystep = false); 193.331 μs (22 allocations: 2.25 KiB)Вот, собственно, и все. При работе со статическими массивами не нужно беспокоиться о выделении, поэтому используйте операции типа * и не думайте об операциях слияния (речь о них пойдет в следующем разделе). Создайте векторизованный код на R/MATLAB/Python, и ваш код будет быстрым, или используйте непосредственно числа/значения.
Пример ускорения жесткого уравнения: уравнение Робертсона
Для следующих примеров решим уравнения Робертсона (также известные как ROBER).
Учитывая, что эти уравнения являются жесткими, нежесткие решатели ODE, такие как Tsit5 или Vern9, не смогут решить эти уравнения. Автоматический алгоритм обнаружит это и автоматически переключится на что-то более надежное и подходящее для решения этих задач. Например:
using DifferentialEquations
using Plots
function rober!(du, u, p, t)
y₁, y₂, y₃ = u
k₁, k₂, k₃ = p
du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
du[2] = k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃
du[3] = k₂ * y₂^2
nothing
end
prob = ODEProblem(rober!, [1.0, 0.0, 0.0], (0.0, 1e5), [0.04, 3e7, 1e4])
sol = solve(prob)
plot(sol, tspan = (1e-2, 1e5), xscale = :log10)using BenchmarkTools
@btime solve(prob); 93.616 μs (667 allocations: 58.44 KiB)Выбор хорошего решателя
Выбор хорошего решателя необходим для получения максимальной скорости. Общие рекомендации можно найти на странице решателя (например, Рекомендации для решателя ODE). Существующие рекомендации можно упростить до метода Розенброка (Rosenbrock23 или Rodas5) для небольших (<50 ODE) задач, методов ESDIRK для несколько больших (TRBDF2 или KenCarp4 для <2000 ODE) и QNDF для еще больших задач. Иногда стоит попробовать lsoda из LSODA.jl для категории среднего размера.
Более подробную информацию о том, какой решатель выбрать, можно получить с помощью сравнительных анализов. Сравнение многих решателей для многих задач см. на странице SciMLBenchmarks.
Исходя из этого, мы пробуем применить рекомендацию Rosenbrock23() для жестких уравнений ODE при стандартных допусках:
@btime solve(prob, Rosenbrock23()); 83.664 μs (500 allocations: 40.70 KiB)Объявление функции Якоби
Для уменьшения стоимости построения якобиана можно описать функцию якобиана, используя аргумент jac для ODEFunction. Сначала необходимо получить якобиан , который равен J[i,j]. Отсюда получаем:
function rober_jac!(J, u, p, t)
y₁, y₂, y₃ = u
k₁, k₂, k₃ = p
J[1, 1] = k₁ * -1
J[2, 1] = k₁
J[3, 1] = 0
J[1, 2] = y₃ * k₃
J[2, 2] = y₂ * k₂ * -2 + y₃ * k₃ * -1
J[3, 2] = y₂ * 2 * k₂
J[1, 3] = k₃ * y₂
J[2, 3] = k₃ * y₂ * -1
J[3, 3] = 0
nothing
end
f! = ODEFunction(rober!, jac = rober_jac!)
prob_jac = ODEProblem(f!, [1.0, 0.0, 0.0], (0.0, 1e5), (0.04, 3e7, 1e4))ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 100000.0)
u0: 3-element Vector{Float64}:
1.0
0.0
0.0@btime solve(prob_jac, Rosenbrock23()); 87.647 μs (422 allocations: 34.70 KiB)Автоматическое выведение функций Якоби
Но это было трудно. Если требуется символический якобиан численного кода, можно воспользоваться пакетом ModelingToolkit.jl для символической модификации численного кода, произвести символическое вычисление и вернуть для этого код Julia.
using ModelingToolkit
de = modelingtoolkitize(prob)Model ##MTKizedODE#235536 with 3 equations
States (3):
x₁(t) [defaults to 1.0]
x₂(t) [defaults to 0.0]
x₃(t) [defaults to 0.0]
Parameters (3):
α₁ [defaults to 0.04]
α₂ [defaults to 3.0e7]
α₃ [defaults to 10000.0]Если нужно увидеть код, можно указать ему вычислить якобиан:
ModelingToolkit.generate_jacobian(de)[2] # Второй на месте:(function (ˍ₋out, ˍ₋arg1, ˍ₋arg2, t)
#= /root/.julia/packages/SymbolicUtils/ssQsQ/src/code.jl:373 =#
#= /root/.julia/packages/SymbolicUtils/ssQsQ/src/code.jl:374 =#
#= /root/.julia/packages/SymbolicUtils/ssQsQ/src/code.jl:375 =#
begin
begin
begin
#= /root/.julia/packages/Symbolics/gBKZv/src/build_function.jl:537 =#
#= /root/.julia/packages/SymbolicUtils/ssQsQ/src/code.jl:422 =# @inbounds begin
#= /root/.julia/packages/SymbolicUtils/ssQsQ/src/code.jl:418 =#
ˍ₋out[1] = (*)(-1, ˍ₋arg2[1])
ˍ₋out[2] = ˍ₋arg2[1]
ˍ₋out[3] = 0
ˍ₋out[4] = (*)(ˍ₋arg1[3], ˍ₋arg2[3])
ˍ₋out[5] = (+)((*)((*)(-2, ˍ₋arg1[2]), ˍ₋arg2[2]), (*)((*)(-1, ˍ₋arg1[3]), ˍ₋arg2[3]))
ˍ₋out[6] = (*)((*)(2, ˍ₋arg1[2]), ˍ₋arg2[2])
ˍ₋out[7] = (*)(ˍ₋arg1[2], ˍ₋arg2[3])
ˍ₋out[8] = (*)((*)(-1, ˍ₋arg1[2]), ˍ₋arg2[3])
ˍ₋out[9] = 0
#= /root/.julia/packages/SymbolicUtils/ssQsQ/src/code.jl:420 =#
nothing
end
end
end
end
end)Теперь воспользуемся этим для получения якобиана аналитического решения:
prob_jac2 = ODEProblem(de, [], (0.0, 1e5), jac = true)ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 100000.0)
u0: 3-element Vector{Float64}:
1.0
0.0
0.0@btime solve(prob_jac2); 93.612 μs (606 allocations: 58.16 KiB)Более подробную информацию см. в документации по ModelingToolkit.jl.
Ускорение решения небольших обыкновенных дифференциальных уравнение с помощью StaticArrays
Если система ODE достаточно невелика (<20 ODE или около того), использование StaticArrays.jl для переменных состояния может значительно повысить производительность. Это делается путем преобразования u0 в StaticArray и написания неизменяемой диспетчеризации не на месте для статических массивов. Для задачи ROBER эти действия выглядят следующим образом.
using StaticArrays
function rober_static(u, p, t)
y₁, y₂, y₃ = u
k₁, k₂, k₃ = p
du1 = -k₁ * y₁ + k₃ * y₂ * y₃
du2 = k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃
du3 = k₂ * y₂^2
SA[du1, du2, du3]
end
prob = ODEProblem(rober_static, SA[1.0, 0.0, 0.0], (0.0, 1e5), SA[0.04, 3e7, 1e4])
sol = solve(prob, Rosenbrock23())retcode: Success
Interpolation: specialized 2nd order "free" stiffness-aware interpolation
t: 61-element Vector{Float64}:
0.0
3.196206628740808e-5
0.00014400709336278452
0.00025605212043816096
0.00048593871402339607
0.0007179482102678373
0.0010819240251828343
0.0014801655107859655
0.0020679567717440095
0.002843584518457066
⋮
25371.93159838571
30784.11718374498
37217.42390396605
44850.61094811346
53893.688830057334
64593.73530179436
77241.71691097679
92180.81843146283
100000.0
u: 61-element Vector{SVector{3, Float64}}:
[1.0, 0.0, 0.0]
[0.9999987215181657, 1.2780900152625978e-6, 3.9181897521319503e-10]
[0.9999942397329006, 5.7185104612947566e-6, 4.175663804739006e-8]
[0.9999897579688383, 9.992106612572491e-6, 2.49924549040571e-7]
[0.9999805626683271, 1.7833623941038088e-5, 1.6037077316934769e-6]
[0.9999712826607852, 2.403488562731424e-5, 4.682453587410618e-6]
[0.9999567250114038, 3.0390689334989113e-5, 1.2884299261094982e-5]
[0.9999407986095145, 3.388427339038224e-5, 2.531711709494679e-5]
[0.9999172960310598, 3.583508669306405e-5, 4.686888224684217e-5]
[0.9998862913763157, 3.6412401619257426e-5, 7.729622206475401e-5]
⋮
[0.05563508171413305, 2.3546322394505495e-7, 0.9443646828226426]
[0.047925352159210115, 2.012149653947115e-7, 0.9520744466258239]
[0.04123342367542567, 1.7192789116847091e-7, 0.9587664043966821]
[0.03543700020207701, 1.4688362762022196e-7, 0.9645628529142937]
[0.03042537309965345, 1.2546809864160592e-7, 0.9695745014322467]
[0.026099133126498978, 1.071560882169501e-7, 0.9739007597174127]
[0.022369692367946337, 9.149845157822593e-8, 0.977630216133602]
[0.019158563494465593, 7.811096455346351e-8, 0.9808413583945704]
[0.017827893845894716, 7.258919980139027e-8, 0.982172033564906]Если провести сравнительный анализ, мы увидим действительно быстрое решение с очень низким количеством выделений:
@btime sol = solve(prob, Rosenbrock23()); 77.852 μs (807 allocations: 46.66 KiB)Таким образом, этот вариант очень удобен для многопоточности и других форм параллелизма.
Пример ускорения полудискретизации дифференциальных уравнений в частных производных линейной алгебры
В этом руководстве мы оптимизируем определение правой части полудискретизации PDE.
|
Настоятельно рекомендуется обратиться к руководству Решение больших жестких уравнений для получения сведений о настройке DifferentialEquations.jl для более эффективного решения больших жестких уравнений ODE. В этом разделе будет рассматриваться только код на стороне пользователя. |
Оптимизируем решение дискретизации реакционно-диффузионного дифференцированного уравнения в частных производных (PDE). В дискретизированном виде это и есть ODE:
где , и являются матрицами. Здесь мы будем использовать упрощенный вариант, где — трехдиагональный шаблон ], т. е. это двухмерная дискретизация лапласиана. Собственный код будет выглядеть примерно так:
using DifferentialEquations, LinearAlgebra, BenchmarkTools
# Создает константы
p = (1.0, 1.0, 1.0, 10.0, 0.001, 100.0) # a, α, ubar, β, D1, D2
N = 100
Ax = Array(Tridiagonal([1.0 for i in 1:(N - 1)], [-2.0 for i in 1:N],
[1.0 for i in 1:(N - 1)]))
Ay = copy(Ax)
Ax[2, 1] = 2.0
Ax[end - 1, end] = 2.0
Ay[1, 2] = 2.0
Ay[end, end - 1] = 2.0
function basic_version!(dr, r, p, t)
a, α, ubar, β, D1, D2 = p
u = r[:, :, 1]
v = r[:, :, 2]
Du = D1 * (Ay * u + u * Ax)
Dv = D2 * (Ay * v + v * Ax)
dr[:, :, 1] = Du .+ a .* u .* u ./ v .+ ubar .- α * u
dr[:, :, 2] = Dv .+ a .* u .* u .- β * v
end
a, α, ubar, β, D1, D2 = p
uss = (ubar + β) / α
vss = (a / β) * uss^2
r0 = zeros(100, 100, 2)
r0[:, :, 1] .= uss .+ 0.1 .* rand.()
r0[:, :, 2] .= vss
prob = ODEProblem(basic_version!, r0, (0.0, 0.1), p)ODEProblem with uType Array{Float64, 3} and tType Float64. In-place: true
timespan: (0.0, 0.1)
u0: 100×100×2 Array{Float64, 3}:
[:, :, 1] =
11.0526 11.0371 11.0016 11.0671 … 11.0646 11.0927 11.0524 11.0093
11.0346 11.0127 11.0966 11.0603 11.0675 11.0779 11.0317 11.0259
11.0973 11.0786 11.0122 11.0395 11.0395 11.0371 11.0182 11.055
11.012 11.0846 11.0813 11.0175 11.0453 11.0735 11.0837 11.0431
11.0383 11.0138 11.04 11.0509 11.0291 11.0244 11.0333 11.0639
11.0259 11.0033 11.099 11.0404 … 11.0187 11.0234 11.0225 11.0917
11.0411 11.0302 11.0342 11.055 11.0398 11.0109 11.0293 11.0036
11.0213 11.0912 11.017 11.021 11.018 11.0483 11.0937 11.0686
11.0182 11.0595 11.0159 11.0497 11.0766 11.059 11.0044 11.0889
11.0087 11.0789 11.006 11.0883 11.0499 11.0549 11.0306 11.0491
⋮ ⋱
11.0523 11.0989 11.0192 11.0676 11.0435 11.0253 11.0076 11.0363
11.0321 11.0323 11.0439 11.0013 11.0328 11.0999 11.0413 11.0565
11.0354 11.0117 11.0551 11.0134 11.0602 11.0361 11.0991 11.0881
11.0726 11.0982 11.0956 11.065 11.04 11.0811 11.0952 11.0517
11.089 11.0327 11.0454 11.095 … 11.0634 11.0918 11.0079 11.0285
11.0761 11.0777 11.0622 11.0641 11.0303 11.0036 11.048 11.0537
11.0986 11.0366 11.0918 11.0772 11.0287 11.0732 11.0879 11.0723
11.0502 11.0094 11.0188 11.0629 11.0331 11.0535 11.0944 11.0599
11.0796 11.0903 11.049 11.0034 11.0749 11.0019 11.0334 11.0385
[:, :, 2] =
12.1 12.1 12.1 12.1 12.1 12.1 … 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 … 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
⋮ ⋮ ⋱ ⋮
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 … 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1
12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1 12.1В этом варианте мы закодировали начальное условие в виде трехмерного массива, где u[:,:,1] является частью A, а u[:,:,2] — частью B.
@btime solve(prob, Tsit5()); 73.064 ms (7349 allocations: 186.83 MiB)Хотя этот вариант не очень эффективен,
рекомендуем сначала написать «высокоуровневый код», а затем итеративно оптимизировать его
Первое, что мы можем сделать, — это избавиться от нарезки выделений. Операция r[:,:,1] создает временный массив вместо представления, т. е. указателя на уже существующую память. Чтобы сделать его представлением, добавьте @view. Обратите внимание, что с представлениями нужно работать аккуратно, поскольку они указывают на одну и ту же память, и, следовательно, при изменении представления изменяются исходные значения:
A = rand(4)
@show A
B = @view A[1:3]
B[2] = 2
@show A4-element Vector{Float64}:
0.7835981747579798
2.0
0.9033405273644081
0.3312967933875772Обратите внимание, что изменение B привело к изменению A. С этим нужно быть осторожным, но в то же время мы хотим использовать эту возможность, поскольку хотим модифицировать вывод dr. Кроме того, последнее выражение является чисто поэлементной операцией, и поэтому здесь можно использовать слияние трансляции. Перепишем basic_version!, чтобы избежать нарезки выделений и использовать слияние трансляции:
function gm2!(dr, r, p, t)
a, α, ubar, β, D1, D2 = p
u = @view r[:, :, 1]
v = @view r[:, :, 2]
du = @view dr[:, :, 1]
dv = @view dr[:, :, 2]
Du = D1 * (Ay * u + u * Ax)
Dv = D2 * (Ay * v + v * Ax)
@. du = Du + a .* u .* u ./ v + ubar - α * u
@. dv = Dv + a .* u .* u - β * v
end
prob = ODEProblem(gm2!, r0, (0.0, 0.1), p)
@btime solve(prob, Tsit5()); 47.886 ms (5879 allocations: 119.71 MiB)Теперь большая часть выделений происходит в Du = D1*(Ay*u + u*Ax), поскольку эти операции векторизованы и не изменяются. Матричные умножения следует заменить mul!. При этом нам понадобятся переменные кэша, в которые будет выполняться запись. Это выглядит следующим образом.
Ayu = zeros(N, N)
uAx = zeros(N, N)
Du = zeros(N, N)
Ayv = zeros(N, N)
vAx = zeros(N, N)
Dv = zeros(N, N)
function gm3!(dr, r, p, t)
a, α, ubar, β, D1, D2 = p
u = @view r[:, :, 1]
v = @view r[:, :, 2]
du = @view dr[:, :, 1]
dv = @view dr[:, :, 2]
mul!(Ayu, Ay, u)
mul!(uAx, u, Ax)
mul!(Ayv, Ay, v)
mul!(vAx, v, Ax)
@. Du = D1 * (Ayu + uAx)
@. Dv = D2 * (Ayv + vAx)
@. du = Du + a * u * u ./ v + ubar - α * u
@. dv = Dv + a * u * u - β * v
end
prob = ODEProblem(gm3!, r0, (0.0, 0.1), p)
@btime solve(prob, Tsit5()); 46.835 ms (4703 allocations: 29.97 MiB)Но временные переменные являются глобальными. Необходимо либо объявить кэши как const, либо локализовать их. Их можно локализовать, добавив в параметры (p). Компилятору проще иметь дело с локальными переменными, чем с глобальными. Локализация переменных помогает обеспечить стабильность типов.
p = (1.0, 1.0, 1.0, 10.0, 0.001, 100.0, Ayu, uAx, Du, Ayv, vAx, Dv) # a, α, ubar, β, D1, D2
function gm4!(dr, r, p, t)
a, α, ubar, β, D1, D2, Ayu, uAx, Du, Ayv, vAx, Dv = p
u = @view r[:, :, 1]
v = @view r[:, :, 2]
du = @view dr[:, :, 1]
dv = @view dr[:, :, 2]
mul!(Ayu, Ay, u)
mul!(uAx, u, Ax)
mul!(Ayv, Ay, v)
mul!(vAx, v, Ax)
@. Du = D1 * (Ayu + uAx)
@. Dv = D2 * (Ayv + vAx)
@. du = Du + a * u * u ./ v + ubar - α * u
@. dv = Dv + a * u * u - β * v
end
prob = ODEProblem(gm4!, r0, (0.0, 0.1), p)
@btime solve(prob, Tsit5()); 36.810 ms (1028 allocations: 29.66 MiB)Мы могли бы использовать BLAS gemmv для дополнительной оптимизации умножения матриц, но вместо этого проведем девекторизацию шаблона.
p = (1.0, 1.0, 1.0, 10.0, 0.001, 100.0, N)
function fast_gm!(du, u, p, t)
a, α, ubar, β, D1, D2, N = p
@inbounds for j in 2:(N - 1), i in 2:(N - 1)
du[i, j, 1] = D1 *
(u[i - 1, j, 1] + u[i + 1, j, 1] + u[i, j + 1, 1] + u[i, j - 1, 1] -
4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
end
@inbounds for j in 2:(N - 1), i in 2:(N - 1)
du[i, j, 2] = D2 *
(u[i - 1, j, 2] + u[i + 1, j, 2] + u[i, j + 1, 2] + u[i, j - 1, 2] -
4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
end
@inbounds for j in 2:(N - 1)
i = 1
du[1, j, 1] = D1 *
(2u[i + 1, j, 1] + u[i, j + 1, 1] + u[i, j - 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
end
@inbounds for j in 2:(N - 1)
i = 1
du[1, j, 2] = D2 *
(2u[i + 1, j, 2] + u[i, j + 1, 2] + u[i, j - 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
end
@inbounds for j in 2:(N - 1)
i = N
du[end, j, 1] = D1 *
(2u[i - 1, j, 1] + u[i, j + 1, 1] + u[i, j - 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
end
@inbounds for j in 2:(N - 1)
i = N
du[end, j, 2] = D2 *
(2u[i - 1, j, 2] + u[i, j + 1, 2] + u[i, j - 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
end
@inbounds for i in 2:(N - 1)
j = 1
du[i, 1, 1] = D1 *
(u[i - 1, j, 1] + u[i + 1, j, 1] + 2u[i, j + 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
end
@inbounds for i in 2:(N - 1)
j = 1
du[i, 1, 2] = D2 *
(u[i - 1, j, 2] + u[i + 1, j, 2] + 2u[i, j + 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
end
@inbounds for i in 2:(N - 1)
j = N
du[i, end, 1] = D1 *
(u[i - 1, j, 1] + u[i + 1, j, 1] + 2u[i, j - 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
end
@inbounds for i in 2:(N - 1)
j = N
du[i, end, 2] = D2 *
(u[i - 1, j, 2] + u[i + 1, j, 2] + 2u[i, j - 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
end
@inbounds begin
i = 1
j = 1
du[1, 1, 1] = D1 * (2u[i + 1, j, 1] + 2u[i, j + 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
du[1, 1, 2] = D2 * (2u[i + 1, j, 2] + 2u[i, j + 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
i = 1
j = N
du[1, N, 1] = D1 * (2u[i + 1, j, 1] + 2u[i, j - 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
du[1, N, 2] = D2 * (2u[i + 1, j, 2] + 2u[i, j - 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
i = N
j = 1
du[N, 1, 1] = D1 * (2u[i - 1, j, 1] + 2u[i, j + 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
du[N, 1, 2] = D2 * (2u[i - 1, j, 2] + 2u[i, j + 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
i = N
j = N
du[end, end, 1] = D1 * (2u[i - 1, j, 1] + 2u[i, j - 1, 1] - 4u[i, j, 1]) +
a * u[i, j, 1]^2 / u[i, j, 2] + ubar - α * u[i, j, 1]
du[end, end, 2] = D2 * (2u[i - 1, j, 2] + 2u[i, j - 1, 2] - 4u[i, j, 2]) +
a * u[i, j, 1]^2 - β * u[i, j, 2]
end
end
prob = ODEProblem(fast_gm!, r0, (0.0, 0.1), p)
@btime solve(prob, Tsit5()); 7.760 ms (440 allocations: 29.62 MiB)Обратите внимание, что в данном случае слияние циклов и отказ от линейных операторов приводит к значительному улучшению — примерно в 10 раз. Это на порядок быстрее, чем наш исходный код в векторном стиле на MATLAB/SciPy/R.
Поскольку делать это вручную утомительно, отметим, что генерация символьного кода ModelingToolkit.jl выполняет эти действия автоматически, начиная с базовой версии.
using ModelingToolkit
function basic_version!(dr, r, p, t)
a, α, ubar, β, D1, D2 = p
u = r[:, :, 1]
v = r[:, :, 2]
Du = D1 * (Ay * u + u * Ax)
Dv = D2 * (Ay * v + v * Ax)
dr[:, :, 1] = Du .+ a .* u .* u ./ v .+ ubar .- α * u
dr[:, :, 2] = Dv .+ a .* u .* u .- β * v
end
a, α, ubar, β, D1, D2 = p
uss = (ubar + β) / α
vss = (a / β) * uss^2
r0 = zeros(100, 100, 2)
r0[:, :, 1] .= uss .+ 0.1 .* rand.()
r0[:, :, 2] .= vss
prob = ODEProblem(basic_version!, r0, (0.0, 0.1), p)
de = modelingtoolkitize(prob)
# Примечание. jac=true,sparse=true приводит к автоматическому построению кода
# разреженного якобиана.
fastprob = ODEProblem(de, [], (0.0, 0.1), jac = true, sparse = true)ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 0.1)
u0: 20000-element Vector{Float64}:
11.074064020465531
11.04322113492956
11.015062072513434
11.008939655202521
11.064008868363828
11.01104464637149
11.011993387604734
11.092503672819479
11.076804445317938
11.036016752519918
⋮
12.100000000000001
12.100000000000001
12.100000000000001
12.100000000000001
12.100000000000001
12.100000000000001
12.100000000000001
12.100000000000001
12.100000000000001Наконец, мы можем делать и другое, например применять многопоточность для основных циклов. LoopVectorization.jl предоставляет макрос @turbo для реализации множества улучшений SIMD, а @tturbo является многопоточной версией.
Оптимизация выбора алгоритма
Последнее, что остается сделать, это оптимизировать выбор алгоритма. В качестве тестового алгоритма мы использовали Tsit5(), но в действительности данная задача является дискретизацией жесткого уравнения PDE, поэтому рекомендуется использовать CVODE_BDF(). Однако вместо плотного якобиана, используемого по умолчанию, следует воспользоваться разреженным якобианом, который предоставляет данная задача. Якобиан представляет собой матрицу , где считается по линейному индексу (т. е. вниз по столбцам). Но поскольку переменные зависят от , размер полосы здесь велик, и, следовательно, это не будет иметь успеха с решателем ленточного якобиана. Вместо этого мы используем алгоритмы с разреженными якобианами. CVODE_BDF позволяет использовать разреженный решатель Ньютона-Крылова. Для этого требуется задать linear_solver = :GMRES.
|
Более подробные сведения можно найти в руководстве Решение больших жестких уравнений. А здесь приведена лишь незначительная часть возможностей для оптимизации. |
Посмотрим, как масштабируется наша быстрая правая часть при увеличении времени интеграции.
prob = ODEProblem(fast_gm!, r0, (0.0, 10.0), p)
@btime solve(prob, Tsit5()); 108.703 s (39331 allocations: 2.76 GiB)using Sundials
@btime solve(prob, CVODE_BDF(linear_solver = :GMRES)); 906.420 ms (13479 allocations: 120.05 MiB)prob = ODEProblem(fast_gm!, r0, (0.0, 100.0), p)
# Если не отключить параметр `save_everystep`, память будет исчерпана.
@btime solve(prob, Tsit5(), save_everystep = false); 6.879 s (68 allocations: 2.90 MiB)@btime solve(prob, CVODE_BDF(linear_solver = :GMRES), save_everystep = false); 2.942 s (35152 allocations: 2.46 MiB)prob = ODEProblem(fast_gm!, r0, (0.0, 500.0), p)
@btime solve(prob, CVODE_BDF(linear_solver = :GMRES), save_everystep = false); 4.253 s (51451 allocations: 3.17 MiB)Обратите внимание, что мы избавились практически от всех выделений, что позволяет оптимизировать код без сборки мусора и замедления производительности
Почему решатель CVODE_BDF так хорошо работает? Это происходит потому, что из-за жесткости задачи быстро растет количество шагов, требуемых явному методу Рунге-Кутта, в то время как CVODE_BDF делает большие шаги. Кроме того, форма линейного решателя GMRES является достаточно эффективным способом решения неявной системы в данном случае. Здесь действует зависимость от конкретной задачи, и во многих случаях использование метода Крылова требует предобусловливателя, поэтому вам нужно протестировать другие алгоритмы и линейные решатели, чтобы найти подходящий вариант для вашей задачи.
Перейдем к руководству Решение больших жестких уравнений, где содержатся более подробные сведения об оптимизации выбора алгоритма для таких кодов.