Improve order estimation utility

Author: tansongchen

examples/Project.toml

@@ -6,10 +6,12 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LayeredLayouts = "f4a74d36-062a-4d48-97cd-1356bad1de4e"
+LinearRegression = "92481ed7-9fb7-40fd-80f2-46fd0f076581"
 NetworkLayout = "46757867-2c16-5918-afeb-47bfcb05e46a"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"

examples/implicit_integration.jl

@@ -1,6 +1,7 @@
 include("integration.jl")
 
-using NonlinearSolve, ODEProblemLibrary, LinearAlgebra
+using NonlinearSolve, ODEProblemLibrary, LinearAlgebra, LinearRegression
+using OrdinaryDiffEq, OrdinaryDiffEqRosenbrock
 
 # a simple fixed time-step driver for rapid algorithm prototyping
 function integrate(prob, solver; h)
@@ -46,18 +47,19 @@ struct MuTaylorExtrapolation{P, T1, T2, T3, T4}
 end
 
 function prepare(prob, solver::MuTaylorExtrapolation{P, T1, T2, T3, T4}) where {P, T1, T2, T3, T4}
-    propagator_all, _ = build_propagator(prob.f, prob.p, ntuple(_ -> 1.0, solver.order), length(prob.u0))
+    propagator_all, d_propagator_all = build_propagator(prob.f, prob.p, ntuple(_ -> 1.0, solver.order), length(prob.u0))
     propagator, _ = propagator_all
-    return (; propagator)
+    d_propagator, _ = d_propagator_all
+    return (; propagator, d_propagator)
 end
 
 function update!(u, prob, solver::MuTaylorExtrapolation{P, T1, T2, T3, T4}, cache, t, h) where {T1, T2, T3, T4, P}
     (; μ1, c1, μ2, c2) = solver
     (; propagator) = cache
     tm1 = t + μ1 * h
     tm2 = t + μ2 * h
-    nlf = (u_next, u_curr) -> propagator(u_next, tm1, -μ1 * h) - u_curr
-    nlprob1 = NonlinearProblem(nlf, copy(u), u)
+    nlf1 = (u_next, u_curr) -> propagator(u_next, tm1, -μ1 * h) - u_curr
+    nlprob1 = NonlinearProblem(nlf1, copy(u), u)
     sol1 = solve(nlprob1)
     u_trial1 = real(propagator(sol1.u, tm1, (1 - μ1) * h))
     nlf2 = (u_next, u_curr) -> propagator(u_next, tm2, -μ2 * h) - u_curr
