[WIP] Relaxation through Controller

Project.toml

@@ -33,6 +33,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"

src/OrdinaryDiffEq.jl

@@ -64,6 +64,8 @@ using ExponentialUtilities
 
 using NonlinearSolve
 
+using Roots
+
 # Required by temporary fix in not in-place methods with 12+ broadcasts
 # `MVector` is used by Nordsieck forms
 import StaticArrays: SArray, MVector, SVector, @SVector, StaticArray, MMatrix, SA
@@ -465,5 +467,8 @@ export ShampineCollocationInit, BrownFullBasicInit, NoInit
 
 export NLNewton, NLAnderson, NLFunctional, NonlinearSolveAlg
 
-export IController, PIController, PIDController
+export IController, PIController, PIDController, RelaxationController
+
+export Relaxation
+
 end # module

src/integrators/controllers.jl

@@ -28,6 +28,8 @@ reset_alg_dependent_opts!(controller::AbstractController, alg1, alg2) = nothing
 
 DiffEqBase.reinit!(integrator::ODEIntegrator, controller::AbstractController) = nothing
 
+@inline next_time_controller(::ODEIntegrator, ::AbstractController, ttmp, dt) = ttmp
+
 # Standard integral (I) step size controller
 """
     IController()
@@ -736,3 +738,81 @@ function step_accept_controller!(integrator, alg::Union{FBDF{max_order}, DFBDF{m
     integrator.cache.iters_from_event += 1
     return integrator.dt / q
 end
+
+
+
+# Relaxation step size controller
+"""
+    RelaxationController(controller, T)
+
+Controller to perform a relaxation on a step of a Runge-Kuttas method.  
+
+## References
+ - Sebastian Bleecke, Hendrik Ranocha (2023)
+    Step size control for explicit relaxation Runge-Kutta methods preserving invariants
+    [DOI: 10.1145/641876.641877](https://doi.org/10.1145/641876.641877)
+"""
+mutable struct RelaxationController{CON, T} <: AbstractController
+    controller::CON
+    gamma::T
+    function RelaxationController(controller::AbstractController, T)
+        new{typeof(controller), T}(controller, one(T))
+    end
+end
+
+mutable struct Relaxation{INV, OPT, T}
+    invariant::INV
+    opt::OPT
+    gamma_min::T
+    gamma_max::T
+    function Relaxation(invariant, opt = AlefeldPotraShi, gamma_min = 4//5, gamma_max = 6//5)
+        new{typeof(invariant), typeof(opt), typeof(gamma_min)}(invariant, opt, gamma_min, gamma_max)
+    end
+end
+
+@muladd function (r::Relaxation)(integrator)
+    @unpack t, dt, uprev, u = integrator
+    @unpack invariant, opt, gamma_min, gamma_max = integrator.opts.relaxation
+    gamma = one(t)
+    S_u = u .- uprev
+    if (invariant(gamma_min * S_u .+ uprev) .- invariant(uprev)) * (invariant(gamma_max * S_u .+ uprev) .- invariant(uprev)) ≤ 0 
+        gamma = find_zero(gamma -> invariant(gamma*S_u .+ uprev) .- invariant(uprev), (gamma_min, gamma_max), opt())
+    end
+    gamma
+end
+
+function Base.show(io::IO, controller::RelaxationController)
+    print(io, controller.controller)
+    print(io, "\n Relaxation parameters = ", controller.gamma)
+end
+
+@inline function next_time_controller(integrator::ODEIntegrator, controller::RelaxationController, ttmp, dt)
+    gamma = integrator.opts.relaxation(integrator)
+    integrator.dt *= oftype(dt, gamma)
+    vdt = integrator.dt
+    modify_dt_for_tstops!(integrator)
+    if integrator.dt != vdt
+        gamma = integrator.dt/dt
+    end
+    @. integrator.u =  integrator.uprev + gamma*(integrator.u .- integrator.uprev)
+    @. integrator.fsallast = integrator.fsalfirst + gamma*(integrator.fsallast - integrator.fsalfirst)
+    controller.gamma = gamma
+    ttmp + integrator.dt - dt
+end
+
+@inline function stepsize_controller!(integrator, controller::RelaxationController, alg)
+    stepsize_controller!(integrator, controller.controller, alg)
+end
+
+@inline function accept_step_controller(integrator, controller::RelaxationController)
+    accept_step_controller(integrator, controller.controller)
+end
+
+function step_accept_controller!(integrator, controller::RelaxationController, alg, dt_factor)
+    step_accept_controller!(integrator, controller.controller, alg, dt_factor)
+end
+
+function step_reject_controller!(integrator, controller::RelaxationController, alg)
+    integrator.dt = integrator.dt * integrator.qold / controller.gamma
+end
+

src/integrators/integrator_utils.jl

@@ -238,6 +238,7 @@ function _loopfooter!(integrator)
                 q)) *
                     oneunit(integrator.dt)
             integrator.tprev = integrator.t
+            ttmp = next_time_controller(integrator, integrator.opts.controller, ttmp, integrator.dt)
             integrator.t = if has_tstop(integrator)
                 tstop = integrator.tdir * first_tstop(integrator)
                 if abs(ttmp - tstop) <

src/integrators/type.jl

@@ -46,6 +46,7 @@ mutable struct DEOptions{absType, relType, QT, tType, Controller, F1, F2, F3, F4
     force_dtmin::Bool
     advance_to_tstop::Bool
     stop_at_next_tstop::Bool
+    relaxation
 end
 
 TruncatedStacktraces.@truncate_stacktrace DEOptions

src/solve.jl

@@ -72,6 +72,7 @@ function DiffEqBase.__init(
         alias_u0 = false,
         alias_du0 = false,
         initializealg = DefaultInit(),
+        relaxation = nothing,
         kwargs...) where {recompile_flag}
     if prob isa DiffEqBase.AbstractDAEProblem && alg isa OrdinaryDiffEqAlgorithm
         error("You cannot use an ODE Algorithm with a DAEProblem")
@@ -367,6 +368,8 @@ function DiffEqBase.__init(
         controller = default_controller(_alg, cache, qoldinit, beta1, beta2)
     end
 
+    controller = relaxation !== nothing ? RelaxationController(controller, eltype(u)) : controller
+
     save_end_user = save_end
     save_end = save_end === nothing ?
                save_everystep || isempty(saveat) || saveat isa Number ||
@@ -415,7 +418,8 @@ function DiffEqBase.__init(
         unstable_check,
         verbose, calck, force_dtmin,
         advance_to_tstop,
-        stop_at_next_tstop)
+        stop_at_next_tstop,
+        relaxation)
 
     stats = SciMLBase.DEStats(0)
     differential_vars = prob isa DAEProblem ? prob.differential_vars :

test/relaxation/harmonic_oscillator.jl

@@ -0,0 +1,23 @@
+using OrdinaryDiffEq, DiffEqDevTools, LinearAlgebra
+
+printstyled("Harmonic Oscillator\n"; bold = true)
+
+dts = (1 / 2) .^ (6:-1:4)
+
+f = (u, p, t) -> [-u[2],u[1]]
+prob = ODEProblem(
+    ODEFunction(f; analytic = (u0, p, t) -> [cos(t), sin(t)]),
+    [1.0, 0.0],
+    (0.0, 1.0))
+
+invariant(x) = norm(x)
+
+# Convergence with the old method Tsit5()
+sim = test_convergence(dts, prob, Tsit5(), adaptive = true)
+println("order of convergence of Tsit5 without relaxation : "*string(sim.𝒪est[:final]))
+
+# Convergence with relaxation with FSAL-R, i.e  f(uᵧ,ₙ₊₁) ≈ f(uᵧ,ₙ) + γ ( f(uₙ₊₁) - f(uᵧ,ₙ)) 
+r = Relaxation(invariant)
+sim_relax = test_convergence(dts, prob, Tsit5(); relaxation = r, adaptive = true)
+println("order with relaxation with FSAL-R modification: "*string(sim_relax.𝒪est[:final]))
+

test/relaxation/non_linear_oscillator.jl

@@ -0,0 +1,22 @@
+using OrdinaryDiffEq, DiffEqDevTools
+
+printstyled("Non linear harmonic oscillator\n"; bold = true)
+
+dts = (1 / 2) .^ (6:-1:4)
+
+f = (u, p, t) -> [-u[2]/(u[1]^2 + u[2]^2),u[1]/(u[1]^2 + u[2]^2)]
+prob = ODEProblem(
+    ODEFunction(f; analytic = (u0, p, t) -> [cos(t), sin(t)]),
+    [1.0, 0.0],
+    (0.0, 1.0))
+
+invariant(x) = norm(x)
+
+# Convergence with the method Tsit5()
+sim = test_convergence(dts, prob, Tsit5())
+println("order of convergence of older perform_step! : "*string(sim.𝒪est[:final]))
+
+# Convergence with relaxation with FSAL-R, i.e  f(uᵧ,ₙ₊₁) ≈ f(uᵧ,ₙ) + γ ( f(uₙ₊₁) - f(uᵧ,ₙ)) 
+r = Relaxation(invariant)
+sim = test_convergence(dts, prob, Tsit5(); relaxation = r)
+println("order with relaxation with FSAL-R modification: "*string(sim.𝒪est[:final]))

test/relaxation/non_linear_pendulum.jl

@@ -0,0 +1,24 @@
+using OrdinaryDiffEq, DiffEqDevTools
+
+printstyled("Non Linear Pendulum\n"; bold = true)
+
+dts = (1 / 2) .^ (6:-1:4)
+
+f = (u, p, t) -> [-sin(u[2]), u[1]]
+prob = ODEProblem(
+    f,
+    [1.0, 0.0],
+    (0.0, 1.0))
+
+invariant(x) = x[1]^2/2 -  cos(x[2])
+
+test_setup = Dict(:alg => Vern9(), :reltol => 1e-14, :abstol => 1e-14)
+
+# Convergence with the method Tsit5()
+sim = analyticless_test_convergence(dts, prob, Tsit5(), test_setup)
+println("order of convergence of older perform_step! : "*string(sim.𝒪est[:final]))
+
+# Convergence with relaxation with FSAL-R, i.e  f(uᵧ,ₙ₊₁) ≈ f(uᵧ,ₙ) + γ ( f(uₙ₊₁) - f(uᵧ,ₙ)) 
+r = Relaxation(invariant)
+sim = analyticless_test_convergence(dts, prob, Tsit5(), test_setup; relaxation = r)
+println("order with relaxation with FSAL-R modification: "*string(sim.𝒪est[:final]))

test/relaxation/test.jl

@@ -0,0 +1,13 @@
+using OrdinaryDiffEq, DiffEqDevTools,  Test, BenchmarkTools, LinearAlgebra
+
+f = (u, p, t) -> [-u[2],u[1]]
+prob = ODEProblem(
+    ODEFunction(f; analytic = (u0, p, t) -> [cos(t), sin(t)]),
+                [1.0, 0.0],
+                (0.0, 1.0))
+invariant(x) = norm(x)
+
+r = Relaxation(invariant)
+
+sol_relax = solve(prob, Tsit5(); relaxation = r)
+sol = solve(prob, Tsit5())

test/relaxation/time_dependant_harmonic_oscillator.jl

@@ -0,0 +1,24 @@
+using OrdinaryDiffEq, DiffEqDevTools
+
+printstyled("Time-dependent harmonic oscillator\n"; bold = true)
+
+dts = (1 / 2) .^ (6:-1:4)
+
+f = (u, p, t) -> [-(1+0.5 * sin(t))*u[2], (1+0.5 * sin(t))*u[1]]
+prob = ODEProblem(
+    ODEFunction(f; 
+                analytic = (u0, p, t)->[cos(0.5)*cos(t-0.5*cos(t))-sin(0.5)*sin(t-0.5*cos(t)), 
+                                        sin(0.5)*cos(t-0.5*cos(t))+cos(0.5)*sin(t-0.5*cos(t))]),
+    [1.0, 0.0],
+    (0.0, 1.0))
+
+invariant(x) = norm(x)
+
+# Convergence with the method Tsit5()
+sim = test_convergence(dts, prob, Tsit5(), adaptative = true)
+println("order of convergence of older perform_step! : "*string(sim.𝒪est[:final]))
+
+# Convergence with relaxation with FSAL-R, i.e  f(uᵧ,ₙ₊₁) ≈ f(uᵧ,ₙ) + γ ( f(uₙ₊₁) - f(uᵧ,ₙ)) 
+r = Relaxation(invariant)
+sim = test_convergence(dts, prob, Tsit5(); relaxation = r, adaptative = true)
+println("order with relaxation with FSAL-R modification: "*string(sim.𝒪est[:final]))