Setup NonlinearSolveAlg with jacobian reuse

[ChrisRackauckas]

No description provided.

[ChrisRackauckas]

ROBER is benchmarking for type inference issues.

using NonlinearSolve, MINPACK, ModelingToolkit, OrdinaryDiffEqBDF, OrdinaryDiffEqNonlinearSolve, Sundials, LSODA
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
using ModelingToolkit: t_nounits as t, D_nounits as D
RUS = RadiusUpdateSchemes

@parameters k₁ k₂ k₃
@variables y₁(t) y₂(t) y₃(t)

eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃,
    D(y₂) ~ k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃,
    D(y₃) ~ k₂ * y₂^2]
rober = ODESystem(eqs, t; name = :rober) |> structural_simplify |> complete
prob = ODEProblem(rober, [y₁, y₂, y₃] .=> [1.0; 0.0; 0.0], (0.0, 1e5), [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4), jac = true)

nlalgs = [
    NonlinearSolveAlg(TrustRegion(autodiff = AutoFiniteDiff()))
    NonlinearSolveAlg(TrustRegion(autodiff = AutoFiniteDiff(), radius_update_scheme = RUS.Bastin))
    NonlinearSolveAlg(NewtonRaphson(autodiff = AutoFiniteDiff()))
    NonlinearSolveAlg(LevenbergMarquardt(autodiff = AutoFiniteDiff()))
    NonlinearSolveAlg(LevenbergMarquardt(autodiff = AutoFiniteDiff()))
    NonlinearSolveAlg(Broyden())
    NonlinearSolveAlg(PseudoTransient(autodiff = AutoFiniteDiff()))
    NonlinearSolveAlg(DFSane())
    NonlinearSolveAlg(CMINPACK(; method=:hybr))
]
sol = solve(prob, FBDF(autodiff=false, nlsolve = nlalgs[3]));
sol = solve(prob, FBDF());

using Plots
plot(sol)
using BenchmarkTools
@btime sol = solve(prob, FBDF(autodiff=false, nlsolve = nlalgs[3]));

@btime sol = solve(prob, FBDF(autodiff=false, nlsolve = nlalgs[1]));
@btime sol = solve(prob, FBDF(autodiff=false));

@btime sol = solve(prob, FBDF());
@btime sol = solve(prob, lsoda());
@btime sol = solve(prob, CVODE_BDF());

sol = solve(prob, FBDF())
sol = solve(prob, CVODE_BDF())
solve(prob, FBDF(autodiff=false, nlsolve = nlalgs[3]));

@profview for i in 1:100 solve(prob, FBDF(autodiff=false, nlsolve = nlalgs[3])) end

OrdinaryDiffEqBDF.nlsolve!(nlsolver, integ, integ.cache, false) 
@inferred SciMLBase.get_du(_cache, _idx)

@which OrdinaryDiffEqBDF.nlsolve!(nlsolver, integ, integ.cache, false)
@which OrdinaryDiffEqBDF.compute_step!(nlsolver, integ)

integ = init(prob, FBDF(autodiff=false, nlsolve = NonlinearSolveAlg(NewtonRaphson(autodiff = AutoFiniteDiff()))))
(;ts, u_history, order, u_corrector, bdf_coeffs, r, nlsolver, weights, terk_tmp, terkp1_tmp, atmp, tmp, equi_ts, u₀, ts_tmp, step_limiter!) = integ.cache;
(;z, tmp, ztmp, γ, α, cache, method) = nlsolver;
(;tstep, invγdt, atmp, ustep )= cache;

@inferred NonlinearSolveBase.Utils.evaluate_f!(nlsolver.cache.cache, nlsolver.cache.cache.u, nlsolver.cache.cache.p)

@which OrdinaryDiffEqNonlinearSolve.step!(nlsolver.cache.cache)
@which NonlinearSolveBase.InternalAPI.step!(nlsolver.cache.cache)
@profview for i in 1:100 OrdinaryDiffEqBDF.nlsolve!(nlsolver, integ, integ.cache, false) end

BRUSS is benchmarking for reuse correctness.

using LinearAlgebra, SparseArrays

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
function brusselator_2d_loop(du, u, p, t)
    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
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)
prob2 = ODEProblem(brusselator_2d_loop, u0, (0.0, 11.5), p)

sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[1]), dt = 1e-6);
sol = solve(prob2, FBDF(autodiff=false));

@btime sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[3]), dt = 1e-6);
@btime sol = solve(prob2, FBDF(autodiff=false));
@btime sol = solve(prob2, FBDF());
@btime sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[3]));

@btime sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[1]), dt=1e-6);

One oddity is that the initializations seem off: this is solved by setting a starting dt, but it should be investigated why that's necessary.

[oscardssmith]

We really should fix our stiff alg initial dt algorithm...