Enable save solution with time intervals for SimpleIntegratorSSP (#1677)

* First attempt to enable save solution with time intervals for SimpleIntegratorSSP * Importing `add_tstop!` * First working version of SimpleIntegratorSSP with SaveSolutionCallback using time intervals * Improved formatting and updated reference solution for test * Modified initialization of tstops to ensure a stop at the end of the simulation * Added missing docstring * Removed OrdinaryDiffEq from Trixi's dependencies * Empty tstops BinaryHeap during the call to `terminate!(integrator::SimpleIntegratorSSP)` * Fixed bug and added explanatory comments * Updated Project.toml * format --------- Co-authored-by: Hendrik Ranocha <[email protected]>
trixi-framework · Dec 14, 2023 · f191444 · f191444
1 parent c7c4cf7
commit f191444
Show file tree

Hide file tree

Showing 5 changed files with 82 additions and 17 deletions.
diff --git a/Project.toml b/Project.toml
@@ -6,6 +6,7 @@ version = "0.6.5-pre"
 [deps]
 CodeTracking = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -52,6 +53,7 @@ TrixiMakieExt = "Makie"
 [compat]
 CodeTracking = "1.0.5"
 ConstructionBase = "1.3"
+DataStructures = "0.18.15"
 DiffEqCallbacks = "2.25"
 EllipsisNotation = "1.0"
 FillArrays = "0.13.2, 1"

diff --git a/examples/tree_2d_dgsem/elixir_euler_shockcapturing_subcell.jl b/examples/tree_2d_dgsem/elixir_euler_shockcapturing_subcell.jl
@@ -67,7 +67,7 @@ analysis_callback = AnalysisCallback(semi, interval = analysis_interval)
 
 alive_callback = AliveCallback(analysis_interval = analysis_interval)
 
-save_solution = SaveSolutionCallback(interval = 100,
+save_solution = SaveSolutionCallback(dt = 0.1,
                                      save_initial_solution = true,
                                      save_final_solution = true,
                                      solution_variables = cons2prim)

diff --git a/src/Trixi.jl b/src/Trixi.jl
@@ -37,7 +37,8 @@ using SciMLBase: CallbackSet, DiscreteCallback,
 import SciMLBase: get_du, get_tmp_cache, u_modified!,
                   AbstractODEIntegrator, init, step!, check_error,
                   get_proposed_dt, set_proposed_dt!,
-                  terminate!, remake
+                  terminate!, remake, add_tstop!, has_tstop, first_tstop
+
 using CodeTracking: CodeTracking
 using ConstructionBase: ConstructionBase
 using DiffEqCallbacks: PeriodicCallback, PeriodicCallbackAffect
@@ -70,6 +71,7 @@ using TriplotBase: TriplotBase
 using TriplotRecipes: DGTriPseudocolor
 @reexport using SimpleUnPack: @unpack
 using SimpleUnPack: @pack!
+using DataStructures: BinaryHeap, FasterForward, extract_all!
 
 # finite difference SBP operators
 using SummationByPartsOperators: AbstractDerivativeOperator,

diff --git a/src/time_integration/methods_SSP.jl b/src/time_integration/methods_SSP.jl
@@ -55,17 +55,25 @@ struct SimpleSSPRK33{StageCallbacks} <: SimpleAlgorithmSSP
 end
 
 # This struct is needed to fake https://github.com/SciML/OrdinaryDiffEq.jl/blob/0c2048a502101647ac35faabd80da8a5645beac7/src/integrators/type.jl#L1
-mutable struct SimpleIntegratorSSPOptions{Callback}
+mutable struct SimpleIntegratorSSPOptions{Callback, TStops}
     callback::Callback # callbacks; used in Trixi
     adaptive::Bool # whether the algorithm is adaptive; ignored
     dtmax::Float64 # ignored
     maxiters::Int # maximal number of time steps
-    tstops::Vector{Float64} # tstops from https://diffeq.sciml.ai/v6.8/basics/common_solver_opts/#Output-Control-1; ignored
+    tstops::TStops # tstops from https://diffeq.sciml.ai/v6.8/basics/common_solver_opts/#Output-Control-1; ignored
 end
 
 function SimpleIntegratorSSPOptions(callback, tspan; maxiters = typemax(Int), kwargs...)
-    SimpleIntegratorSSPOptions{typeof(callback)}(callback, false, Inf, maxiters,
-                                                 [last(tspan)])
+    tstops_internal = BinaryHeap{eltype(tspan)}(FasterForward())
+    # We add last(tspan) to make sure that the time integration stops at the end time
+    push!(tstops_internal, last(tspan))
+    # We add 2 * last(tspan) because add_tstop!(integrator, t) is only called by DiffEqCallbacks.jl if tstops contains a time that is larger than t
+    # (https://github.com/SciML/DiffEqCallbacks.jl/blob/025dfe99029bd0f30a2e027582744528eb92cd24/src/iterative_and_periodic.jl#L92)
+    push!(tstops_internal, 2 * last(tspan))
+    SimpleIntegratorSSPOptions{typeof(callback), typeof(tstops_internal)}(callback,
+                                                                          false, Inf,
+                                                                          maxiters,
+                                                                          tstops_internal)
 end
 
 # This struct is needed to fake https://github.com/SciML/OrdinaryDiffEq.jl/blob/0c2048a502101647ac35faabd80da8a5645beac7/src/integrators/type.jl#L77
@@ -78,6 +86,7 @@ mutable struct SimpleIntegratorSSP{RealT <: Real, uType, Params, Sol, F, Alg,
     du::uType
     r0::uType
     t::RealT
+    tdir::RealT
     dt::RealT # current time step
     dtcache::RealT # ignored
     iter::Int # current number of time steps (iteration)
@@ -87,8 +96,29 @@ mutable struct SimpleIntegratorSSP{RealT <: Real, uType, Params, Sol, F, Alg,
     alg::Alg
     opts::SimpleIntegratorSSPOptions
     finalstep::Bool # added for convenience
+    dtchangeable::Bool
+    force_stepfail::Bool
 end
 
+"""
+    add_tstop!(integrator::SimpleIntegratorSSP, t)
+Add a time stop during the time integration process. 
+This function is called after the periodic SaveSolutionCallback to specify the next stop to save the solution.
+"""
+function add_tstop!(integrator::SimpleIntegratorSSP, t)
+    integrator.tdir * (t - integrator.t) < zero(integrator.t) &&
+        error("Tried to add a tstop that is behind the current time. This is strictly forbidden")
+    # We need to remove the first entry of tstops when a new entry is added.
+    # Otherwise, the simulation gets stuck at the previous tstop and dt is adjusted to zero.
+    if length(integrator.opts.tstops) > 1
+        pop!(integrator.opts.tstops)
+    end
+    push!(integrator.opts.tstops, integrator.tdir * t)
+end
+
+has_tstop(integrator::SimpleIntegratorSSP) = !isempty(integrator.opts.tstops)
+first_tstop(integrator::SimpleIntegratorSSP) = first(integrator.opts.tstops)
+
 # Forward integrator.stats.naccept to integrator.iter (see GitHub PR#771)
 function Base.getproperty(integrator::SimpleIntegratorSSP, field::Symbol)
     if field === :stats
@@ -113,11 +143,13 @@ function solve(ode::ODEProblem, alg = SimpleSSPRK33()::SimpleAlgorithmSSP;
     du = similar(u)
     r0 = similar(u)
     t = first(ode.tspan)
+    tdir = sign(ode.tspan[end] - ode.tspan[1])
     iter = 0
-    integrator = SimpleIntegratorSSP(u, du, r0, t, dt, zero(dt), iter, ode.p,
+    integrator = SimpleIntegratorSSP(u, du, r0, t, tdir, dt, zero(dt), iter, ode.p,
                                      (prob = ode,), ode.f, alg,
                                      SimpleIntegratorSSPOptions(callback, ode.tspan;
-                                                                kwargs...), false)
+                                                                kwargs...),
+                                     false, true, false)
 
     # resize container
     resize!(integrator.p, nelements(integrator.p.solver, integrator.p.cache))
@@ -160,6 +192,8 @@ function solve!(integrator::SimpleIntegratorSSP)
             terminate!(integrator)
         end
 
+        modify_dt_for_tstops!(integrator)
+
         @. integrator.r0 = integrator.u
         for stage in eachindex(alg.c)
             t_stage = integrator.t + integrator.dt * alg.c[stage]
@@ -198,6 +232,10 @@ function solve!(integrator::SimpleIntegratorSSP)
         end
     end
 
+    # Empty the tstops array. 
+    # This cannot be done in terminate!(integrator::SimpleIntegratorSSP) because DiffEqCallbacks.PeriodicCallbackAffect would return at error.
+    extract_all!(integrator.opts.tstops)
+
     for stage_callback in alg.stage_callbacks
         finalize_callback(stage_callback, integrator.p)
     end
@@ -226,7 +264,30 @@ end
 # stop the time integration
 function terminate!(integrator::SimpleIntegratorSSP)
     integrator.finalstep = true
-    empty!(integrator.opts.tstops)
+end
+
+"""
+    modify_dt_for_tstops!(integrator::SimpleIntegratorSSP)
+Modify the time-step size to match the time stops specified in integrator.opts.tstops.
+To avoid adding OrdinaryDiffEq to Trixi's dependencies, this routine is a copy of
+https://github.com/SciML/OrdinaryDiffEq.jl/blob/d76335281c540ee5a6d1bd8bb634713e004f62ee/src/integrators/integrator_utils.jl#L38-L54
+"""
+function modify_dt_for_tstops!(integrator::SimpleIntegratorSSP)
+    if has_tstop(integrator)
+        tdir_t = integrator.tdir * integrator.t
+        tdir_tstop = first_tstop(integrator)
+        if integrator.opts.adaptive
+            integrator.dt = integrator.tdir *
+                            min(abs(integrator.dt), abs(tdir_tstop - tdir_t)) # step! to the end
+        elseif iszero(integrator.dtcache) && integrator.dtchangeable
+            integrator.dt = integrator.tdir * abs(tdir_tstop - tdir_t)
+        elseif integrator.dtchangeable && !integrator.force_stepfail
+            # always try to step! with dtcache, but lower if a tstop
+            # however, if force_stepfail then don't set to dtcache, and no tstop worry
+            integrator.dt = integrator.tdir *
+                            min(abs(integrator.dtcache), abs(tdir_tstop - tdir_t)) # step! to the end
+        end
+    end
 end
 
 # used for AMR

diff --git a/test/test_tree_2d_euler.jl b/test/test_tree_2d_euler.jl
@@ -214,16 +214,16 @@ end
     @test_trixi_include(joinpath(EXAMPLES_DIR,
                                  "elixir_euler_shockcapturing_subcell.jl"),
                         l2=[
-                            0.08508147906199143,
-                            0.04510299017724501,
-                            0.045103019801950375,
-                            0.6930704343869766,
+                            0.08508152653623638,
+                            0.04510301725066843,
+                            0.04510304668512745,
+                            0.6930705064715306,
                         ],
                         linf=[
-                            0.31123546471463326,
-                            0.5616274869594462,
-                            0.5619692712224448,
-                            2.88670199345138,
+                            0.31136518019691406,
+                            0.5617651935473419,
+                            0.5621200790240503,
+                            2.8866869108596056,
                         ])
     # Ensure that we do not have excessive memory allocations
     # (e.g., from type instabilities)