New Controller Interface

lib/ImplicitDiscreteSolve/src/ImplicitDiscreteSolve.jl

@@ -3,14 +3,20 @@ module ImplicitDiscreteSolve
 using SciMLBase
 using NonlinearSolveBase
 using NonlinearSolveFirstOrder
-using ConcreteStructs
 using SymbolicIndexingInterface: parameter_symbols
 import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, alg_cache, OrdinaryDiffEqMutableCache,
     OrdinaryDiffEqConstantCache, get_fsalfirstlast, isfsal,
     initialize!, perform_step!, isdiscretecache, isdiscretealg,
     alg_order, beta2_default, beta1_default, dt_required,
     _initialize_dae!, DefaultInit, BrownFullBasicInit, OverrideInit,
-    @muladd, @.., _unwrap_val, OrdinaryDiffEqCore, isadaptive
+    @muladd, @.., _unwrap_val, OrdinaryDiffEqCore, isadaptive,
+    AbstractController
+
+@static if Base.pkgversion(OrdinaryDiffEqCore) >= v"3.3"
+    @eval begin
+        import OrdinaryDiffEqCore: AbstractControllerCache
+    end
+end
 
 using Reexport
 @reexport using SciMLBase

lib/ImplicitDiscreteSolve/src/alg_utils.jl

@@ -19,25 +19,3 @@ dt_required(alg::IDSolve) = false
 isdiscretealg(alg::IDSolve) = true
 
 isadaptive(alg::IDSolve) = true
-
-# @concrete struct ConvergenceRateTracing
-#     inner_tracing
-# end
-
-# @concrete struct ConvergenceRateTraceTrick
-#     incrementL2norms
-#     residualL2norms
-#     trace_wrapper
-# end
-
-# function NonlinearSolveBase.init_nonlinearsolve_trace(
-#         prob, alg::IDSolve, u, fu, J, δu;
-#         trace_level::ConvergenceRateTracing, kwargs... # This kind of dispatch does not work. Need to figure out a different way.
-# )
-#     inner_trace = NonlinearSolveBase.init_nonlinearsolve_trace(
-#         prob, alg, u, fu, J, δu;
-#         trace_level.inner_tracing, kwargs...
-#     )
-
-#     return ConvergenceRateTraceTrick(eltype(δu)[], eltype(fu)[], inner_trace)
-# end

lib/ImplicitDiscreteSolve/src/controller.jl

