Rethinking the linear solver and Jacobian interface

[ChrisRackauckas]

With the new advanced ODE solving tutorial, one thing that has become more clear is that we need to think about how we're doing some of the linear solver pieces.

• the preconditioners really need to be a function so that they can use the most current Jacobians.
• there is a weird interaction between some undocumented keyword arguments thrown towards the linear solvers.
• for preconditioned Krylov you may not want to build the full Jacobian and instead do an approximation. Right now you do that by passing a jac function that would be approximate, and then use a JacVecOperator.
• It would be nice if JacVecOperator was actually automatic when doing something Newton-Krylov.

[j-fu]

Hi, just looked for an open issue where this fits... Is there a way to calculate f and J at once ? I put together my system matrix and the residual vector by a number of calls to ForwardDiff.jacobian, so this is coming out naturally. The following works with VoronoiFVM.jl on unstructured grids (I just show the pattern here, will release that code soon):

using DifferentialEquations
using LinearAlgebra
using SparseArrays

mutable struct SolverContext
    J
    f
    uhash
end

function solvercontext(u)
    n=length(u)
    SolverContext(spzeros(n,n), zeros(n),0x0)
end

function lorenz!(J,du,u)
    du[1] = 10.0*(u[2]-u[1])
    du[2] = u[1]*(28.0-u[3]) - u[2]
    du[3] = u[1]*u[2] - (8/3)*u[3]

    J[1,1]=-10
    J[1,2]=10
    J[2,1]=28-u[3]
    J[2,2]=-1
    J[2,3]=-u[1]
    J[3,1]=u[2]
    J[3,2]=u[1]
    J[3,3]=-8/3
end

function lorenz_eval!(p,u)
    uhash=hash(u)
    if uhash!=p.uhash
        lorenz!(p.J,p.f,u)
        p.uhash=uhash
    end
end

function lorenz_f!(du, u, p,t)
    lorenz_eval!(p,u)
    du.=p.f
    nothing
end

function lorenz_jac!(J, u, p,t)
    lorenz_eval!(p,u)
    J.=p.J
    nothing
end

function run()
    u0 = [1.0;0.0;0.0]
    ctx=solvercontext(u0)
    tspan = (0.0,100.0)
    f = ODEFunction(lorenz_f!, jac=lorenz_jac!,jac_prototype=ctx.J)
    prob = ODEProblem(f,u0,tspan,ctx)
    sol = solve(prob,Rodas5(),reltol=1.0e-10,abstol=1.0e-10)
end

The main idea is that in order to prevent to run the assembly ode twice I just check if it was already run with a given vector by comparing hashes.

Some questions:

• Is there another way in the current interface ?
• How many copies of the Jacobi matrix are around ? IMHO it should be possible to go with just one.
• Yeah and the preconditioner thing. Still need to get my head around this...

Nevertheless I am just lightning struck that everything appears to work: It seems that I can run now my pde evolution problems with DifferentialEquations.jl by just adding a couple of methods to my code which provide the right interface.

And yes, I avoid costly sparsity detection on larger grids here, and my ExtendableSparse.jl matrix class seems to play well, too.

[j-fu]

Here is the first running code of this kind (see run_diffeq()) :https://j-fu.github.io/VoronoiFVM.jl/dev/examples/DiffEqCompare/

[ChrisRackauckas]

Is there a way to calculate f and J at once ?

Implicit Euler is probably the only place where that would make sense, and even then I'm a little skeptical. If it takes 20 f calls to calculate the columns of a matrix, then you at most get <5% speedup, and a Jacobian with 20 columns isn't too uncommon. This interface change would allow you to do these kinds of optimizations though, but this is one that I don't think will make much of a difference because just from the statistics, most of the primal calculations for Jacobians are repeated many times for the columns so the maximal total gain is low.

And yes, I avoid costly sparsity detection on larger grids here, and my ExtendableSparse.jl matrix class seems to play well, too.

Interesting. It might be good to hardcode some matrix coloring solutions for FVM, because indeed it can take a bit to compute. There's an interface for matrix types to overload and provide a fast coloring solution:

https://github.com/SciML/ArrayInterface.jl/blob/master/src/ArrayInterface.jl#L440-L474

How many copies of the Jacobi matrix are around ? IMHO it should be possible to go with just one.

You need two if you're not doing a Krylov method, one if you are. The reason is because factorizations are disconnected from J updates, so J and fact(J) can be at different type points as a way to reduce the total factorizations. This is one of the main ways that the number of factorizations is reduced many fold, so it's a pretty big performance optimization for a 2-fold memory increase. It's common in other stiff ODE solvers as well.

Here is the first running code of this kind (see run_diffeq()) :https://j-fu.github.io/VoronoiFVM.jl/dev/examples/DiffEqCompare/

Thanks for sharing. I think in that case you Rodas5 won't scale well (because it cannot reuse factorizations well). You might want an interface to DAEProblem to use IDA. On the pure Julia side, to really perfect our suite for this kind of problem, RadauIIA5 needs sparse Jacobian support, and we need a pure Julia BDF that is better than the one we had before. We're almost there.