use LinearSolve precs interface

[oscardssmith]

This is a post OrdinaryDiffEq solver split version of #2247 (remaking was substantially easier than rebasing). It still needs a little bit of work (and for the LinearSolve side to be merged first), but it's getting pretty close.

[avik-pal]

and for the LinearSolve side to be merged first)

Is there a link for this? I might want to do a similar cleanup of the boilerplate in NonlinearSolve

[oscardssmith]

yeah it's SciML/LinearSolve.jl#514. It's not yet fully documented but there are some docs for it and it works.

[oscardssmith]

I believe this now works (once SciML/SciMLBase.jl#769 is merged so remake on the Linear solver works properly)

[oscardssmith]

Tests all pass locally (with deved SciMLBase). @ChrisRackauckas can you confirm that this is the API we want?

[ChrisRackauckas]

It looks right.

[oscardssmith]

hmm. The Interface II test seems to be failing stochastically.

[oscardssmith]

planning on merging as soon as tests pass

[oscardssmith]

Ah, it's not random, but it is dependent on the order of runs. Does Krylov.jl depend on the global random number generator? IMO that's pretty bad UI.

[oscardssmith]

@ChrisRackauckas This error does actually appear to be at least somewhat stochastic. Do you know where the randomness would be coming from and how to work around it? Currently, it looks like this test will fail roughly 1/60 solve calls that gets run (roughly equally likely for various solvers and precs types).

[ChrisRackauckas]

Are you factorizing any uninitialized memory? The reason for the init stuff in ArrayInterface is to avoid that, so double check it is still in play.

[oscardssmith]

I'm pretty sure that's not the problem (since I'm not getting undef errors, but rather just timesteps failing repeatedly)

[ChrisRackauckas]

You wouldn't get undef errors from similar memory, you'd just get random non-zeroed values.

[oscardssmith]

It looks like the undef J prototype that we were creating was causing all sorts of problems. Changing it to a zeromatrix seems to fix them.

[oscardssmith]

Looks like tests are finally mostly passing

[oscardssmith]

The last failure is StabilizedIRK and I honestly have absolutely no idea why it's any different. @ChrisRackauckas any thoughts?

[oscardssmith]

I think this is finally good to merge.

[ChrisRackauckas]

Show some before and after WPDs from the stiff ODE set.

[oscardssmith]

WPDs with or without precs?

[oscardssmith]

also why work-precision? this shouldn't change accuracy

[ChrisRackauckas]

It should take like 10 minutes to make them and eyeball if there's a change. This is a pretty large change so it's worth double checking.

[oscardssmith]

Here's the WP results:

using OrdinaryDiffEq, LinearSolve, Test, IncompleteLU, ModelingToolkit, DiffEqDevTools
import AlgebraicMultigrid

# Required due to https://github.com/JuliaSmoothOptimizers/Krylov.jl/pull/477
Base.eltype(::IncompleteLU.ILUFactorization{Tv, Ti}) where {Tv, Ti} = Tv
Base.eltype(::AlgebraicMultigrid.Preconditioner) = Float64

const N = 32
const xyd_brusselator = range(0, stop = 1, length = N)
brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0
limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a

const iter = Ref(0)
function brusselator_2d_loop(du, u, p, t)
    global iter[] += 1
    A, B, alpha, dx = p
    alpha = alpha / dx^2
    @inbounds for I in CartesianIndices((N, N))
        i, j = Tuple(I)
        x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
        ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
        limit(j - 1, N)
        du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
                       4u[i, j, 1]) +
                      B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
                      brusselator_f(x, y, t)
        du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] -
                       4u[i, j, 2]) +
                      A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
    end
    nothing
end
p = (3.4, 1.0, 10.0, step(xyd_brusselator))

function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I, 1] = 22 * (y * (1 - y))^(3 / 2)
        u[I, 2] = 27 * (x * (1 - x))^(3 / 2)
    end
    u
end
u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop, u0, (0.0, 11.5), p)

du0 = copy(u0)
jac = ModelingToolkit.Symbolics.jacobian_sparsity(
    (du, u) -> brusselator_2d_loop(du, u, p,
        0.0), du0,
    u0)

prob_ode_brusselator_2d_sparse = ODEProblem(
    ODEFunction(brusselator_2d_loop,
        jac_prototype = float.(jac)),
    u0, (0.0, 11.5), p)
function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
    if newW === nothing || newW
        Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(convert(
            AbstractMatrix,
            W)))
    else
        Pl = Plprev
    end
    Pl, nothing
end
abstols = 1.0 ./ 10.0 .^ (4:11)
reltols = 1.0 ./ 10.0 .^ (1:8)
setups = [Dict(:alg=>Rosenbrock23(linsolve = KrylovJL_GMRES(precs = algebraicmultigrid), concrete_jac=true))]
test_sol = [solve(prob_ode_brusselator_2d_sparse, Rodas5P(), save_everystep = false, abstol=1e-10, reltol=1e-10)]
probs=[prob_ode_brusselator_2d_sparse]
wp = WorkPrecisionSet(probs,abstols,reltols,setups;
                             save_everystep=false,appxsol=test_sol,maxiters=Int(1e5),numruns=10)
plot(wp)

master:

pr:

[ChrisRackauckas]

It does look like there are some significant regressions?

[ChrisRackauckas]

What's the current status here?

[oscardssmith]

I need to figure out what's causing this to sometimes be slower, but it's not getting merged urgently since it's not going in until OrdinaryDiffEqv7