I have a system that I would like to solve which is basically a reaction advection system.
The full-size version is here, but the sake of this issue, let’s just say that we have a system state that consists of a number of grid cells and a number of chemical species in each grid cell.
There are reactions between the species in each grid cell, and also advection of other operations that happen between the grid cells for each species.
For the sake of this issue, it doesn’t actually matter what the advection or spatial operators are, it only matters that they exist.
For the reaction operator, let’s use the Robertson system:
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])
Remember that we also have grid cells and a (non-stiff) spatial operator, so we will want to eventually use an IMEX method.
Therefore, we’re restricted to the solvers that are compatible with IMEX methods.
sol = solve(prob, KenCarp47())
plot(sol, tspan = (1e-2, 1e5), xscale = :log10, legend = :topleft,
xlabel="Number of grid cells", ylabel="Time to solve")
Now let’s account for the fact that we need to integrate this system for all of the grid cells:
using Test
function solve_grid(u, p)
prob, solver = p
N = size(u, 2)
o = similar(u)
for i in 1:N
prob = remake(prob, u0=view(u, :, i))
sol = solve(prob, solver, saveat = 1e5)
o[:, i] .= sol[end]
end
o
end
u0 = repeat([1.0 0.5
0.0 0.0
0.0 0.0], 1, 5)
sol_grid = solve_grid(u0, [prob, KenCarp47()])
# Check for correctness
for i in 1:2:size(u0, 2)
@test sol_grid[:, i] ≈ sol[end]
end
Now, let’s check how long it takes to solve as a function of the number of grid cells:
using ProgressLogging
function create_u0(N)
u = zeros(3, N)
u[1,:] = rand(N)
u
end
ngrid = Int.(round.((10.0.^collect(0.5:0.2:3.6))./2)) .* 2
u0s = create_u0.(ngrid)
ts = []
@progress for (i, N) in enumerate(ngrid)
t = @elapsed solve_grid(u0s[i], [prob, KenCarp47()])
push!(ts, t)
end
plot(ngrid, ts, xscale = :log10, yscale = :log10, label = "Independent solves",
legend = :topleft, xlabel="Number of grid cells", ylabel="Time to solve")
plot!(ngrid, ngrid .* ts[1] ./ ngrid[1], xscale = :log10, yscale = :log10, label = "Linear scaling")
As we would expect, there is a linear relationship between the number of grid cells and the time it takes to solve the system.
However, recall that we also have a spatial operator, so we will need to use an IMEX method.
IMEX methods require our system to be in the form du = f(u, p, t) + g(u, p, t), where f is the stiff part and g is the non-stiff part.
Let’s reorganize it in that manner, still just considering the stiff part for now:
function rober_grid!(du, u, p, t)
u = reshape(u, 3, :)
du = reshape(du, 3, :)
N = size(u, 2)
prob = p[1]
for i in 1:N
prob.f(view(du, :, i), view(u, :, i), prob.p, t)
end
end
u0 = repeat([1.0, 0.0, 0.0], 1, 10)
prob2 = ODEProblem(rober_grid!, u0[:], (0.0, 1e5), [prob])
sol_grid2 = solve(prob2, KenCarp47())
# Check for correctness
for i in 1:size(u0, 2)
@test reshape(sol_grid2[end], 3, :)[:, i] ≈ sol[end]
end
Now, let’s see how the time scaling with this method compares to the independent solving method:
ts = []
@progress for (i, N) in enumerate(ngrid)
prob2 = ODEProblem(rober_grid!, u0s[i][:], (0.0, 1e5), [prob])
t = @elapsed solve(prob2, KenCarp47())
push!(ts, t)
end
plot!(ngrid, ts, xscale = :log10, yscale = :log10, label = "Block solves")
Let’s check what is taking the most time in the block solve:
u0 = repeat([1.0, 0.0, 0.0], 1, 1000)
prob2 = ODEProblem(rober_grid!, u0[:], (0.0, 1e5), [prob])
@profview solve(prob2, KenCarp47())
I don’t completely understand what’s going on here, but it looks like most of the time is being spent on the linear solve.
Anyway, let’s move on to the advice in “Solving Large Stiff Equations”.
First, we can try adding a Jacobian prototype.
My understanding is that for this system, the Jacobian is a block banded matrix:
using BlockBandedMatrices
using LinearSolve
u0 = repeat([1.0, 0.0, 0.0], 1, 10)
b = repeat([3], size(u0, 2))
j = BlockBandedMatrix{eltype(u0)}(undef, b, b, (0, 0))
prob3 = ODEProblem(ODEFunction(rober_grid!, jac_prototype=j), u0[:], (0.0, 1e5), [prob])
sol_grid3 = solve(prob3, KenCarp47())
This gives the error ERROR: Special factorization not handled in current default algorithm, apparently referring to the default algorithm choice for the linear solve.
Let’s try using a different linear solver:
sol_grid3 = solve(prob3, KenCarp47(linsolve=KLUFactorization()))
This gives the error: MethodError: no method matching nonzeros(::BlockBandedMatrix{Float64}).
sol_grid3 = solve(prob3, KenCarp47(linsolve=UMFPACKFactorization()))
This gives the error: MethodError: no method matching getcolptr(::BlockBandedMatrix{Float64}).
sol_grid3 = solve(prob3, KenCarp47(linsolve=RFLUFactorization()))
This gives the error: TypeError: in setfield!, expected Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, got a value of type Tuple{LinearAlgebra.LU{Float64, BlockBandedMatrix{Float64}, Vector{Int64}}, Vector{Int64}}.
using LinearAlgebra
sol_grid3 = solve(prob3, KenCarp47(linsolve=LUFactorization(pivot=LinearAlgebra.RowMaximum())))
This works!
sol_grid3 = solve(prob3, KenCarp47(linsolve=QRFactorization()))
This gives the error: MethodError: qr!(::BlockBandedMatrix{Float64}, ::NoPivot) is ambiguous.
Since we are able to get one of the linear solvers to work, let’s see how the time scaling compares to the other methods:
ts = []
@progress for (i, N) in enumerate(ngrid)
u0 = u0s[i]
b = repeat([3], size(u0, 2))
j = BlockBandedMatrix{eltype(u0)}(undef, b, b, (0, 0))
prob2 = ODEProblem(ODEFunction(rober_grid!, jac_prototype=j), u0[:], (0.0, 1e5), [prob])
t = @elapsed solve(prob2, KenCarp47(linsolve=LUFactorization(pivot=LinearAlgebra.RowMaximum())))
push!(ts, t)
end
plot!(ngrid, ts, xscale = :log10, yscale = :log10, label = "Block banded J; LU")
Let’s keep following the stiff ODE example, now trying krylov methods:
ts = []
@progress for (i, N) in enumerate(ngrid)
u0 = u0s[i]
b = repeat([3], size(u0, 2))
j = BlockBandedMatrix{eltype(u0)}(undef, b, b, (0, 0))
prob2 = ODEProblem(ODEFunction(rober_grid!, jac_prototype=j), u0[:], (0.0, 1e5), [prob])
t = @elapsed solve(prob2, KenCarp47(linsolve=KrylovJL_GMRES()))
push!(ts, t)
end
plot!(ngrid, ts, xscale = :log10, yscale = :log10, label = "KrylovJL_GMRES")
ts = []
@progress for (i, N) in enumerate(ngrid)
u0 = u0s[i]
b = repeat([3], size(u0, 2))
j = BlockBandedMatrix{eltype(u0)}(undef, b, b, (0, 0))
prob2 = ODEProblem(ODEFunction(rober_grid!, jac_prototype=j), u0[:], (0.0, 1e5), [prob])
t = @elapsed solve(prob2, KenCarp47(linsolve=SimpleGMRES()))
push!(ts, t)
end
plot!(ngrid, ts, xscale = :log10, yscale = :log10, label = "SimpleGMRES")
(SimpleGMRES results in unstable solutions with default settings.)
Let’s keep following the stiff ODE tutorial, now adding a precoditioner:
using AlgebraicMultigrid
function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
if newW === nothing || newW
Pl = aspreconditioner(AlgebraicMultigrid.ruge_stuben(convert(AbstractMatrix, W)))
else
Pl = Plprev
end
Pl, nothing
end
ts = []
@progress for N in ngrid
u0 = repeat([1.0 0.5
0.0 0.0
0.0 0.0], 1, Int(N/2))
b = repeat([3], size(u0, 2))
j = BlockBandedMatrix{eltype(u0)}(undef, b, b, (0, 0))
prob2 = ODEProblem(ODEFunction(rober_grid!, jac_prototype=j), u0[:], (0.0, 1e5), [prob])
t = @elapsed solve(prob2,
KenCarp47(linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid,
concrete_jac = true))
push!(ts, t)
end
# plot!(ngrid, ts, xscale = :log10, yscale = :log10, label = "GMRES AMG preconditioner")
This gives the error: ERROR: MethodError: no method matching ruge_stuben(::BlockBandedMatrix{Float64}).
Results
This is the plot that is created from all of the commands above:
As you can see, the “independent solves” case scales approximately linearly, which is what we would expect. We would also hope that the other methods scale in more or less the same manner that “independent solves” does, but that is not the case. It could also be okay if they scaled linearly but just with a larger slope, but as you can see they seem to be curving upwards.
This is obviously just a toy problem, but the real problem I would like to solve has somewhere between 15 and 400 chemical species (instead of three) and around 1 million grid cells (instead of 5,000 here). And the current scaling behavior isn’t workable at that scale.
I’d greatly appreciate any suggestions regarding what we can do to make it go faster. Thanks!
