Add GlobalErrorTransport global error estimation and control

Implements the linearized error-transport estimator
(SciML/GlobalDiffEq.jl#5; Shampine 1986 / Berzins 1988 / Lang-Verwer
2007): after a dense forward solve, the companion linear ODE
ε' = Jε + d(t) driven by the dense-output defect d(t) = f(P(t)) - P'(t)
is integrated for the endpoint global error, with matrix-free
Jacobian-vector products through DifferentiationInterface (autodiff
accepts any ADTypes backend, default AutoForwardDiff). gtol-controlled
solves tighten local tolerances via the shared refinement loop in
companion.jl, mirroring GlobalAdjoint. global_error_estimate(prob, alg)
is the standalone estimator, also added for GlobalAdjoint as an alias
of adjoint_error_estimate.

Verified locally on Lotka-Volterra vs a 1e-13 reference: estimate to
true-error ratios 0.995-0.999 at tolerances 1e-4/1e-6, and gtol=1e-7
control yields 9.8e-9 final error.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

docs/src/globalerrorcontrol/GlobalDiffEq.md

@@ -23,10 +23,12 @@ using GlobalDiffEq
     only a few extra stages per step over a plain method of the same order and
     require nothing beyond the right-hand side `f`.
   - To *control* the endpoint global error to a tolerance `gtol`, wrap any
-    adaptive solver in [`GlobalAdjoint`](@ref) (adjoint-based, for endpoint
-    functionals; requires SciMLSensitivity and QuadGK to be loaded). It solves
-    the problem, estimates the endpoint global error, and tightens the local
-    tolerances until the requested global tolerance is met.
+    adaptive solver in [`GlobalErrorTransport`](@ref) (linearized
+    error-transport equation, Jacobian-vector products via automatic
+    differentiation) or [`GlobalAdjoint`](@ref) (adjoint-based, for endpoint
+    functionals; requires SciMLSensitivity and QuadGK to be loaded). Each
+    solves the problem, estimates the endpoint global error, and tightens the
+    local tolerances until the requested global tolerance is met.
   - [`GlobalRichardson`](@ref) wraps any fixed-step method in global Richardson
     extrapolation over whole solves, interpreting `abstol` and `reltol` as
     global tolerances. It is the most robust and most expensive option.
@@ -60,6 +62,7 @@ global_error_estimate
 
 ```@docs
 GlobalRichardson
+GlobalErrorTransport
 GlobalAdjoint
 adjoint_error_estimate
 ```

lib/GlobalDiffEq/Project.toml

@@ -1,11 +1,14 @@
 name = "GlobalDiffEq"
 uuid = "1d72d19b-84cc-4cb7-a099-7cbdb9ccc67c"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "1.5.0"
+version = "1.6.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
@@ -31,8 +34,11 @@ OrdinaryDiffEqCore = {path = "../OrdinaryDiffEqCore"}
 OrdinaryDiffEqTsit5 = {path = "../OrdinaryDiffEqTsit5"}
 
 [compat]
+ADTypes = "1"
 DiffEqBase = "7"
+DifferentiationInterface = "0.6, 0.7"
 FastBroadcast = "1.3"
+ForwardDiff = "0.10, 1"
 LinearAlgebra = "1"
 MuladdMacro = "0.2.1"
 OrdinaryDiffEqCore = "4"

lib/GlobalDiffEq/README.md

@@ -17,6 +17,9 @@ provides:
     together with an asymptotically correct estimate of its global error.
   - `MM5GEE`: the Makazaga-Murua (2003) Dormand-Prince-based order-5 scheme
     with a cheap built-in global error estimate.
+  - `GlobalErrorTransport`: global error estimation and `gtol`-based control by
+    integrating the linearized error-transport equation driven by the
+    dense-output defect (Shampine 1986, Berzins 1988, Lang-Verwer 2007).
   - `GlobalAdjoint`: adjoint-based a posteriori endpoint error estimation and
     control (Cao and Petzold 2004), available as a package extension when
     SciMLSensitivity and QuadGK are loaded.

lib/GlobalDiffEq/src/GlobalDiffEq.jl

@@ -3,8 +3,9 @@ module GlobalDiffEq
 using Reexport: @reexport
 @reexport using DiffEqBase
 
-import LinearAlgebra, OrdinaryDiffEqCore, OrdinaryDiffEqTsit5, Random,
-    RecursiveArrayTools, Richardson, SciMLBase, SciMLStructures
+import ADTypes, DifferentiationInterface, ForwardDiff, LinearAlgebra,
+    OrdinaryDiffEqCore, OrdinaryDiffEqTsit5, Random, RecursiveArrayTools,
+    Richardson, SciMLBase, SciMLStructures
 import DiffEqBase: initialize!, calculate_residuals, calculate_residuals!
 import OrdinaryDiffEqCore: perform_step!, @cache
 using FastBroadcast: FastBroadcast, @..
