[WIP] Relaxation Runge-Kutta

src/OrdinaryDiffEq.jl

@@ -218,6 +218,7 @@ include("perform_step/prk_perform_step.jl")
 include("perform_step/pdirk_perform_step.jl")
 include("perform_step/dae_perform_step.jl")
 include("perform_step/qprk_perform_step.jl")
+include("perform_step/test_for_relaxation_perform_step.jl")
 
 include("dense/generic_dense.jl")
 include("dense/interpolants.jl")
@@ -377,7 +378,7 @@ export FunctionMap, Euler, Heun, Ralston, Midpoint, RK4, ExplicitRK, OwrenZen3,
        OwrenZen5,
        BS3, BS5, RK46NL, DP5, Tsit5, DP8, Vern6, Vern7, Vern8, TanYam7, TsitPap8,
        Vern9, Feagin10, Feagin12, Feagin14, CompositeAlgorithm, Anas5, RKO65, FRK65, PFRK87,
-       RKM, MSRK5, MSRK6, Stepanov5, SIR54, QPRK98, PSRK4p7q6, PSRK3p6q5, PSRK3p5q4
+       RKM, MSRK5, MSRK6, Stepanov5, SIR54, QPRK98, PSRK4p7q6, PSRK3p6q5, PSRK3p5q4, Tsit5_for_relaxation
 
 export SSPRK22, SSPRK33, KYKSSPRK42, SSPRK53, SSPRK53_2N1, SSPRK53_2N2, SSPRK53_H, SSPRK63,
        SSPRK73, SSPRK83, SSPRK43, SSPRK432,

src/alg_utils.jl

@@ -588,6 +588,7 @@
 alg_order(alg::OwrenZen5) = 5
 alg_order(alg::DP5) = 5
 alg_order(alg::Tsit5) = 5
+alg_order(alg::Tsit5_for_relaxation) = 5
 alg_order(alg::DP8) = 8
 alg_order(alg::Vern6) = 6
 alg_order(alg::Vern7) = 7

src/algorithms/explicit_rk.jl

@@ -361,6 +361,19 @@
     Tsit5(stage_limiter!, step_limiter!, False())
 end
 
