# Solving Stiff Equations

##### Chris Rackauckas

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.

## Code Optimization for Differential Equations

### Writing Efficient Code

For a detailed tutorial on how to optimize one's DifferentialEquations.jl code, please see the Optimizing DiffEq Code tutorial.

### Choosing a Good Solver

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.

### Check Out the Speed FAQ

See this FAQ for information on common pitfalls and how to improve performance.

### Setting Up Your Julia Installation for Speed

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

8


If I want to set this directly to 4 threads, I would use:

using LinearAlgebra


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.

### Use Accelerator Hardware

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.

## Speeding Up Jacobian Calculations

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.

### Declaring Jacobian Functions

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]


### Automatic Derivation of Jacobian Functions

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


### Declaring a Sparse Jacobian

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


### Automatic Sparsity Detection

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)