Reaction tangent controller

docs/src/api-reference/solver.md

@@ -34,3 +34,11 @@ Thunderbolt.OS.LieTrotterGodunov
 Thunderbolt.OS.GenericSplitFunction
 Thunderbolt.OS.OperatorSplittingIntegrator
 ```
+
+## Operator Splitting Adaptivity
+
+```@docs
+Thunderbolt.AdaptiveOperatorSplittingAlgorithm
+Thunderbolt.ReactionTangentController
+Thunderbolt.get_reaction_tangent
+```

docs/src/assets/references.bib

@@ -278,3 +278,13 @@ @article{PotDubRicVinGul:2006:cmb
   pages={2425-2435},
   doi={10.1109/TBME.2006.880875}
 }
+@article{OgiBalPer:2023:seats,
+author = {Ogiermann, Dennis and Perotti, Luigi E. and Balzani, Daniel},
+title = {A simple and efficient adaptive time stepping technique for low-order operator splitting schemes applied to cardiac electrophysiology},
+journal = {International Journal for Numerical Methods in Biomedical Engineering},
+volume = {39},
+number = {2},
+pages = {e3670},
+url = {https://onlinelibrary.wiley.com/doi/abs/10.1002/cnm.3670},
+year = {2023}
+}

examples/conduction-velocity-benchmark.jl

@@ -63,18 +63,22 @@ steady_state_initializer!(u₀, odeform)
 
 # io = ParaViewWriter("spiral-wave-test")
 
-
-timestepper = OS.LieTrotterGodunov((
-    BackwardEulerSolver(
-        solution_vector_type=Vector{Float32},
-        system_matrix_type=Thunderbolt.ThreadedSparseMatrixCSR{Float32, Int32},
-        inner_solver=LinearSolve.KrylovJL_CG(atol=1.0f-6, rtol=1.0f-5),
-    ),
-    AdaptiveForwardEulerSubstepper(
-        solution_vector_type=Vector{Float32},
-        reaction_threshold=0.1f0,
-    ),
-))
+timestepper = Thunderbolt.AdaptiveOperatorSplittingAlgorithm(
+    OS.LieTrotterGodunov((
+        BackwardEulerSolver(
+            solution_vector_type=Vector{Float32},
+            system_matrix_type=Thunderbolt.ThreadedSparseMatrixCSR{Float32, Int32},
+            inner_solver=LinearSolve.KrylovJL_CG(atol=1.0f-6, rtol=1.0f-5),
+        ),
+        AdaptiveForwardEulerSubstepper(
+            solution_vector_type=Vector{Float32},
+            reaction_threshold=0.1f0,
+        )
+    )),
+    Thunderbolt.ReactionTangentController(
+        0.5, 1.0, (0.01, 0.3)
+    )
+)
 
 problem = OS.OperatorSplittingProblem(odeform, u₀, tspan)
 

src/Thunderbolt.jl

@@ -2,7 +2,7 @@ module Thunderbolt
 
 using TimerOutputs
 
-import Unrolled: @unroll
+import Unrolled: @unroll, unrolled_filter
 
 using Reexport, UnPack
 import LinearAlgebra: mul!
@@ -72,6 +72,8 @@ include("solver/interface.jl")
 include("solver/linear.jl")
 include("solver/nonlinear.jl")
 include("solver/time_integration.jl")
+include("solver/adaptivity.jl")
+
 
 include("processing/ecg.jl")
 

src/solver/adaptivity.jl

@@ -0,0 +1,99 @@
+abstract type AbstractTimeAdaptionAlgorithm end
+
+"""
+    ReactionTangentController{T <: Real} <: AbstractTimeAdaptionAlgorithm
+A timestep length controller for [`LieTrotterGodunov`](@ref) [Lie:1880:tti,Tro:1959:psg,God:1959:dmn](@cite)
+operator splitting using the reaction tangent as proposed in [OgiBalPer:2023:seats](@cite)
+# Fields
+- `σ_s::T`: steepness
+- `σ_c::T`: offset in R axis
+- `Δt_bounds::NTuple{2,T}`: lower and upper timestep length bounds
+- `Rₙ₊₁::T`: updated maximal reaction magnitude
+- `Rₙ::T`: previous reaction magnitude
+"""
+mutable struct ReactionTangentController{T <: Real} <: AbstractTimeAdaptionAlgorithm
+    const σ_s::T
+    const σ_c::T
+    const Δt_bounds::NTuple{2,T}
+    Rₙ₊₁::T
+    Rₙ::T
+end
+
+function ReactionTangentController(σ_s::T, σ_c::T, Δt_bounds::NTuple{2,T}) where {T <: Real}
+    return ReactionTangentController(σ_s, σ_c, Δt_bounds, 0.0, 0.0)
+end
+
+
+"""
+    AdaptiveOperatorSplittingAlgorithm{TOperatorSplittingAlg <: OS.AbstractOperatorSplittingAlgorithm, TTimeAdaptionAlgorithm <: AbstractTimeAdaptionAlgorithm} <: OS.AbstractOperatorSplittingAlgorithm
+A generic operator splitting algorithm `operator_splitting_algorithm` with adaptive timestepping using the controller `controller`.
+# Fields
+- `operator_splitting_algorithm::TOperatorSplittingAlg`: steepness
+- `controller::TTimeAdaptionAlgorithm`: offset in R axis
+"""
+struct AdaptiveOperatorSplittingAlgorithm{TOperatorSplittingAlg <: OS.AbstractOperatorSplittingAlgorithm, TTimeAdaptionAlgorithm <: AbstractTimeAdaptionAlgorithm} <: OS.AbstractOperatorSplittingAlgorithm
+    operator_splitting_algorithm::TOperatorSplittingAlg
+    controller::TTimeAdaptionAlgorithm
+end
+
+@inline DiffEqBase.isadaptive(::AdaptiveOperatorSplittingAlgorithm) = true
+
+"""
+    get_reaction_tangent(integrator::OS.OperatorSplittingIntegrator)
+Returns the maximal reaction magnitude using the [`PointwiseODEFunction`](@ref) of an operator splitting integrator that uses [`LieTrotterGodunov`](@ref) [Lie:1880:tti,Tro:1959:psg,God:1959:dmn](@cite).
+It is assumed that the problem containing the reaction tangent is a [`PointwiseODEFunction`](@ref).
+"""
+@inline function get_reaction_tangent(integrator::OS.OperatorSplittingIntegrator)
+    reaction_subintegrators = unrolled_filter(subintegrator -> subintegrator.f isa PointwiseODEFunction, integrator.subintegrators)
+    isempty(reaction_subintegrators) && @error "PointwiseODEFunction not found"
+    _get_reaction_tangent(reaction_subintegrators)
+end
+
+@inline #=@unroll=# function _get_reaction_tangent(subintegrators)
+    subintegrator = subintegrators[1]
+    # @unroll for subintegrator in subintegrators
+        #TODO: It should be either all the the same type or just one subintegrator after filtering, don't unroll?
+    φₘidx = transmembranepotential_index(subintegrator.f.ode)
+    return maximum(@view subintegrator.cache.dumat[:, φₘidx])
+    # end
+end
+
+@inline function OS.stepsize_controller!(integrator::OS.OperatorSplittingIntegrator, alg::AdaptiveOperatorSplittingAlgorithm)
+    OS.stepsize_controller!(integrator, alg.controller, alg)
+end
+
+@inline function OS.step_accept_controller!(integrator::OS.OperatorSplittingIntegrator, alg::AdaptiveOperatorSplittingAlgorithm, q)
+    OS.step_accept_controller!(integrator, alg.controller, alg, q)
+end
+
+@inline function OS.step_reject_controller!(integrator::OS.OperatorSplittingIntegrator, alg::AdaptiveOperatorSplittingAlgorithm, q)
+    OS.step_reject_controller!(integrator, alg.controller, alg, q)
+end
+
+@inline function OS.stepsize_controller!(integrator::OS.OperatorSplittingIntegrator, controller::ReactionTangentController, alg::AdaptiveOperatorSplittingAlgorithm{<:OS.LieTrotterGodunov})
+    @unpack σ_s, σ_c, Δt_bounds, Rₙ₊₁, Rₙ = controller
+    controller.Rₙ = controller.Rₙ₊₁
+    controller.Rₙ₊₁ = get_reaction_tangent(integrator)
+    return nothing
+end
+
+@inline function OS.step_accept_controller!(integrator::OS.OperatorSplittingIntegrator, controller::ReactionTangentController, alg::AdaptiveOperatorSplittingAlgorithm{<:OS.LieTrotterGodunov}, q)
+    @unpack σ_s, σ_c, Δt_bounds, Rₙ₊₁, Rₙ = controller
+    R = max(Rₙ, Rₙ₊₁)
+    integrator.dt = (1 - 1/(1+exp((σ_c - R)*σ_s)))*(Δt_bounds[2] - Δt_bounds[1]) + Δt_bounds[1]
+    return nothing
+end
+
+@inline function OS.step_reject_controller!(integrator::OS.OperatorSplittingIntegrator, controller::ReactionTangentController, alg::AdaptiveOperatorSplittingAlgorithm{<:OS.LieTrotterGodunov}, q)
+    return nothing # Do nothing
+end
+
+# Dispatch for outer construction
+function OS.init_cache(prob::OS.OperatorSplittingProblem, alg::AdaptiveOperatorSplittingAlgorithm; dt, kwargs...) # TODO
+    OS.init_cache(prob, alg.operator_splitting_algorithm;dt = dt, kwargs...)
+end
+
+# Dispatch for recursive construction
+function OS.construct_inner_cache(f::OS.AbstractOperatorSplitFunction, alg::AdaptiveOperatorSplittingAlgorithm, u::AbstractArray, uprev::AbstractArray)
+    OS.construct_inner_cache(f, alg.operator_splitting_algorithm, u, uprev)
+end

src/solver/operator_splitting/integrator.jl

@@ -54,7 +54,7 @@ function DiffEqBase.__init(
     callback = nothing,
     advance_to_tstop = false,
     save_func = (u, t) -> copy(u), # custom kwarg
-    dtchangeable = true,           # custom kwarg
+    dtchangeable = DiffEqBase.isadaptive(alg),           # custom kwarg
     kwargs...,
 )
     (; u0, p) = prob
@@ -111,6 +111,7 @@ function DiffEqBase.reinit!(
     tstops = integrator._tstops,
     saveat = integrator._saveat,
     reinit_callbacks = true,
+    reinit_retcode = true
 )
     (t0,tf) = tspan
     integrator.u .= u0
@@ -126,6 +127,9 @@ function DiffEqBase.reinit!(
         saving_callback = integrator.callback.discrete_callbacks[end]
         DiffEqBase.initialize!(saving_callback, u0, t0, integrator)
     end
+    if reinit_retcode
+        integrator.sol = DiffEqBase.solution_new_retcode(integrator.sol, DiffEqBase.ReturnCode.Default)
+    end
 end
 
 # called by DiffEqBase.solve
@@ -166,8 +170,6 @@ function DiffEqBase.step!(integrator::OperatorSplittingIntegrator, dt, stop_at_t
     end
 end
 
-
-
 # TimeChoiceIterator API
 @inline function DiffEqBase.get_tmp_cache(integrator::OperatorSplittingIntegrator)
     DiffEqBase.get_tmp_cache(integrator, integrator.alg, integrator.cache)
@@ -184,7 +186,35 @@ function (integrator::OperatorSplittingIntegrator)(tmp, t)
     linear_interpolation!(tmp, t, integrator.uprev, integrator.u, integrator.tprev, integrator.t)
 end
 
+"""
+    stepsize_controller!(::OperatorSplittingIntegrator)
+Updates the controller using the current state of the integrator if the operator splitting algorithm is adaptive.
+"""
+@inline function stepsize_controller!(integrator::OperatorSplittingIntegrator) 
+    algorithm = integrator.alg
+    DiffEqBase.isadaptive(algorithm) || return nothing
+    stepsize_controller!(integrator, algorithm)
+end
 
+"""
+    step_accept_controller!(::OperatorSplittingIntegrator)
+Updates `_dt` of the integrator if the step is accepted and the operator splitting algorithm is adaptive.
+"""
+@inline function step_accept_controller!(integrator::OperatorSplittingIntegrator) 
+    algorithm = integrator.alg
+    DiffEqBase.isadaptive(algorithm) || return nothing
+    step_accept_controller!(integrator, algorithm, nothing)
+end
+
+"""
+    step_reject_controller!(::OperatorSplittingIntegrator)
+Updates `_dt` of the integrator if the step is rejected and the the operator splitting algorithm is adaptive.
+"""
+@inline function step_reject_controller!(integrator::OperatorSplittingIntegrator) 
+    algorithm = integrator.alg
+    DiffEqBase.isadaptive(algorithm) || return nothing
+    step_reject_controller!(integrator, algorithm, nothing)
+end
 
 # helper functions for dealing with time-reversed integrators in the same way
 # that OrdinaryDiffEq.jl does
@@ -220,18 +250,21 @@ end
 
 function __step!(integrator)
     (; dtchangeable, tstops) = integrator
-    _dt = DiffEqBase.get_dt(integrator)
+    _dt = DiffEqBase.isadaptive(integrator.alg) ? DiffEqBase.get_dt(integrator) : integrator.dt
 
     # update dt before incrementing u; if dt is changeable and there is
     # a tstop within dt, reduce dt to tstop - t
     integrator.dt =
         !isempty(tstops) && dtchangeable ? tdir(integrator) * min(_dt, abs(first(tstops) - integrator.t)) :
         tdir(integrator) * _dt
 
+    # Propagate information down into the subintegrators
+    synchronize_subintegrators!(integrator)
     tnext = integrator.t + integrator.dt
 
      # Solve inner problems
     advance_solution_to!(integrator, tnext)
+    stepsize_controller!(integrator)
 
     # Update integrator
     # increment t by dt, rounding to the first tstop if that is roughly
@@ -242,8 +275,7 @@ function __step!(integrator)
     integrator.tprev = integrator.t
     integrator.t = !isempty(tstops) && abs(first(tstops) - tnext) < max_t_error ? first(tstops) : tnext
 
-    # Propagate information down into the subintegrators
-    synchronize_subintegrators!(integrator)
+    step_accept_controller!(integrator)
 
     # remove tstops that were just reached
     while !isempty(tstops) && reached_tstop(integrator, first(tstops))

src/solver/operator_splitting/solver.jl

@@ -1,4 +1,3 @@
-
 # Lie-Trotter-Godunov Splitting Implementation
 """
     LieTrotterGodunov <: AbstractOperatorSplittingAlgorithm
@@ -9,6 +8,8 @@ struct LieTrotterGodunov{AlgTupleType} <: AbstractOperatorSplittingAlgorithm
     # transfer_algs::TransferTupleType # Tuple of transfer algorithms from the master solution into the individual ones
 end
 
+@inline DiffEqBase.isadaptive(::AbstractOperatorSplittingAlgorithm) = false
+
 struct LieTrotterGodunovCache{uType, tmpType, iiType} <: AbstractOperatorSplittingCache
     u::uType
     uprev::uType # True previous solution
@@ -54,5 +55,4 @@ end
         advance_solution_to!(subinteg, inner_caches[i], tnext)
         finalize_local_step!(subinteg)
     end
-end
-
+end 

src/utils.jl

@@ -377,7 +377,7 @@ end
 # Transfer the element data into a vector
 function ea_collapse!(b::Vector, bes::EAVector)
     ndofs = size(b, 1)
-    @batch minbatch= max(1, ndofs÷Threads.nthreads()) for dof ∈ 1:ndofs
+    @batch minbatch=max(1, ndofs÷Threads.nthreads()) for dof ∈ 1:ndofs
         _ea_collapse_kernel!(b, dof, bes)
     end
 end

test/integration/test_electrophysiology.jl

@@ -24,7 +24,7 @@ using Thunderbolt
         end
     end
 
-    function solve_waveprop(mesh, coeff, subdomains)
+    function solve_waveprop(mesh, coeff, subdomains, isadaptive = false)
         cs = CoordinateSystemCoefficient(CartesianCoordinateSystem(mesh))
         model = MonodomainModel(
             ConstantCoefficient(1.0),
@@ -45,10 +45,18 @@ using Thunderbolt
             mesh
         )
 
-        timestepper = LieTrotterGodunov((
+        _timestepper = LieTrotterGodunov((
             BackwardEulerSolver(),
             ForwardEulerCellSolver()
         ))
+        if isadaptive
+            timestepper = Thunderbolt.AdaptiveOperatorSplittingAlgorithm(
+                _timestepper,
+                Thunderbolt.ReactionTangentController(0.5, 1.0, (0.01, 0.3))
+            )
+        else
+            timestepper = _timestepper
+        end
 
         u₀ = zeros(Float64, OS.function_size(odeform))
         simple_initializer!(u₀, odeform)
@@ -60,26 +68,43 @@ using Thunderbolt
         DiffEqBase.solve!(integrator)
         @test integrator.sol.retcode == DiffEqBase.ReturnCode.Success
         @test integrator.u ≉ u₀
+        return integrator.u
     end
 
     mesh  = generate_mesh(Hexahedron, (4, 4, 4), Vec{3}((0.0,0.0,0.0)), Vec{3}((1.0,1.0,1.0)))
     coeff = ConstantCoefficient(SymmetricTensor{2,3,Float64}((4.5e-5, 0, 0, 2.0e-5, 0, 1.0e-5)))
-    solve_waveprop(mesh, coeff, [""])
+    u = solve_waveprop(mesh, coeff, [""])
+    u_adaptive = solve_waveprop(mesh, coeff, [""], true)
+    @test u ≈ u_adaptive
 
     mesh  = generate_ideal_lv_mesh(4,1,1)
     coeff = ConstantCoefficient(SymmetricTensor{2,3,Float64}((4.5e-5, 0, 0, 2.0e-5, 0, 1.0e-5)))
-    solve_waveprop(mesh, coeff, [""])
-
+    u = solve_waveprop(mesh, coeff, [""])
+    u_adaptive = solve_waveprop(mesh, coeff, [""], true)
+    @test u ≈ u_adaptive
+
     mesh = to_mesh(generate_mixed_grid_2D())
     coeff = ConstantCoefficient(SymmetricTensor{2,2,Float64}((4.5e-5, 0, 2.0e-5)))
-    solve_waveprop(mesh, coeff, ["Pacemaker", "Myocardium"])
-    solve_waveprop(mesh, coeff, ["Pacemaker"])
-    solve_waveprop(mesh, coeff, ["Myocardium"])
+    u = solve_waveprop(mesh, coeff, ["Pacemaker", "Myocardium"])
+    u_adaptive = solve_waveprop(mesh, coeff, ["Pacemaker", "Myocardium"], true)
+    @test u ≈ u_adaptive
+    u = solve_waveprop(mesh, coeff, ["Pacemaker"])
+    u_adaptive = solve_waveprop(mesh, coeff, ["Pacemaker"], true)
+    @test u ≈ u_adaptive
+    u = solve_waveprop(mesh, coeff, ["Myocardium"])
+    u_adaptive = solve_waveprop(mesh, coeff, ["Myocardium"], true)
+    @test u ≈ u_adaptive
 
     mesh = to_mesh(generate_mixed_dimensional_grid_3D())
     coeff = ConstantCoefficient(SymmetricTensor{2,3,Float64}((4.5e-5, 0, 0, 2.0e-5, 0, 1.0e-5)))
-    solve_waveprop(mesh, coeff, ["Ventricle"])
+    u = solve_waveprop(mesh, coeff, ["Ventricle"])
+    u_adaptive = solve_waveprop(mesh, coeff, ["Ventricle"])
+    @test u ≈ u_adaptive
     coeff = ConstantCoefficient(SymmetricTensor{2,3,Float64}((5e-5, 0, 0, 5e-5, 0, 5e-5)))
-    solve_waveprop(mesh, coeff, ["Purkinje"])
-    solve_waveprop(mesh, coeff, ["Ventricle", "Purkinje"])
+    u = solve_waveprop(mesh, coeff, ["Purkinje"])
+    u_adaptive = solve_waveprop(mesh, coeff, ["Purkinje"])
+    @test u ≈ u_adaptive
+    u = solve_waveprop(mesh, coeff, ["Ventricle", "Purkinje"])
+    u_adaptive = solve_waveprop(mesh, coeff, ["Ventricle", "Purkinje"])
+    @test u ≈ u_adaptive
 end