Description
Is your feature request related to a problem? Please describe.
I'm trying to use IMEX to solve a reaction-advection system, where the reaction is the stiff part and the advection is the non-stiff part.
The issue is described in detail here.
Describe the solution you’d like
The state of the system is a S x N matrix, where S is the number of chemical species and N is the number of grid cells. While solving the stiff part, each column of the state matrix is independent of all the other columns (because the reaction happens within an individual grid cell), so the Jacobian is a Block diagonal matrix.
Because all the grid cells are independent, the solve time should theoretically scale linearly with number of grid cells. However linear solve part of it doesn't scale linearly:
using LinearSolve
using BlockBandedMatrices
using ProgressLogging
using LinearAlgebra
using Plots
ngrid = Int.(round.((10.0.^collect(0.5:0.2:4))./2)) .* 2
ts = []
@progress for N in ngrid
B = rand(3 * N)
b = repeat([3], N)
x = rand(N*3, N*3)
A = BlockBandedMatrix{eltype(x)}(x, b, b, (0, 0))
prob = LinearProblem(A, B)
t = @elapsed solve(prob, LUFactorization())
push!(ts, t)
end
plot(ngrid, ts, xaxis=:log, yaxis=:log, xlabel="N", legend=:topleft,
label="Block Diagonal Matrix", ylabel="Time (s)")
plot!(ngrid, ts[3] .* ngrid ./ ngrid[3],
label="Linear Scaling")
The red line is linear scaling, but as you can see the linear solve with the block banded matrix is more like quadratic.
So I like like it to scale linearly, and just in general be as fast as possible in this case.
Describe alternatives you’ve considered
If we use a BandedMatrix instead of a BlockBandedMatrix, we do get linear scaling, which is the green line in the figure above:
using BandedMatrices
ts = []
@progress for N in ngrid
B = rand(3 * N)
b = repeat([3], N)
x = rand(N*3, N*3)
A = BandedMatrix{eltype(x)}(x, (2, 2))
prob = LinearProblem(A, B)
t = @elapsed solve(prob, LUFactorization())
push!(ts, t)
end
plot!(ngrid, ts, label="Banded Matrix")However, the extra zeros in the banded matrix mean that something like 2x the amount of computations are being done compared to what is strictly necessary. This may be significant when the number of chemical species S is large.
Additional context
I'm told there is a GPU kernel for this here.