@@ -15,13 +16,16 @@ abstract type GlobalDiffEqAlgorithm <: SciMLBase.AbstractODEAlgorithm end
 
 include("richardson.jl")
 include("adjoint.jl")
+include("companion.jl")
+include("transport.jl")
 include("glee/tableaus.jl")
 include("glee/algorithms.jl")
 include("glee/solve.jl")
 include("glee/caches.jl")
 include("glee/perform_step.jl")
 
 export GlobalAdjoint, GlobalRichardson, adjoint_error_estimate
+export GlobalErrorTransport
 export GLEE23, GLEE24, GLEE35, MM5GEE, global_error_estimate
 
 @setup_workload begin

lib/GlobalDiffEq/src/adjoint.jl

@@ -312,6 +312,17 @@ function adjoint_error_estimate(
     )
 end
 
+"""
+    global_error_estimate(prob, alg::GlobalAdjoint; kwargs...)
+
+Alias for [`adjoint_error_estimate`](@ref).
+"""
+function global_error_estimate(
+        prob::SciMLBase.AbstractODEProblem, alg::GlobalAdjoint, args...; kwargs...
+    )
+    return adjoint_error_estimate(prob, alg, args...; kwargs...)
+end
+
 function SciMLBase.__solve(
         prob::SciMLBase.AbstractODEProblem, alg::GlobalAdjoint, args...;
         abstol = something(alg.gtol, 1.0e-6),

lib/GlobalDiffEq/src/companion.jl

@@ -0,0 +1,143 @@
+# Shared machinery for the defect-companion global error estimators
+# (GlobalErrorTransport and GlobalDefectCorrection): both solve an auxiliary
+# ODE for the global error driven by the defect of the dense numerical
+# solution, d(t) = f(P(t)) - P'(t), where P is the solution interpolant.
+
+function _validate_estimation_problem(prob, name)
+    prob.u0 isa AbstractVector{<:AbstractFloat} ||
+        throw(ArgumentError("$name requires a real floating-point vector state"))
+    isempty(prob.u0) && throw(ArgumentError("$name requires a nonempty state"))
+    prob.tspan[1] < prob.tspan[2] ||
+        throw(ArgumentError("$name requires a forward-time ODEProblem"))
+    prob.f.mass_matrix == LinearAlgebra.I ||
+        throw(ArgumentError("$name currently requires the standard mass matrix"))
+    problem_kwargs = values(prob.kwargs)
+    if haskey(problem_kwargs, :callback) && problem_kwargs.callback !== nothing
+        throw(ArgumentError("$name does not currently support callbacks"))
+    end
+    return nothing
+end
+
+function _validate_estimation_solution(sol, name)
+    SciMLBase.successful_retcode(sol) ||
+        throw(ErrorException("the forward solve failed with retcode $(sol.retcode)"))
+    length(sol.t) >= 2 ||
+        throw(ArgumentError("the forward solution must save at least its two endpoints"))
+    return nothing
+end
+
+# Out-of-place view of the (possibly in-place) user RHS, safe for dual numbers.
+function _oop_rhs(prob)
+    f = SciMLBase.unwrapped_f(prob.f)
+    if SciMLBase.isinplace(prob)
+        return let f = f
+            function (u, p, t)
+                du = similar(u)
+                f(du, u, p, t)
+                return du
+            end
+        end
+    else
+        return f
+    end
+end
+
+const _DENSE_SOLVE_KWARGS = (;
+    dense = true,
+    save_everystep = true,
+    save_start = true,
+    save_end = true,
+    saveat = (),
+    save_idxs = nothing,
+)
+
+# Endpoint-error refinement loop shared by the companion estimators: solve,
+# estimate the endpoint global error, and tighten the local tolerances until
+# the estimate is at most gtol; then redo the solve with the user's original
+# saving options.
+function _refine_to_gtol(
+        estimator, prob, inner_alg, gtol, options, args...;
+        abstol, reltol, kwargs...
+    )
+    local_abstol = abstol
+    local_reltol = reltol
+    last_estimate = oftype(float(gtol), Inf)
+
+    for _ in 1:options.maxiters
+        last_estimate = estimator(local_abstol, local_reltol)
+        isfinite(last_estimate) ||
+            throw(ErrorException("the global error estimate is not finite"))
+
+        if last_estimate <= gtol
+            return SciMLBase.solve(
+                prob, inner_alg, args...;
+                abstol = local_abstol, reltol = local_reltol, kwargs...
+            )
+        end
+
+        tolerance_scale = min(0.5, options.safety * gtol / last_estimate)
+        local_abstol *= tolerance_scale
+        local_reltol *= tolerance_scale
+        iszero(local_abstol) && iszero(local_reltol) &&
+            throw(ErrorException("local tolerances underflowed during global error control"))
+    end
+
+    throw(
+        ErrorException(
+            "failed to meet gtol=$(gtol) after $(options.maxiters) iterations; " *
+                "last estimated global error was $(last_estimate)"
+        )
+    )
+end
+
+function _companion_options(gtol, maxiters, safety, companion_abstol, companion_reltol)
+    maxiters isa Integer && maxiters > 0 ||
+        throw(ArgumentError("maxiters must be a positive integer"))
+    safety isa Real && isfinite(safety) && 0 < safety < 1 ||
+        throw(ArgumentError("safety must be between zero and one"))
+    (gtol === nothing || _positive_finite_real(gtol)) ||
+        throw(ArgumentError("gtol must be a positive finite real number"))
+    _validate_tolerances(companion_abstol, companion_reltol, "companion")
+    return (;
+        maxiters = Int(maxiters), safety,
+        companion_abstol, companion_reltol,
+    )
+end
+
+# Dense forward solve + companion error solve, returning the endpoint global
+# error 2-norm estimate. rhs_for(sol) builds the companion RHS from the dense
+# forward solution.
+function _companion_error_estimate(
+        rhs_for, name, prob, inner_alg, companion_alg, args...;
+        abstol, reltol, companion_abstol, companion_reltol, kwargs...
+    )
+    haskey(kwargs, :callback) &&
+        throw(ArgumentError("$name does not currently support callbacks"))
+    _validate_estimation_problem(prob, name)
+    solve_kwargs = merge((; kwargs...), _DENSE_SOLVE_KWARGS)
+    sol = SciMLBase.solve(prob, inner_alg, args...; abstol, reltol, solve_kwargs...)
+    _validate_estimation_solution(sol, name)
+    endpoint_error = _companion_endpoint(
+        rhs_for(sol), sol, companion_alg;
+        abstol = companion_abstol, reltol = companion_reltol
+    )
+    return LinearAlgebra.norm(endpoint_error)
+end
+
+# Solve the companion error ODE ε' = rhs(ε, t), ε(t0) = 0, over the forward
+# solution's time span and return the endpoint error estimate ε(T).
+function _companion_endpoint(rhs, sol, companion_alg; abstol, reltol)
+    ε0 = zero(sol.prob.u0)
+    companion_prob = SciMLBase.ODEProblem{false}(rhs, ε0, sol.prob.tspan)
+    companion_sol = SciMLBase.solve(
+        companion_prob, companion_alg;
+        abstol, reltol, dense = false, save_everystep = false,
+        save_start = false, save_end = true
+    )
+    SciMLBase.successful_retcode(companion_sol) || throw(
+        ErrorException(
+            "the companion error solve failed with retcode $(companion_sol.retcode)"
+        )
+    )
+    return companion_sol.u[end]
+end

lib/GlobalDiffEq/src/transport.jl

@@ -0,0 +1,136 @@
+"""
+    GlobalErrorTransport(alg; gtol=nothing, transport_alg=alg,
+                         autodiff=ADTypes.AutoForwardDiff(), maxiters=6,
+                         safety=0.8, companion_abstol, companion_reltol)
+
+Wrap an ODE algorithm with global error estimation and control based on the
+linearized error-transport (first variational) equation. After a dense forward
+solve producing the interpolant `P(t)`, the global error `ε(t) ≈ y(t) - P(t)`
+is estimated by integrating the linear companion ODE
+
+```math
+ε'(t) = J(P(t), p, t) \\, ε(t) + d(t), \\qquad ε(t_0) = 0,
+```
+
+where `d(t) = f(P(t), p, t) - P'(t)` is the defect of the dense output and `J`
+is the Jacobian of `f`, applied matrix-free through Jacobian-vector products
+computed with `autodiff` (any `ADTypes` backend supported by
+DifferentiationInterface). This is the error-transport approach of Shampine
+(1986) and Berzins (1988), in the continuous defect-driven form summarized by
+Lang and Verwer (2007).
+
+Set the requested absolute endpoint error with `gtol`:
+
+```julia
+solve(prob, GlobalErrorTransport(Tsit5(); gtol = 1.0e-6))
+```
+
+The solver then tightens local tolerances until the estimated endpoint global
+error 2-norm is at most `gtol`, like [`GlobalAdjoint`](@ref). Omit `gtol` when
+using the algorithm only with [`global_error_estimate`](@ref). `transport_alg`
+selects the solver for the companion equation, and `companion_abstol` /
+`companion_reltol` control its tolerances.
+
+This implementation supports forward-time, standard-mass-matrix ODEs with real
+vector states and no callbacks. The wrapped solver must provide a dense
+first-derivative interpolation.
+
+## References
+
+  - L. F. Shampine, Global error estimation with one-step methods, Computers &
+    Mathematics with Applications 12A (1986).
+  - M. Berzins, Global error estimation in the method of lines for parabolic
+    equations, SIAM Journal on Scientific and Statistical Computing 9 (1988).
+  - J. Lang and J. Verwer, On global error estimation and control for initial
+    value problems, SIAM Journal on Scientific Computing 29 (2007).
+"""
+struct GlobalErrorTransport{A, TA, AD, G, V} <: GlobalDiffEqAlgorithm
+    alg::A
+    transport_alg::TA
+    autodiff::AD
+    gtol::G
+    options::V
+end
+
+function GlobalErrorTransport(
+        alg;
+        transport_alg = alg,
+        autodiff = ADTypes.AutoForwardDiff(),
+        gtol = nothing,
+        maxiters = 6,
+        safety = 0.8,
+        companion_abstol = gtol === nothing ? 1.0e-10 : min(gtol / 100, 1.0e-10),
+        companion_reltol = gtol === nothing ? 1.0e-8 : min(gtol / 100, 1.0e-8)
+    )
+    options = _companion_options(
+        gtol, maxiters, safety, companion_abstol, companion_reltol
+    )
+    return GlobalErrorTransport(alg, transport_alg, autodiff, gtol, options)
+end
+
+SciMLBase.allows_arbitrary_number_types(::GlobalErrorTransport) = false
+SciMLBase.allowscomplex(::GlobalErrorTransport) = false
+SciMLBase.isautodifferentiable(::GlobalErrorTransport) = false
+
+function _transport_rhs(sol, autodiff)
+    prob = sol.prob
+    foop = _oop_rhs(prob)
+    return let sol = sol, foop = foop, p = prob.p, backend = autodiff
+        function (ε, _, t)
+            u = sol(t, continuity = :right)
+            du = sol(t, Val{1}, continuity = :right)
+            defect = foop(u, p, t) - du
+            jv = only(
+                DifferentiationInterface.pushforward(
+                    x -> foop(x, p, t), backend, u, (ε,)
+                )
+            )
+            return jv + defect
+        end
+    end
+end
+
+"""
+    global_error_estimate(prob, alg::GlobalErrorTransport; abstol=1e-6, reltol=1e-3, kwargs...)
+
+Solve an ODE problem and estimate the 2-norm of its global error at the final
+time by integrating the linearized error-transport equation driven by the
+dense-output defect. `abstol` and `reltol` control the forward solve;
+`companion_abstol` and `companion_reltol` control the companion error solve.
+"""
+function global_error_estimate(
+        prob::SciMLBase.AbstractODEProblem, alg::GlobalErrorTransport, args...;
+        abstol = 1.0e-6,
+        reltol = 1.0e-3,
+        companion_abstol = alg.options.companion_abstol,
+        companion_reltol = alg.options.companion_reltol,
+        kwargs...
+    )
+    return _companion_error_estimate(
+        sol -> _transport_rhs(sol, alg.autodiff), "GlobalErrorTransport",
+        prob, alg.alg, alg.transport_alg, args...;
+        abstol, reltol, companion_abstol, companion_reltol, kwargs...
+    )
+end
+
+function SciMLBase.__solve(
+        prob::SciMLBase.AbstractODEProblem, alg::GlobalErrorTransport, args...;
+        abstol = something(alg.gtol, 1.0e-6),
+        reltol = something(alg.gtol, 1.0e-3),
+        kwargs...
+    )
+    alg.gtol === nothing && throw(
+        ArgumentError("GlobalErrorTransport requires a positive `gtol` constructor keyword")
+    )
+    _validate_tolerances(abstol, reltol, "local")
+    haskey(kwargs, :callback) &&
+        throw(ArgumentError("GlobalErrorTransport does not currently support callbacks"))
+    estimator = (local_abstol, local_reltol) -> global_error_estimate(
+        prob, alg, args...;
+        abstol = local_abstol, reltol = local_reltol, kwargs...
+    )
+    return _refine_to_gtol(
+        estimator, prob, alg.alg, alg.gtol, alg.options, args...;
+        abstol, reltol, kwargs...
+    )
+end