@@ -25,67 +25,144 @@ The baseline algorithm has been derived in Peter Deuflhard's book "Newton Method
 Nonlinear Problems" in Section 5.1.3 (Adaptive pathfollowing algorithms). Please note
 that some implementation details deviate from the original algorithm.
 """
-Base.@kwdef struct KantorovichTypeController <: OrdinaryDiffEqCore.AbstractController
-    Θmin::Float64
+Base.@kwdef struct KantorovichTypeController{T} <: AbstractController
+    Θmin::T
     p::Int64
-    Θreject::Float64 = 0.95
-    Θbar::Float64 = 0.5
-    γ::Float64 = 0.95
-    qmin::Float64 = 1 / 5
-    qmax::Float64 = 5.0
+    Θreject::T = 0.95
+    Θbar::T = 0.5
+    γ::T = 0.95
+    qmin::T = 1 / 5
+    qmax::T = 5.0
     strict::Bool = true
 end
 
-function OrdinaryDiffEqCore.default_controller(
-        alg::IDSolve, cache::IDSolveCache, _1, _2, _3
-    )
-    return KantorovichTypeController(; Θmin = 1 // 8, p = 1)
-end
+@static if Base.pkgversion(OrdinaryDiffEqCore) >= v"3.3"
+    @eval begin
+        mutable struct KantorovichTypeControllerCache{T} <: AbstractControllerCache
+            controller::KantorovichTypeController{T}
+            # Proposed scaling factor for the time step length
+            q::T
+        end
 
-function OrdinaryDiffEqCore.stepsize_controller!(
-        integrator, controller::KantorovichTypeController, alg::IDSolve
-    )
-    @inline g(x) = √(1 + 4x) - 1
+        function OrdinaryDiffEqCore.default_controller_v7(
+                QT, alg::IDSolve,
+            )
+            return KantorovichTypeController{QT}(; Θmin = QT(1 // 8), p = 1)
+        end
 
-    # Adapt dt with a priori estimate (Eq. 5.24)
-    (; Θks) = integrator.cache
-    (; Θbar, γ, Θmin, qmin, qmax, p) = controller
+        function OrdinaryDiffEqCore.setup_controller_cache(alg, atmp, controller::KantorovichTypeController{T}) where {T}
+            return KantorovichTypeControllerCache(
+                controller,
+                T(1),
+            )
+        end
 
-    Θ₀ = length(Θks) > 0 ? max(first(Θks), Θmin) : Θmin
-    q = clamp(γ * (g(Θbar) / (g(Θ₀)))^(1 / p), qmin, qmax)
+        function OrdinaryDiffEqCore.stepsize_controller!(
+                integrator, cache::KantorovichTypeControllerCache, alg::IDSolve
+            )
+            (; controller) = cache
+            @inline g(x) = √(1 + 4x) - 1
 
-    return q
-end
+            # Adapt dt with a priori estimate (Eq. 5.24)
+            (; Θks) = integrator.cache
+            (; Θbar, γ, Θmin, qmin, qmax, p) = controller
 
-function OrdinaryDiffEqCore.step_accept_controller!(
-        integrator, controller::KantorovichTypeController, alg::IDSolve, q
-    )
-    return q * integrator.dt
-end
+            Θ₀ = length(Θks) > 0 ? max(first(Θks), Θmin) : Θmin
+            cache.q = clamp(γ * (g(Θbar) / (g(Θ₀)))^(1 / p), qmin, qmax)
+
+            return cache.q
+        end
+
+        function OrdinaryDiffEqCore.step_accept_controller!(
+                integrator, cache::KantorovichTypeControllerCache, alg::IDSolve, q
+            )
+            @assert q ≈ cache.q "Controller cache went out of sync with time stepping logic."
+            return cache.q * integrator.dt
+        end
+
+        function OrdinaryDiffEqCore.step_reject_controller!(
+                integrator, cache::KantorovichTypeControllerCache, alg::IDSolve
+            )
+            (; controller) = cache
+            @inline g(x) = √(1 + 4x) - 1
 
-function OrdinaryDiffEqCore.step_reject_controller!(
-        integrator, controller::KantorovichTypeController, alg::IDSolve
-    )
-    @inline g(x) = √(1 + 4x) - 1
-
-    # Shorten dt according to (Eq. 5.24)
-    (; Θks) = integrator.cache
-    (; Θbar, Θreject, γ, Θmin, qmin, qmax, p) = controller
-    for Θk in Θks
-        if Θk > Θreject
-            q = clamp(γ * (g(Θbar) / g(Θk))^(1 / p), qmin, qmax)
-            integrator.dt = q * integrator.dt
+            # Shorten dt according to (Eq. 5.24)
+            (; Θks) = integrator.cache
+            (; Θbar, Θreject, γ, Θmin, qmin, qmax, p) = controller
+            for Θk in Θks
+                if Θk > Θreject
+                    q = clamp(γ * (g(Θbar) / g(Θk))^(1 / p), qmin, qmax)
+                    integrator.dt = q * integrator.dt
+                    return
+                end
+            end
             return
         end
+
+        function OrdinaryDiffEqCore.accept_step_controller(integrator, cache::KantorovichTypeControllerCache)
+            (; controller) = cache
+            (; Θks) = integrator.cache
+            if controller.strict
+                return all(controller.Θreject .< Θks)
+            else
+                return true
+            end
+        end
     end
-    return
-end
+else
+    @eval begin
+        function OrdinaryDiffEqCore.default_controller(
+                alg::IDSolve, cache::IDSolveCache, _1, _2, _3
+            )
+            return KantorovichTypeController(; Θmin = float(1 // 8), p = 1)
+        end
+
+        function OrdinaryDiffEqCore.stepsize_controller!(
+                integrator, controller::KantorovichTypeController, alg::IDSolve
+            )
+            @inline g(x) = √(1 + 4x) - 1
+
+            # Adapt dt with a priori estimate (Eq. 5.24)
+            (; Θks) = integrator.cache
+            (; Θbar, γ, Θmin, qmin, qmax, p) = controller
+
+            Θ₀ = length(Θks) > 0 ? max(first(Θks), Θmin) : Θmin
+            q = clamp(γ * (g(Θbar) / (g(Θ₀)))^(1 / p), qmin, qmax)
 
-function OrdinaryDiffEqCore.accept_step_controller(integrator, controller::KantorovichTypeController)
-    (; Θks) = integrator.cache
-    if controller.strict
-        return all(controller.Θreject .< Θks)
-    else
-        return true
+            return q
+        end
+
+        function OrdinaryDiffEqCore.step_accept_controller!(
+                integrator, controller::KantorovichTypeController, alg::IDSolve, q
+            )
+            return q * integrator.dt
+        end
+
+        function OrdinaryDiffEqCore.step_reject_controller!(
+                integrator, controller::KantorovichTypeController, alg::IDSolve
+            )
+            @inline g(x) = √(1 + 4x) - 1
+
+            # Shorten dt according to (Eq. 5.24)
+            (; Θks) = integrator.cache
+            (; Θbar, Θreject, γ, Θmin, qmin, qmax, p) = controller
+            for Θk in Θks
+                if Θk > Θreject
+                    q = clamp(γ * (g(Θbar) / g(Θk))^(1 / p), qmin, qmax)
+                    integrator.dt = q * integrator.dt
+                    return
+                end
+            end
+            return
+        end
+
+        function OrdinaryDiffEqCore.accept_step_controller(integrator, controller::KantorovichTypeController)
+            (; Θks) = integrator.cache
+            if controller.strict
+                return all(controller.Θreject .< Θks)
+            else
+                return true
+            end
+        end
     end
 end

lib/OrdinaryDiffEqBDF/src/OrdinaryDiffEqBDF.jl

@@ -16,7 +16,8 @@ import OrdinaryDiffEqCore: alg_order, calculate_residuals!,
     trivial_limiter!,
     issplit, qsteady_min_default, qsteady_max_default,
     get_current_alg_order, get_current_adaptive_order,
-    default_controller, stepsize_controller!,
+    default_controller_v7,
+    legacy_default_controller, stepsize_controller!,
     step_accept_controller!,
     step_reject_controller!, post_newton_controller!,
     u_modified!, DAEAlgorithm, _unwrap_val, DummyController,

lib/OrdinaryDiffEqBDF/src/controllers.jl

@@ -1,4 +1,8 @@
-function default_controller(alg::Union{QNDF, FBDF}, args...)
+function legacy_default_controller(alg::Union{QNDF, FBDF}, args...)
+    return DummyController()
+end
+
+function default_controller_v7(QT, alg::Union{QNDF, FBDF}, args...)
     return DummyController()
 end
 
@@ -9,7 +13,11 @@ stepsize_controller!(integrator, alg::QNDF) = nothing
 # Implementation of an Adaptive BDF2 Formula and Comparison with the MATLAB Ode15s paper
 # E. Alberdi Celaya, J. J. Anza Aguirrezabala, and P. Chatzipantelidis
 
-function step_accept_controller!(integrator, alg::QNDF{max_order}, q) where {max_order}
+function step_accept_controller!(integrator, alg::QNDF, q)
+    return step_accept_controller!(integrator, integrator.cache, alg, q)
+end
+
+function step_accept_controller!(integrator, cache::Union{QNDFCache, QNDFConstantCache}, alg::QNDF{max_order}, q) where {max_order}
     #step is accepted, reset count of consecutive failed steps
     integrator.cache.consfailcnt = 0
     integrator.cache.nconsteps += 1
@@ -250,15 +258,22 @@ end
 
 function stepsize_controller!(
         integrator,
+        alg::FBDF
+    )
+    return stepsize_controller!(integrator, integrator.cache, alg)
+end
+
+function stepsize_controller!(
+        integrator,
+        cache::Union{FBDFCache, FBDFConstantCache},
         alg::FBDF{max_order}
     ) where {
         max_order,
     }
-    (; cache) = integrator
     cache.prev_order = cache.order
     k, terk = choose_order!(alg, integrator, cache, Val(max_order))
     if k != cache.order
-        integrator.cache.nconsteps = 0
+        cache.nconsteps = 0
         cache.order = k
     end
     if iszero(terk)
@@ -270,16 +285,20 @@ function stepsize_controller!(
     return q
 end
 
+function step_accept_controller!(integrator, alg::FBDF, q)
+    return step_accept_controller!(integrator, integrator.cache, alg, q)
+end
+
 function step_accept_controller!(
-        integrator, alg::FBDF{max_order},
+        integrator, cache::Union{FBDFCache, FBDFConstantCache}, alg::FBDF{max_order},
         q
     ) where {max_order}
-    integrator.cache.consfailcnt = 0
+    cache.consfailcnt = 0
     if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
         q = one(q)
     end
-    integrator.cache.nconsteps += 1
-    integrator.cache.iters_from_event += 1
+    cache.nconsteps += 1
+    cache.iters_from_event += 1
     return integrator.dt / q
 end
 
@@ -396,15 +415,22 @@ end
 
 function stepsize_controller!(
         integrator,
+        alg::DFBDF
+    )
+    return stepsize_controller!(integrator, integrator.cache, alg)
+end
+
+function stepsize_controller!(
+        integrator,
+        cache::Union{DFBDFCache, DFBDFConstantCache},
         alg::DFBDF{max_order}
     ) where {
         max_order,
     }
-    (; cache) = integrator
     cache.prev_order = cache.order
     k, terk = choose_order!(alg, integrator, cache, Val(max_order))
     if k != cache.order
-        integrator.cache.nconsteps = 0
+        cache.nconsteps = 0
         cache.order = k
     end
     if iszero(terk)
@@ -416,15 +442,19 @@ function stepsize_controller!(
     return q
 end
 
+function step_accept_controller!(integrator, alg::DFBDF, q)
+    return step_accept_controller!(integrator, integrator.cache, alg, q)
+end
+
 function step_accept_controller!(
-        integrator, alg::DFBDF{max_order},
+        integrator, cache::Union{DFBDFCache, DFBDFConstantCache}, alg::DFBDF{max_order},
         q
     ) where {max_order}
-    integrator.cache.consfailcnt = 0
+    cache.consfailcnt = 0
     if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
         q = one(q)
     end
-    integrator.cache.nconsteps += 1
-    integrator.cache.iters_from_event += 1
+    cache.nconsteps += 1
+    cache.iters_from_event += 1
     return integrator.dt / q
 end

lib/OrdinaryDiffEqCore/Project.toml

@@ -1,7 +1,7 @@
 name = "OrdinaryDiffEqCore"
 uuid = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 authors = ["ParamThakkar123 <[email protected]>"]
-version = "3.2.0"
+version = "3.3.0"
 
 [deps]
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"