@@ -67,27 +69,22 @@ function update!(u, prob, solver::MuTaylorExtrapolation{P, T1, T2, T3, T4}, cach
     return u .= (c1 * u_trial1 + c2 * u_trial2)
 end
 
-function estimate_order(solver, prob; dt1 = 1/12, dt2 = 1/16)
-    ref = prob.f.analytic(prob.u0, prob.p, prob.tspan[2])
-    println("Reference: ", ref)
-    u1 = integrate(prob, solver; h = dt1)
-    println("u1: ", u1)
-    u2 = integrate(prob, solver; h = dt2)
-    println("u2: ", u2)
-    err1 = norm(u1 - ref)
-    err2 = norm(u2 - ref)
-    order = log(err1 / err2) / log(dt1 / dt2)
-    return order
-end
+prob = ODEProblemLibrary.prob_ode_lotkavolterra
+refsol = solve(prob, Tsit5(); reltol = 1e-14, abstol = 1e-14)
 
-u0 = rand(2)
-linear_f = ODEFunction(
-    ODEProblemLibrary.f_2dlinear, analytic = ODEProblemLibrary.f_2dlinear_analytic
-)
-prob = ODEProblem(linear_f, u0, (0.0, 1.0), 1.01)
-prob_complex = ODEProblem(
-    linear_f, complex(u0), (0.0, 1.0), 1.01
-)
+function estimate_order(solver, prob; dts = 2. .^ (-10:-4))
+    println("Reference: ", refsol.u[end])
+    errs = Float64[]
+    for dt in dts
+        u = integrate(prob, solver; h = dt)
+        err = norm(u - refsol.u[end])
+        push!(errs, err)
+    end
+    logdts = log.(dts)
+    logerrs = log.(errs)
+    lr = linregress(logdts, logerrs)
+    return LinearRegression.slope(lr)
+end
 
 ImplicitTaylor = MuTaylor(Val(1), 1.0)
 ImplicitTaylorMidpoint = MuTaylor(Val(1), 0.5)
@@ -119,9 +116,9 @@ function normalized_pade(p::Int, q::Int)
     RI = Rational{BigInt}
     N = p + q
     c = Vector{RI}(undef, N + 1)
-    c[1] = big(1) // big(1)
+    c[1] = 1 // 1
     for k in 1:N
-        c[k + 1] = c[k] // big(k)
+        c[k + 1] = c[k] // k
     end
     if q == 0
         dres = RI[]
@@ -133,7 +130,7 @@ function normalized_pade(p::Int, q::Int)
             rhs[i] = -c[k + 1]             # -c_k
             for j in 1:q
                 idx = k - j               # 0-indexed Taylor index
-                A[i, j] = idx >= 0 ? c[idx + 1] : big(0) // big(1)
+                A[i, j] = idx >= 0 ? c[idx + 1] : 0 // 1
             end
         end
         dres = A \ rhs                       # Rational 精确求解
@@ -145,7 +142,7 @@ function normalized_pade(p::Int, q::Int)
             n[k + 1] += dres[j] * c[k - j + 1]
         end
     end
-    d = vcat(big(1) // big(1), dres)
+    d = vcat(1 // 1, dres)
     normalized_n = [x * factorial(k - 1) for (k, x) in enumerate(n)]
     normalized_d = [x * factorial(k - 1) for (k, x) in enumerate(d)]
     return normalized_n, normalized_d
@@ -159,22 +156,25 @@ end
 
 function prepare(prob, solver::TaylorPade{P, Q}) where {P, Q}
     polynomial_p, polynomial_q = normalized_pade(P, Q)
-    tuple_p = Base.tail(tuple(map(Float64, polynomial_p)...))
-    tuple_q = Base.tail(tuple(map(Float64, polynomial_q)...))
-    println("Pade coefficients (numerator): ", tuple_p)
-    println("Pade coefficients (denominator): ", tuple_q)
-    propagator_p_all, _ = build_propagator(prob.f, prob.p, tuple_p, length(prob.u0))
-    propagator_q_all, _ = build_propagator(prob.f, prob.p, tuple_q, length(prob.u0))
+    println("Pade coefficients: P = $polynomial_p, Q = $polynomial_q")
+    tuple_p = tuple(map(Float64, polynomial_p)...)
+    tuple_q = tuple(map(Float64, polynomial_q)...)
+    propagator_p_all, d_propagator_p_all = build_propagator(prob.f, prob.p, tuple_p, length(prob.u0))
+    propagator_q_all, d_propagator_q_all = build_propagator(prob.f, prob.p, tuple_q, length(prob.u0))
     propagator_p, _ = propagator_p_all
     propagator_q, _ = propagator_q_all
-    return (; propagator_p, propagator_q)
+    d_propagator_p, _ = d_propagator_p_all
+    d_propagator_q, _ = d_propagator_q_all
+    return (; propagator_p, propagator_q, d_propagator_p, d_propagator_q)
 end
 
 function update!(u, prob, solver::TaylorPade{P, Q}, cache, t, h) where {P, Q}
-    (; propagator_p, propagator_q) = cache
+    (; propagator_p, propagator_q, d_propagator_p, d_propagator_q) = cache
     rhs = propagator_p(u, t, h)
     nlf = (u_next, rhs) -> propagator_q(u_next, t + h, h) .- rhs
-    nlprob = NonlinearProblem(nlf, copy(u), rhs)
+    jac = (u_next, rhs) -> d_propagator_q(u_next, t + h, h)
+    nlf_wrapped = NonlinearFunction(nlf; jac)
+    nlprob = NonlinearProblem(nlf_wrapped, copy(u), rhs)
     sol = solve(nlprob)
     return u .= sol.u
 end

examples/integration.jl

@@ -122,12 +122,13 @@ function build_jet(f, ::Val{iip}, p, order::Val{P}, length = nothing) where {P,
     return jet
 end
 
-function build_propagator(f::ODEFunction{iip}, p, coeffs::NTuple{P, Float64}, length = nothing) where {P, iip}
+function build_propagator(f::ODEFunction{iip}, p, coeffs::NTuple{P1, Float64}, length = nothing) where {P1, iip}
     f = unwrapped_f(f)
     return build_propagator(f, Val{iip}(), p, coeffs, length)
 end
 
-function build_propagator(f, ::Val{iip}, p, coeffs::NTuple{P, Float64}, length = nothing) where {P, iip}
+function build_propagator(f, ::Val{iip}, p, coeffs::NTuple{P1, Float64}, length = nothing) where {P1, iip}
+    P = P1 - 1
     @variables t0::Real dt::Real
     u0 = isnothing(length) ? Symbolics.variable(:u0) : Symbolics.variables(:u0, 1:length)
     if iip
@@ -148,7 +149,7 @@ function build_propagator(f, ::Val{iip}, p, coeffs::NTuple{P, Float64}, length =
         d = get_coefficient(fu, index - 1) / index
         u = append_coefficient(u, d)
     end
-    ut = eval_taylor_polynomial(u, (1., coeffs...), dt)
+    ut = eval_taylor_polynomial(u, coeffs, dt)
     propagator = build_function(ut, u0, t0, dt; expression = Val(false), cse = true)
     jacobian = Symbolics.jacobian(ut, u0)
     d_propagator = build_function(jacobian, u0, t0, dt; expression = Val(false), cse = true)