Using stiff solvers for a large system of ode

Hi,

I am trying to solve the following set of partial differential equation:

\begin{cases} \partial_t u(x, y, t) = \partial_x f(x, y, t, u, \partial_x u) \\ \partial_t v(x, y, t) = \partial_y f(x, y, t, u, \partial_x u) \end{cases}

for some known function f. I want to use the a finite discretisation scheme of the shape

\begin{cases} \partial_t u(x, y, t) = D_x^f f(x, y, t, u, D_x^d u) \\ \partial_t v(x, y, t) = D_y^f f(x, y, t, u, D_x^d u) \end{cases}

where D^f stands for forward first order finite difference derivative and D^d for downward first order finite difference derivative, with the boundary condition that D^f=D^u on the right and D^d = 0 on the left.

Using the DifferentialEquations.jl package I arrived at the following working implementation:

using DifferentialEquations

function flux(i, j, t, f, p)
    T, μ, xl, yl = p

    m2π = f[i, j, 1]

    if i == 1
        m2σ = f[i, j, 1]
    else
        m2σ = f[i, j, 1] + xl[i] * (f[i-1, j, 1] - f[i, j, 1]) / (xl[i-1] / 2 - xl[i] / 2)
    end

    # details of the function `f`
    ϵπ = sqrt(t^2 + m2π)
    ϵσ = sqrt(t^2 + m2σ)
    ϵf = sqrt(t^2 + 3.25^2 * xl[i])
    Eplus = sqrt((ϵf + μ)^2 + 2 * 3.25^2 * yl[j])
    Eminus = sqrt((ϵf - μ)^2 + 2 * 3.25^2 * yl[j])

    return -t^4 / (3π^2 * ϵf) * (
        2(ϵf + μ) / Eplus * tanh(Eplus / 2T) +
        2(ϵf - μ) / Eminus * tanh(Eminus / 2T) +
        tanh((ϵf + μ) / 2T) +
        tanh((ϵf - μ) / 2T)
    ) + t^4 / 12π^2 * (3 / ϵπ * coth(ϵπ / 2T) + 1 / ϵσ * coth(ϵσ / 2T))
end


function rhs_pde!(dfdt, f, p, t)
    for I in CartesianIndices((length(p.xl), length(p.yl)))
        i, j = Tuple(I)

        # dudt
        if i == length(p.xl)
            dfdt[i, j, 1] =
                (flux(i - 1, j, t, f, p) - flux(i, j, t, f, p)) /
                (p.xl[i-1] / 2 - p.xl[i] / 2)
        else
            dfdt[i, j, 1] =
                (flux(i + 1, j, t, f, p) - flux(i, j, t, f, p)) /
                (p.xl[i+1] / 2 - p.xl[i] / 2)
        end

        # dvdt
        if j == length(p.yl)
            dfdt[i, j, 2] =
                (flux(i, j - 1, t, f, p) - flux(i, j, t, f, p)) /
                (p.yl[j-1] / 2 - p.yl[j] / 2)
        else
            dfdt[i, j, 2] =
                (flux(i, j + 1, t, f, p) - flux(i, j, t, f, p)) /
                (p.yl[j+1] / 2 - p.yl[j] / 2)
        end
    end
end

Nx = 80
Ny = 20

xl = range(0, 170, Nx) .^ 2
yl = range(0, 100, Ny) .^ 2

u_UV = [500359 + 1 / 2 * 21.0115 * x for x in xl, y in yl]
v_UV = [2 * 8.2e5 + 40 * y for x in xl, y in yl]

f0 = zeros(Nx, Ny, 2)
f0[:, :, 1] = u_UV
f0[:, :, 2] = v_UV

tspan = (1000, 200)

params = (T = 1, μ = 0, xl = xl, yl = yl)
prob = ODEProblem(rhs_pde!, f0, tspan, params)

sol = solve(
    prob,
    Tsit5(),
    save_everystep = false,
    reltol = 1e-12,
    abstol = 1e-12,
    dt = 1e-12,
    progress = true,
    progress_steps = 500,
)

I tried to adapt my code to a minimal working example in the context that it is used, hopefully the main structure is still understandable. This give the expected result in a reasonable amount of time using non-stiff algorithm such as DP5() or Tsit5().

My problem is that the equation get stiff for the value small values of t that I want to reach, as such I would like to use a stiff algorithm, but my code seems to be extremely slow with stiff solver. I checked that my function rhs_pde!() is non allocating but I am short on idea of how to further improve the performances. Any recommendation of what I could try, or I am doing something wrong?

I highly recommend reading Solving Large Stiff Equations · DifferentialEquations.jl

2 Likes

With great surprise the solution to my problem was on the documentation page carrying the name of my problem… I should have read it more carefully. Declaring a sparse Jacobian greatly improved the performances.

1 Like

One last thing worth mentioning is that the autoswitch methods may be valuable if your problem is only sometimes stiff.