+Base.@kwdef struct Tsit5_for_relaxation{StageLimiter, StepLimiter, Thread} <:
+    OrdinaryDiffEqAdaptiveAlgorithm
+stage_limiter!::StageLimiter = trivial_limiter!
+step_limiter!::StepLimiter = trivial_limiter!
+thread::Thread = False()
+end
+
+# for backwards compatibility
+function Tsit5_for_relaxation(stage_limiter!, step_limiter! = trivial_limiter!)
+Tsit5_for_relaxation(stage_limiter!, step_limiter!, False())
+end
+
+
 @doc explicit_rk_docstring(
     "Hairer's 8/5/3 adaption of the Dormand-Prince Runge-Kutta method. (7th order interpolant).",
     "DP8",

src/caches/low_order_rk_caches.jl

@@ -448,6 +448,44 @@
     Base.Experimental.silence!(Tsit5Cache)
 end
 
+@cache struct Tsit5Cache_for_relaxation{uType, rateType, uNoUnitsType, StageLimiter, StepLimiter,
+    Thread} <: OrdinaryDiffEqMutableCache
+    u::uType
+    uprev::uType
+    k1::rateType
+    k2::rateType
+    k3::rateType
+    k4::rateType
+    k5::rateType
+    k6::rateType
+    k7::rateType
+    utilde::uType
+    tmp::uType
+    atmp::uNoUnitsType
+    stage_limiter!::StageLimiter
+    step_limiter!::StepLimiter
+    thread::Thread
+end
+
+function alg_cache(alg::Tsit5_for_relaxation, u, rate_prototype, ::Type{uEltypeNoUnits},
+    ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+    dt, reltol, p, calck,
+    ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    k1 = zero(rate_prototype)
+    k2 = zero(rate_prototype)
+    k3 = zero(rate_prototype)
+    k4 = zero(rate_prototype)
+    k5 = zero(rate_prototype)
+    k6 = zero(rate_prototype)
+    k7 = zero(rate_prototype)
+    utilde = zero(u)
+    atmp = similar(u, uEltypeNoUnits)
+    recursivefill!(atmp, false)
+    tmp = zero(u)
+    Tsit5Cache_for_relaxation(u, uprev, k1, k2, k3, k4, k5, k6, k7, utilde, tmp, atmp,
+        alg.stage_limiter!, alg.step_limiter!, alg.thread)
+end
+
 @cache struct RK46NLCache{uType, rateType, TabType, StageLimiter, StepLimiter, Thread} <:
               OrdinaryDiffEqMutableCache
     u::uType
@@ -547,6 +585,13 @@
     Tsit5ConstantCache()
 end
 
+function alg_cache(alg::Tsit5_for_relaxation, u, rate_prototype, ::Type{uEltypeNoUnits},
+    ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+    dt, reltol, p, calck,
+    ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+Tsit5ConstantCache_for_relaxation()
+end
+
 @cache struct DP5Cache{uType, rateType, uNoUnitsType, StageLimiter, StepLimiter,
     Thread} <: OrdinaryDiffEqMutableCache
     u::uType

src/integrators/type.jl

@@ -1,6 +1,6 @@
 mutable struct DEOptions{absType, relType, QT, tType, Controller, F1, F2, F3, F4, F5, F6,
     F7, tstopsType, discType, ECType, SType, MI, tcache, savecache,
-    disccache}
+    disccache, MT}
     maxiters::MI
     save_everystep::Bool
     adaptive::Bool
@@ -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
+    modif::MT
 end
 
 TruncatedStacktraces.@truncate_stacktrace DEOptions
@@ -134,6 +135,12 @@ mutable struct ODEIntegrator{algType <: Union{OrdinaryDiffEqAlgorithm, DAEAlgori
     stats::SciMLBase.DEStats
     initializealg::IA
     differential_vars::DV
+
+    u_propose::uType
+    u_changed::uType
+    dt_has_changed::Bool
+    dt_changed::tType
+
     fsalfirst::FSALType
     fsallast::FSALType
 
@@ -155,7 +162,8 @@ mutable struct ODEIntegrator{algType <: Union{OrdinaryDiffEqAlgorithm, DAEAlgori
             accept_step, isout, reeval_fsal, u_modified,
             reinitialize, isdae,
             opts, stats,
-            initializealg, differential_vars) where {algType, IIP, uType,
+            initializealg, differential_vars, 
+            u_propose, u_changed, dt_has_changed, dt_changed) where {algType, IIP, uType,
             duType, tType, pType,
             eigenType, EEstT,
             tTypeNoUnits, tdirType,
@@ -177,7 +185,7 @@ mutable struct ODEIntegrator{algType <: Union{OrdinaryDiffEqAlgorithm, DAEAlgori
             do_error_check,
             event_last_time, vector_event_last_time, last_event_error,
             accept_step, isout, reeval_fsal, u_modified, reinitialize, isdae,
-            opts, stats, initializealg, differential_vars) # Leave off fsalfirst and last
+            opts, stats, initializealg, differential_vars, u_propose, u_changed, dt_has_changed, dt_changed) # Leave off fsalfirst and last
     end
 end
 

src/perform_step/test_for_relaxation_perform_step.jl

@@ -0,0 +1,200 @@
+function initialize!(integrator, ::Tsit5ConstantCache_for_relaxation)
+    integrator.kshortsize = 7
+    integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
+    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+    integrator.stats.nf += 1
+
+    # Avoid undefined entries if k is an array of arrays
+    integrator.fsallast = zero(integrator.fsalfirst)
+    integrator.k[1] = integrator.fsalfirst
+    @inbounds for i in 2:(integrator.kshortsize - 1)
+        integrator.k[i] = zero(integrator.fsalfirst)
+    end
+    integrator.k[integrator.kshortsize] = integrator.fsallast
+end
+
+
+function perform_step!(integrator, cache::Tsit5ConstantCache_for_relaxation, repeat_step = false)
+
+    # Variable to know if dt has changed during perform_step
+    integrator.dt_has_changed = false
+
+    # computations! will only contain the mathematical scheme
+    # i.e the computations of the u(t+dt)
+    # the result is store not in integrator.u but integrator.u_propose
+    computations!(integrator, cache, repeat_step)
+
+    # modif_step! enables to modify the step like when we want to perform a relaxation
+    # for this we give a new struture that can be defined either by us for already known
+    # modification we want to do or by a user (see below)
+    modif_step!(integrator)
+
+    # finalize_step! will do staff related to the solver like integrator.stats, register integrator.fsal
+    # and register integrator.u
+    finalize_step!(integrator, cache)
+end
+
+
+@muladd function computations!(integrator, ::Tsit5ConstantCache_for_relaxation, repeat_step = false)
+    @unpack t, dt, uprev, u, f, p = integrator
+    T = constvalue(recursive_unitless_bottom_eltype(u))
+    T2 = constvalue(typeof(one(t)))
+    @OnDemandTableauExtract Tsit5ConstantCacheActual T T2
+    k1 = integrator.fsalfirst
+    a = dt * a21
+    k2 = f(uprev + a * k1, p, t + c1 * dt)
+    k3 = f(uprev + dt * (a31 * k1 + a32 * k2), p, t + c2 * dt)
+    k4 = f(uprev + dt * (a41 * k1 + a42 * k2 + a43 * k3), p, t + c3 * dt)
+    k5 = f(uprev + dt * (a51 * k1 + a52 * k2 + a53 * k3 + a54 * k4), p, t + c4 * dt)
+    g6 = uprev + dt * (a61 * k1 + a62 * k2 + a63 * k3 + a64 * k4 + a65 * k5)
+    k6 = f(g6, p, t + dt)
+    u = uprev + dt * (a71 * k1 + a72 * k2 + a73 * k3 + a74 * k4 + a75 * k5 + a76 * k6)
+    k7 = f(u, p, t + dt)
+    integrator.k[1] = k1
+    integrator.k[2] = k2
+    integrator.k[3] = k3
+    integrator.k[4] = k4
+    integrator.k[5] = k5
+    integrator.k[6] = k6
+    integrator.k[7] = k7
+    integrator.u_propose = u
+
+    if integrator.opts.adaptive
+        utilde = dt *
+                 (btilde1 * integrator.k[1] + btilde2 * integrator.k[2] + btilde3 * integrator.k[3] + btilde4 * integrator.k[4] + btilde5 * integrator.k[5] +
+                  btilde6 * integrator.k[6] + btilde7 * integrator.k[7])
+        atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t)
+        integrator.EEst = integrator.opts.internalnorm(atmp, t)
+    end
+end
+
+
+function modif_step!(integrator)
+
+    # Perform the modifications
+    if !(integrator.opts.modif isa Nothing)
+        integrator.opts.modif(integrator)
+
+        # Here we check the validity of chaging dt if it has changed
+        # if it is valid integrator.changed_valid will be true, if not it will be false
+        changed_valid = true
+        if integrator.dt_has_changed
+            # check dt in [dtmin, dtmax]
+            # things related to tstops
+            # surely other things
+            if changed_valid
+                integrator.u_propose = integrator.u_changed
+                integrator.dt = integrator.dt_changed
+            else
+                # print error or warning
+            end
+        end
+    end
+end
+
+
+function finalize_step!(integrator, cache::Tsit5ConstantCache_for_relaxation)
+    @unpack t, dt, uprev, u_propose, f, p = integrator
+    integrator.u = u_propose
+    integrator.fsallast = f(u_propose, p, t + dt)
+
+    integrator.stats.nf += 7
+end
+
+
+## Non Constant cache
+function initialize!(integrator, cache::Tsit5Cache_for_relaxation)
+    integrator.kshortsize = 7
+    integrator.fsalfirst = cache.k1
+    integrator.fsallast = cache.k7 # setup pointers
+    resize!(integrator.k, integrator.kshortsize)
+    # Setup k pointers
+    integrator.k[1] = cache.k1
+    integrator.k[2] = cache.k2
+    integrator.k[3] = cache.k3
+    integrator.k[4] = cache.k4
+    integrator.k[5] = cache.k5
+    integrator.k[6] = cache.k6
+    integrator.k[7] = cache.k7
+    integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+    integrator.stats.nf += 1
+    return nothing
+end
+
+
+function perform_step!(integrator, cache::Tsit5Cache_for_relaxation, repeat_step = false)
+
+    # Variable to know if dt has changed during perform_step
+    integrator.dt_has_changed = false
+
+    # computations! will only contain the mathematical scheme
+    # i.e the computations of the u(t+dt)
+    # the result is store not in integrator.u but integrator.u_propose
+    computations!(integrator, cache, repeat_step)
+
+    # modif_step! enables to modify the step like when we want to perform a relaxation
+    # for this we give a new struture that can be defined either by us for already known
+    # modification we want to do or by a user (see below)
+    modif_step!(integrator)
+
+    # finalize_step! will do staff related to the solver like integrator.stats, register integrator.fsal
+    # and register integrator.u
+    finalize_step!(integrator, cache)
+end
+
+@muladd function computations!(integrator, cache::Tsit5Cache_for_relaxation, repeat_step = false)
+    @unpack t, dt, uprev, u_propose, f, p = integrator
+    T = constvalue(recursive_unitless_bottom_eltype(u_propose))
+    T2 = constvalue(typeof(one(t)))
+    @OnDemandTableauExtract Tsit5ConstantCacheActual T T2
+    @unpack k1, k2, k3, k4, k5, k6, k7, utilde, tmp, atmp, stage_limiter!, step_limiter!, thread = cache
+    a = dt * a21
+    @.. broadcast=false thread=thread tmp=uprev + a * k1
+    stage_limiter!(tmp, f, p, t + c1 * dt)
+    f(k2, tmp, p, t + c1 * dt)
+    @.. broadcast=false thread=thread tmp=uprev + dt * (a31 * k1 + a32 * k2)
+    stage_limiter!(tmp, f, p, t + c2 * dt)
+    f(k3, tmp, p, t + c2 * dt)
+    @.. broadcast=false thread=thread tmp=uprev + dt * (a41 * k1 + a42 * k2 + a43 * k3)
+    stage_limiter!(tmp, f, p, t + c3 * dt)
+    f(k4, tmp, p, t + c3 * dt)
+    @.. broadcast=false thread=thread tmp=uprev +
+                                          dt * (a51 * k1 + a52 * k2 + a53 * k3 + a54 * k4)
+    stage_limiter!(tmp, f, p, t + c4 * dt)
+    f(k5, tmp, p, t + c4 * dt)
+    @.. broadcast=false thread=thread tmp=uprev +
+                                          dt * (a61 * k1 + a62 * k2 + a63 * k3 + a64 * k4 +
+                                           a65 * k5)
+    stage_limiter!(tmp, f, p, t + dt)
+    f(k6, tmp, p, t + dt)
+    @.. broadcast=false thread=thread u_propose = uprev +
+                                        dt * (a71 * k1 + a72 * k2 + a73 * k3 + a74 * k4 +
+                                         a75 * k5 + a76 * k6)
+    stage_limiter!(u_propose, integrator, p, t + dt)
+    step_limiter!(u_propose, integrator, p, t + dt)
+
+    f(k7, u_propose, p, t + dt)
+
+    if integrator.opts.adaptive
+        @.. broadcast=false thread=thread utilde=dt * (btilde1 * k1 + btilde2 * k2 +
+                                                  btilde3 * k3 + btilde4 * k4 +
+                                                  btilde5 * k5 + btilde6 * k6 +
+                                                  btilde7 * k7)
+        calculate_residuals!(atmp, utilde, uprev, u_propose, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t,
+            thread)
+        integrator.EEst = integrator.opts.internalnorm(atmp, t)
+    end
+    return nothing
+end
+
+function finalize_step!(integrator, cache::Tsit5Cache_for_relaxation)
+    @unpack t, dt, u_propose, u, f, p = integrator
+    @unpack k7, stage_limiter!, step_limiter!, thread = cache
+    @.. broadcast=false thread=thread u = u_propose 
+    stage_limiter!(u, integrator, p, t + dt)
+    step_limiter!(u, integrator, p, t + dt)
+    f(k7, u_propose, p, t + dt)
+    integrator.stats.nf += 7
+end