GlobalDiffEq: add adjoint-based global error control as a SciMLSensitivity extension

docs/src/globalerrorcontrol/GlobalDiffEq.md

@@ -18,8 +18,16 @@ using GlobalDiffEq
 extrapolation over whole solves, interpreting `abstol` and `reltol` as global
 tolerances. It is the most robust and most expensive option.
 
+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.
+
 ## Global error controlling wrappers
 
 ```@docs
 GlobalRichardson
+GlobalAdjoint
+adjoint_error_estimate
 ```

lib/GlobalDiffEq/Project.toml

@@ -1,15 +1,25 @@
 name = "GlobalDiffEq"
 uuid = "1d72d19b-84cc-4cb7-a099-7cbdb9ccc67c"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "1.2.1"
+version = "1.3.0"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Richardson = "708f8203-808e-40c0-ba2d-98a6953ed40d"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
+
+[weakdeps]
+QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+
+[extensions]
+GlobalDiffEqSciMLSensitivityExt = ["SciMLSensitivity", "QuadGK"]
 
 [sources]
 DiffEqBase = {path = "../DiffEqBase"}
@@ -21,10 +31,14 @@ LinearAlgebra = "1"
 OrdinaryDiffEqSSPRK = "2"
 OrdinaryDiffEqTsit5 = "2"
 PrecompileTools = "1.1"
+QuadGK = "2.11"
+Random = "1"
 Reexport = "0.2, 1.0"
 Richardson = "1.2"
 SafeTestsets = "0.0.1, 0.1"
 SciMLBase = "3"
+SciMLSensitivity = "7.116"
+SciMLStructures = "1.10"
 SciMLTesting = "2.1"
 Test = "1"
 julia = "1.10"
@@ -33,9 +47,12 @@ julia = "1.10"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrdinaryDiffEqSSPRK = "669c94d9-1f4b-4b64-b377-1aa079aa2388"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["LinearAlgebra", "OrdinaryDiffEqSSPRK", "OrdinaryDiffEqTsit5", "SafeTestsets", "SciMLTesting", "Test"]
+test = ["LinearAlgebra", "OrdinaryDiffEqSSPRK", "OrdinaryDiffEqTsit5", "QuadGK", "Random", "SafeTestsets", "SciMLSensitivity", "SciMLTesting", "Test"]

lib/GlobalDiffEq/README.md

@@ -12,6 +12,9 @@ provides:
 
   - `GlobalRichardson`: global Richardson extrapolation of whole solves of any
     fixed-step method, interpreting `abstol`/`reltol` as global tolerances.
+  - `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.
 
 See the [OrdinaryDiffEq.jl documentation](https://docs.sciml.ai/OrdinaryDiffEq/stable/)
 for the full API reference.

lib/GlobalDiffEq/ext/GlobalDiffEqSciMLSensitivityExt.jl

@@ -0,0 +1,65 @@
+module GlobalDiffEqSciMLSensitivityExt
+
+import GlobalDiffEq, LinearAlgebra, QuadGK, SciMLBase, SciMLSensitivity
+
+function GlobalDiffEq._default_quadrature_sensealg()
+    return SciMLSensitivity.QuadratureAdjoint(autojacvec = true)
+end
+
+GlobalDiffEq._is_quadrature_adjoint(::SciMLSensitivity.QuadratureAdjoint) = true
+
+function GlobalDiffEq._adjoint_solution(
+        sol, sensealg::SciMLSensitivity.QuadratureAdjoint, adjoint_alg, direction;
+        abstol, reltol
+    )
+    terminal_gradient! = let direction = direction
+        function (out, u, p, t, i)
+            copyto!(out, direction)
+            return nothing
+        end
+    end
+    terminal_time = sol.prob.tspan[2]
+    adjoint_prob = SciMLSensitivity.ODEAdjointProblem(
+        sol, sensealg, adjoint_alg, [terminal_time], terminal_gradient!
+    )
+    adjoint_sol = SciMLBase.solve(
+        adjoint_prob, adjoint_alg;
+        abstol, reltol, dense = true, save_everystep = true
+    )
+    SciMLBase.successful_retcode(adjoint_sol) ||
+        throw(ErrorException("the adjoint solve failed with retcode $(adjoint_sol.retcode)"))
+    return adjoint_sol
+end
+
+function GlobalDiffEq._defect_projection(sol, adjoint_sol; abstol, reltol)
+    prob = sol.prob
+    isinplace = SciMLBase.isinplace(prob)
+
+    function integrand(t)
+        u = sol(t, continuity = :right)
+        du = sol(t, Val{1}, continuity = :right)
+        rhs = if isinplace
+            value = similar(u)
+            prob.f(value, u, prob.p, t)
+            value
+        else
+            prob.f(u, prob.p, t)
+        end
+        lambda = adjoint_sol(t, continuity = :right)
+        return LinearAlgebra.dot(lambda, du - rhs)
+    end
+
+    projection = zero(eltype(prob.u0))
+    for i in 1:(length(sol.t) - 1)
+        left = sol.t[i]
+        right = sol.t[i + 1]
+        left == right && continue
+        interval_projection, _ = QuadGK.quadgk(
+            integrand, left, right; atol = abstol, rtol = reltol
+        )
+        projection += interval_projection
+    end
+    return projection
+end
+
+end

lib/GlobalDiffEq/src/GlobalDiffEq.jl

@@ -3,14 +3,16 @@ module GlobalDiffEq
 using Reexport: @reexport
 @reexport using DiffEqBase
 
-import OrdinaryDiffEqTsit5, Richardson, SciMLBase
+import LinearAlgebra, OrdinaryDiffEqTsit5, Random, Richardson, SciMLBase,
+    SciMLStructures
 using PrecompileTools: @setup_workload, @compile_workload
 
 abstract type GlobalDiffEqAlgorithm <: SciMLBase.AbstractODEAlgorithm end
 
 include("richardson.jl")
+include("adjoint.jl")
 
-export GlobalRichardson
+export GlobalAdjoint, GlobalRichardson, adjoint_error_estimate
 
 @setup_workload begin
     # Simple test ODE: exponential decay du/dt = -u