Method for solving advection-reaction equations

[ctessum]

Hi,

The code below implements an example of a simplified advection-reaction system, meant to represent an air pollution model. (A larger-scale version is here.)

using OrdinaryDiffEq, ModelingToolkit
t = ModelingToolkit.t_nounits
D = ModelingToolkit.D_nounits
using EarthSciMLBase
using DomainSets, MethodOfLines
using Plots
using ForwardDiff
using LinearAlgebra, Statistics

x_min = y_min = t_min = 0.0
x_max = y_max = 1.0
t_max = 4.0

@parameters begin
    x
    y
end

islocation(x, y) = x > x_max / 2 - dx && x < x_max / 2 + dx &&
                   y > y_max / 10 - dx &&
                   y < y_max / 10 + dx
emission(x, y) = ifelse(islocation(x, y), 10, 0)
@register_symbolic emission(x, y)

function RoberPDE(; advection=true)
    @parameters begin
        k₁ = 0.04
        k₂ = 3e7
        k₃ = 1e4
    end
    @variables begin
        y₁(..)
        y₂(..)
        y₃(..)
        θ(..)
        u(..)
        v(..)
    end
    function advect(var)
        if advection
            -u(x, y, t) * Differential(x)(var) - v(x, y, t) * Differential(y)(var)
        else
            0.0
        end
    end
    eqs = [
        θ(x, y, t) ~ atan(y .- 0.5, x .- 0.5)
        u(x, y, t) ~ -sin(θ(x, y, t))
        v(x, y, t) ~ cos(θ(x, y, t))
        D(y₁(x, y, t)) ~ advect(y₁(x, y, t)) - k₁ * y₁(x, y, t) +
                         k₃ * y₂(x, y, t) * y₃(x, y, t) + emission(x, y)
        D(y₂(x, y, t)) ~ advect(y₂(x, y, t)) + k₁ * y₁(x, y, t) -
                         k₂ * y₂(x, y, t)^2 - k₃ * y₂(x, y, t) * y₃(x, y, t)
        D(y₃(x, y, t)) ~ advect(y₃(x, y, t)) + k₂ * y₂(x, y, t)^2
    ]
    domain = [
        x ∈ Interval(x_min, x_max),
        y ∈ Interval(y_min, y_max),
        t ∈ Interval(t_min, t_max),
    ]
    bcs = vcat([[u(x, y, t_min) ~ 0.0,
        u(x_min, y, t) ~ 0,
        u(x_max, y, t) ~ 0,
        u(x, y_min, t) ~ 0,
        u(x, y_max, t) ~ 0] for u in (y₁, y₂, y₃)]...)

    PDESystem(eqs, bcs, domain, [x, y, t], [y₁(x, y, t), y₂(x, y, t), y₃(x, y, t),
            θ(x, y, t), u(x, y, t), v(x, y, t)], [k₁, k₂, k₃]; name=:rober_pde)
end


rober_pde = RoberPDE()
rober_pde_noadvection = RoberPDE(advection=false)

N = 10 # Ideally this would be a larger number.
dx = (x_max - x_min) / N
dy = (y_max - y_min) / N

discretization = MOLFiniteDifference([x => dx, y => dy], t)
prob = discretize(rober_pde, discretization)
prob_noadvect = discretize(rober_pde_noadvection, discretization)

@assert all(isequal.(unknowns(prob.f.sys), unknowns(prob_noadvect.f.sys)))

sol = solve(prob, Rosenbrock23())

The code above implements a version of the Robertson equations on a 2D grid with (optional) advection driven by wind traveling in a circle and emissions forcing of variable y1 at a specific location. The result looks like this:

y1 = only(filter(k -> occursin("y₁", string(Symbol(k))), keys(sol.u)))
y2 = only(filter(k -> occursin("y₂", string(Symbol(k))), keys(sol.u)))
y3 = only(filter(k -> occursin("y₃", string(Symbol(k))), keys(sol.u)))

anim = @animate for i in eachindex(sol[t])
    plot(
        heatmap(sol[x], sol[y], sol[y1][:, :, i], xlabel="x", ylabel="y", title="y₁ Concentration",
            clim=(0, maximum(sol[y1]))),
        heatmap(sol[x], sol[y], sol[y2][:, :, i], title="y₂ (t = $(round(sol.t[i], sigdigits=2)))",
            clim=(0, maximum(sol[y2]))),
        heatmap(sol[x], sol[y], sol[y3][:, :, i], title="y₃",
            clim=(0, maximum(sol[y3]))),
        size=(1000, 800)
    )
