This tutorial is for getting into the extra features for solving stiff ordinary differential equations in an efficient manner. Solving stiff ordinary differential equations requires specializing the linear solver on properties of the Jacobian in order to cut down on the O(n^3) linear solve and the O(n^2) back-solves. Note that these same functions and controls also extend to stiff SDEs, DDEs, DAEs, etc.
For a detailed tutorial on how to optimize one's DifferentialEquations.jl code, please see the Optimizing DiffEq Code tutorial.
Choosing a good solver is required for getting top notch speed. General recommendations can be found on the solver page (for example, the ODE Solver Recommendations). The current recommendations can be simplified to a Rosenbrock method (Rosenbrock23
or Rodas5
) for smaller (<50 ODEs) problems, ESDIRK methods for slightly larger (TRBDF2
or KenCarp4
for <2000 ODEs), and Sundials CVODE_BDF
for even larger problems. lsoda
from LSODA.jl is generally worth a try.
More details on the solver to choose can be found by benchmarking. See the DiffEqBenchmarks to compare many solvers on many problems.
See this FAQ for information on common pitfalls and how to improve performance.
Julia uses an underlying BLAS implementation for its matrix multiplications and factorizations. This library is automatically multithreaded and accelerates the internal linear algebra of DifferentialEquations.jl. However, for optimality, you should make sure that the number of BLAS threads that you are using matches the number of physical cores and not the number of logical cores. See this issue for more details.
To check the number of BLAS threads, use:
using LinearAlgebra LinearAlgebra.BLAS.get_num_threads()
8
If I want to set this directly to 4 threads, I would use:
using LinearAlgebra LinearAlgebra.BLAS.set_num_threads(4)
Additionally, in some cases Intel's MKL might be a faster BLAS than the standard BLAS that ships with Julia (OpenBLAS). To switch your BLAS implementation, you can use MKL.jl which will accelerate the linear algebra routines. Please see the package for the limitations.
When possible, use GPUs. If your ODE system is small and you need to solve it with very many different parameters, see the ensembles interface and DiffEqGPU.jl. If your problem is large, consider using a CuArray for the state to allow for GPU-parallelism of the internal linear algebra.
When one is using an implicit or semi-implicit differential equation solver, the Jacobian must be built at many iterations and this can be one of the most expensive steps. There are two pieces that must be optimized in order to reach maximal efficiency when solving stiff equations: the sparsity pattern and the construction of the Jacobian. The construction is filling the matrix J
with values, while the sparsity pattern is what J
to use.
The sparsity pattern is given by a prototype matrix, the jac_prototype
, which will be copied to be used as J
. The default is for J
to be a Matrix
, i.e. a dense matrix. However, if you know the sparsity of your problem, then you can pass a different matrix type. For example, a SparseMatrixCSC
will give a sparse matrix. Additionally, structured matrix types like Tridiagonal
, BandedMatrix
(from BandedMatrices.jl), BlockBandedMatrix
(from BlockBandedMatrices.jl), and more can be given. DifferentialEquations.jl will internally use this matrix type, making the factorizations faster by utilizing the specialized forms.
For the construction, there are 3 ways to fill J
:
The default, which uses normal finite/automatic differentiation
A function jac(J,u,p,t)
which directly computes the values of J
A colorvec
which defines a sparse differentiation scheme.
We will now showcase how to make use of this functionality with growing complexity.
Let's solve the Rosenbrock equations:
\[ \begin{align} dy_1 &= -0.04y₁ + 10^4 y_2 y_3 \\ dy_2 &= 0.04 y_1 - 10^4 y_2 y_3 - 3*10^7 y_{2}^2 \\ dy_3 &= 3*10^7 y_{3}^2 \\ \end{align} \]
In order to reduce the Jacobian construction cost, one can describe a Jacobian function by using the jac
argument for the ODEFunction
. First, let's do a standard ODEProblem
:
using DifferentialEquations 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,Rosenbrock23()) using Plots plot(sol, xscale=:log10, tspan=(1e-6, 1e5), layout=(3,1))
using BenchmarkTools @btime solve(prob)
250.618 μs (2563 allocations: 186.94 KiB) retcode: Success Interpolation: automatic order switching interpolation t: 115-element Vector{Float64}: 0.0 0.0014148468219250373 0.0020449182545311173 0.0031082402716566307 0.004077787050059496 0.005515332443361059 0.007190040962774541 0.009125372578778032 0.011053912492732977 0.012779077276958607 ⋮ 47335.55742427336 52732.00629853751 58693.72275675742 65277.99247104326 72548.193682209 80574.55524404174 89435.04313420167 99216.40190401232 100000.0 u: 115-element Vector{Vector{Float64}}: [1.0, 0.0, 0.0] [0.9999434113193613, 3.283958829839966e-5, 2.374909234028646e-5] [0.9999182177783585, 3.5542680136344576e-5, 4.6239541505020636e-5] [0.999875715036629, 3.630246933484973e-5, 8.798249403609502e-5] [0.9998369766077329, 3.646280308115454e-5, 0.00012656058918590176] [0.9997795672444667, 3.6466430856422276e-5, 0.00018396632467683696] [0.9997127287139348, 3.644727999289594e-5, 0.0002508240060722832] [0.9996355450022019, 3.6366816179962554e-5, 0.0003280881816181881] [0.9995586925734838, 3.6018927453310745e-5, 0.00040528849906290245] [0.9994899965196854, 3.468694637784628e-5, 0.00047531653393682193] ⋮ [0.03394368533138813, 1.4047985954839008e-7, 0.9660561741887524] [0.031028978802635446, 1.280360882615162e-7, 0.9689708931612764] [0.02835436649772506, 1.1668210763639173e-7, 0.971645516820168] [0.02590132862868119, 1.0632277796078e-7, 0.9740985650485414] [0.023652547707489525, 9.687113505658095e-8, 0.9763473554213756] [0.021591864255513585, 8.824768851310993e-8, 0.9784080474967981] [0.01970422745458613, 8.037977845291491e-8, 0.9802956921656356] [0.017975643191251073, 7.320098956514714e-8, 0.9820242836077591] [0.01785056623499974, 7.26838436117136e-8, 0.9821493610811564]
Now we want to add the Jacobian. First we have to derive the Jacobian $\frac{df_i}{du_j}$ which is J[i,j]
. From this we get:
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)) @btime solve(prob_jac)
193.828 μs (2002 allocations: 126.28 KiB) retcode: Success Interpolation: automatic order switching interpolation t: 115-element Vector{Float64}: 0.0 0.0014148468219250373 0.0020449182545311173 0.0031082402716566307 0.004077787050059496 0.005515332443361059 0.007190040962774541 0.009125372578778032 0.011053912492732977 0.012779077276958607 ⋮ 45964.060340548356 51219.40381376205 57025.01899700374 63436.021374561584 70513.1073617524 78323.14229130604 86939.82338876331 96444.41085674686 100000.0 u: 115-element Vector{Vector{Float64}}: [1.0, 0.0, 0.0] [0.9999434113193613, 3.283958829839966e-5, 2.374909234028646e-5] [0.9999182177783585, 3.5542680136344576e-5, 4.6239541505020636e-5] [0.999875715036629, 3.630246933484973e-5, 8.798249403609502e-5] [0.9998369766077329, 3.646280308115454e-5, 0.00012656058918590176] [0.9997795672444667, 3.6466430856422276e-5, 0.00018396632467683696] [0.9997127287139348, 3.644727999289594e-5, 0.0002508240060722832] [0.9996355450022019, 3.6366816179962554e-5, 0.0003280881816181881] [0.9995586925734838, 3.6018927453310745e-5, 0.00040528849906290245] [0.9994899965196854, 3.468694637784628e-5, 0.00047531653393682193] ⋮ [0.03478048133177493, 1.4406682005231008e-7, 0.9652193746014031] [0.03179591062189176, 1.313038656880417e-7, 0.9682039580742408] [0.029057356622057315, 1.1966100432939363e-7, 0.9709425237169371] [0.02654597011713668, 1.0904070990251299e-7, 0.9734539208421517] [0.024244118287194777, 9.935385522693504e-8, 0.9757557823589477] [0.022135344621501105, 9.05190025093182e-8, 0.9778645648594945] [0.02020432071854, 8.246174295748071e-8, 0.9797955968197154] [0.018436796681356796, 7.511410189106845e-8, 0.9815631282045397] [0.01785426048218692, 7.269900678199638e-8, 0.9821456668188047]
But that was hard! If you want to take the symbolic Jacobian of numerical code, we can make use of ModelingToolkit.jl to symbolicify the numerical code and do the symbolic calculation and return the Julia code for this.
using ModelingToolkit de = modelingtoolkitize(prob) ModelingToolkit.generate_jacobian(de)[2] # Second is in-place
:(function (var"##out#1817", var"##arg#1815", var"##arg#1816", t) #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-28aa5 7c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:282 =# #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-28aa5 7c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:283 =# let var"x₁(t)" = #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78 b-41ea-bc97-28aa57c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:169 =# @i nbounds(var"##arg#1815"[1]), var"x₂(t)" = #= /cache/julia-buildkite-plugin/ depots/a6029d3a-f78b-41ea-bc97-28aa57c6c6ea/packages/SymbolicUtils/9iQGH/sr c/code.jl:169 =# @inbounds(var"##arg#1815"[2]), var"x₃(t)" = #= /cache/juli a-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-28aa57c6c6ea/packages/Sym bolicUtils/9iQGH/src/code.jl:169 =# @inbounds(var"##arg#1815"[3]), α₁ = #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-28aa57c6c6ea/p ackages/SymbolicUtils/9iQGH/src/code.jl:169 =# @inbounds(var"##arg#1816"[1] ), α₂ = #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-28a a57c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:169 =# @inbounds(var"##a rg#1816"[2]), α₃ = #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41 ea-bc97-28aa57c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:169 =# @inbou nds(var"##arg#1816"[3]) #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-2 8aa57c6c6ea/packages/Symbolics/h8kPL/src/build_function.jl:331 =# #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41ea-bc97-2 8aa57c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:329 =# @inbounds begin #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41e a-bc97-28aa57c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:325 =# var"##out#1817"[1] = (*)(-1, α₁) var"##out#1817"[2] = α₁ var"##out#1817"[3] = 0 var"##out#1817"[4] = (*)(α₃, var"x₃(t)") var"##out#1817"[5] = (+)((*)(-2, α₂, var"x₂(t)"), (*)(-1, α₃, var"x₃(t)")) var"##out#1817"[6] = (*)(2, α₂, var"x₂(t)") var"##out#1817"[7] = (*)(α₃, var"x₂(t)") var"##out#1817"[8] = (*)(-1, α₃, var"x₂(t)") var"##out#1817"[9] = 0 #= /cache/julia-buildkite-plugin/depots/a6029d3a-f78b-41e a-bc97-28aa57c6c6ea/packages/SymbolicUtils/9iQGH/src/code.jl:327 =# nothing end end end)
which outputs:
:((##MTIIPVar#376, u, p, t)->begin #= C:\Users\accou\.julia\packages\ModelingToolkit\czHtj\src\utils.jl:65 =# #= C:\Users\accou\.julia\packages\ModelingToolkit\czHtj\src\utils.jl:66 =# let (x₁, x₂, x₃, α₁, α₂, α₃) = (u[1], u[2], u[3], p[1], p[2], p[3]) ##MTIIPVar#376[1] = α₁ * -1 ##MTIIPVar#376[2] = α₁ ##MTIIPVar#376[3] = 0 ##MTIIPVar#376[4] = x₃ * α₃ ##MTIIPVar#376[5] = x₂ * α₂ * -2 + x₃ * α₃ * -1 ##MTIIPVar#376[6] = x₂ * 2 * α₂ ##MTIIPVar#376[7] = α₃ * x₂ ##MTIIPVar#376[8] = α₃ * x₂ * -1 ##MTIIPVar#376[9] = 0 end #= C:\Users\accou\.julia\packages\ModelingToolkit\czHtj\src\utils.jl:67 =# nothing end)
Now let's use that to give the analytical solution Jacobian:
jac = eval(ModelingToolkit.generate_jacobian(de)[2]) f = ODEFunction(rober, jac=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
Jacobian sparsity is declared by the jac_prototype
argument in the ODEFunction
. Note that you should only do this if the sparsity is high, for example, 0.1% of the matrix is non-zeros, otherwise the overhead of sparse matrices can be higher than the gains from sparse differentiation!
But as a demonstration, let's build a sparse matrix for the Rober problem. We can do this by gathering the I
and J
pairs for the non-zero components, like:
I = [1,2,1,2,3,1,2] J = [1,1,2,2,2,3,3] using SparseArrays jac_prototype = sparse(I,J,1.0)
3×3 SparseArrays.SparseMatrixCSC{Float64, Int64} with 7 stored entries: 1.0 1.0 1.0 1.0 1.0 1.0 ⋅ 1.0 ⋅
Now this is the sparse matrix prototype that we want to use in our solver, which we then pass like:
f = ODEFunction(rober, jac=jac, jac_prototype=jac_prototype) 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
One of the useful companion tools for DifferentialEquations.jl is SparsityDetection.jl. This allows for automatic declaration of Jacobian sparsity types. To see this in action, let's look at the 2-dimensional Brusselator equation:
const N = 32 const xyd_brusselator = range(0,stop=1,length=N) brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5. limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a function brusselator_2d_loop(du, u, p, t) A, B, alpha, dx = p alpha = alpha/dx^2 @inbounds for I in CartesianIndices((N, N)) i, j = Tuple(I) x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]] ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N) du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) + B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t) du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) + A*u[i,j,1] - u[i,j,1]^2*u[i,j,2] end end p = (3.4, 1., 10., step(xyd_brusselator))
(3.4, 1.0, 10.0, 0.03225806451612903)
Given this setup, we can give and example input
and output
and call sparsity!
on our function with the example arguments and it will kick out a sparse matrix with our pattern, that we can turn into our jac_prototype
.
using SparsityDetection, SparseArrays input = rand(32,32,2) output = similar(input) sparsity_pattern = jacobian_sparsity(brusselator_2d_loop,output,input,p,0.0) jac_sparsity = Float64.(sparse(sparsity_pattern))
Explored path: SparsityDetection.Path(Bool[], 1) 2048×2048 SparseArrays.SparseMatrixCSC{Float64, Int64} with 12288 stored en tries: ⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠀⠀ ⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄ ⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙ ⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀⠀⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷⣄⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣷
Let's double check what our sparsity pattern looks like:
using Plots spy(jac_sparsity,markersize=1,colorbar=false,color=:deep)