end
gif(anim, fps=8)

I think for this type of system, the current (and perhaps near-future) MTK solution approach may be slower than what is used in the current generation of these types of models (i.e. written in fortran).

One issue is that there are (to a first approximation) two parts of the system, represented by the robertson and advection equations in the system above. The robertson-like equations are usually solved implicitly (and separately for each grid cell) and the advection equations are solved explicitly, and they are recombined using strang splitting.

We've implemented a version of this here, and it works well enough, but it doesn't integrate well with the rest of the SciML ecosystem.

Something similar can be done with IMEX solvers, which we have implemented here. This doesn't work at all unless one takes advantage of the block-diagonal sparsity pattern of the Jacobian, which we handle here.

However, even after taking advantage of the jacobian sparsity, the IMEX solution method is still much slower than our self-implemented Strang method. Assuming it's not just some implementation mistake on our part, I think it might have to do with spatial variability in the stiffness of the system:

function getvalfromsol(sol, v, t_index)
    vv, i, j = getfield(v, 5)
    if string(vv) == "y₃(t)"
        return sol[y3][i, j, t_index]
    elseif string(vv) == "y₂(t)"
        return sol[y2][i, j, t_index]
    elseif string(vv) == "y₁(t)"
        return sol[y1][i, j, t_index]
    else
        error(vv)
    end
end

function jac2array(jac)
    o = zeros(3, 3, N - 1, N - 1)
    for (rr, v) in enumerate(unknowns(prob.f.sys))
        vv, i, j = getfield(v, 5)
        if string(vv) == "y₃(t)"
            k1 = 3
        elseif string(vv) == "y₂(t)"
            k1 = 2
        elseif string(vv) == "y₁(t)"
            k1 = 1
        else
            error(vv)
        end
        for (cc, w) in enumerate(unknowns(prob.f.sys))
            ww, m, n = getfield(w, 5)
            (i != m || j != n) && continue # only diagonal blocks
            if string(ww) == "y₃(t)"
                k2 = 3
            elseif string(ww) == "y₂(t)"
                k2 = 2
            elseif string(ww) == "y₁(t)"
                k2 = 1
            else
                error(ww)
            end
            o[k1, k2, i-1, j-1] = jac[rr, cc]
        end
    end
    o
end

anim = @animate for (i, _t) in enumerate(sol.t)
    u = [getvalfromsol(sol, v, i) for v in unknowns(prob.f.sys)]
    jac = ForwardDiff.jacobian((u) -> prob_noadvect.f(u, prob.p, _t), u)
    jac_a = jac2array(jac)
    jac_eig = [eigvals(jac_a[:, :, xi, yi]) for xi in 1:N-1, yi in 1:N-1]

    jac_eig_i(el, i) = log10(abs(el[i]))
    plot(
        heatmap(jac_eig_i.(jac_eig, (1,)), xlabel="x", ylabel="y", title="y₁ log10 Jacobian eigvals"),
        heatmap(jac_eig_i.(jac_eig, (2,)), title="y₂ (t = $(round(_t, sigdigits=2)))"),
        heatmap(jac_eig_i.(jac_eig, (3,)), title="y₃");
    )
end
gif(anim, fps=8)

The figure above shows that the eigenvalues of the jacobian for this simulation can vary by 4 orders of magnitude across space, especially at the beginning of the simulation. I think this kind of thing also happens in "real" environmental dynamics, where reactions are much slower in some locations than others. Therefore the speed descrepancy may be happening because the Strang splitting method allows the solver to take larger time steps in some locations than others, whereas with the IMEX method all locations are integrated together in lockstep, with the timestep determined by the requirements of the stiffest grid cell.

Anyway, I thought I would post this here as a starting point for a discussion about a roadmap to be able to solve this type of system in MTK without needing external workarounds, and also to see if there are existing solutions that could maybe be leveraged here, such as https://github.com/SciML/OrdinaryDiffEqOperatorSplitting.jl.

Thanks