Handles restarting problem of adaptive time integration methods

Project.toml

@@ -19,6 +19,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 P4est = "7d669430-f675-4ae7-b43e-fab78ec5a902"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"

src/Trixi.jl

@@ -30,7 +30,7 @@ using Reexport: @reexport
 # MPI needs to be imported before HDF5 to be able to use parallel HDF5
 # as long as HDF5.jl uses Requires.jl to enable parallel HDF5 with MPI
 using MPI: MPI
-
+using OrdinaryDiffEq
 using SciMLBase: CallbackSet, DiscreteCallback,
                  ODEProblem, ODESolution, ODEFunction,
                  SplitODEProblem
@@ -245,7 +245,7 @@ export SummaryCallback, SteadyStateCallback, AnalysisCallback, AliveCallback,
        AveragingCallback,
        AMRCallback, StepsizeCallback,
        GlmSpeedCallback, LBMCollisionCallback, EulerAcousticsCouplingCallback,
-       TrivialCallback, AnalysisCallbackCoupled
+       TrivialCallback, AnalysisCallbackCoupled, LoadRestartCallback
 
 export load_mesh, load_time, load_timestep, load_dt
 

src/callbacks_step/callbacks_step.jl

@@ -54,6 +54,7 @@ include("steady_state.jl")
 include("analysis.jl")
 include("alive.jl")
 include("save_restart.jl")
+include("load_restart.jl")
 include("save_solution.jl")
 include("time_series.jl")
 include("visualization.jl")

src/callbacks_step/load_restart.jl

@@ -0,0 +1,112 @@
+# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
+# Since these FMAs can increase the performance of many numerical algorithms,
+# we need to opt-in explicitly.
+# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
+@muladd begin
+
+    #! format: noindent
+
+    mutable struct LoadRestartCallback
+        restart_file::String
+    end
+
+    function Base.show(io::IO, cb::DiscreteCallback{<:Any, <:LoadRestartCallback})
+        @nospecialize cb # reduce precompilation time
+
+        load_callback = cb.affect!
+        print(io, "LoadRestartCallback(restart file=", load_callback.restart_file, ")")
+    end
+
+    function Base.show(io::IO, ::MIME"text/plain",
+                       cb::DiscreteCallback{<:Any, <:LoadRestartCallback})
+        @nospecialize cb # reduce precompilation time
+
+        if get(io, :compact, false)
+            show(io, cb)
+        else
+            save_load_callback = cb.affect!
+            setup = [
+                "restart file" => save_load_callback.restart_file]
+            summary_box(io, "LoadRestartCallback", setup)
+        end
+    end
+
+    function LoadRestartCallback(; restart_file = "out.h5")
+        load_callback = LoadRestartCallback(restart_file)
+
+        DiscreteCallback(load_callback, load_callback, # the first one is the condition, the second the affect!
+                         save_positions = (false, false),
+                         initialize = initialize!)
+    end
+
+    function initialize!(cb::DiscreteCallback{Condition, Affect!}, u, t,
+                         integrator) where {Condition, Affect! <: LoadRestartCallback}
+        load_callback = cb.affect!
+        if integrator.opts.adaptive
+            load_integrator!(load_callback.restart_file, integrator,
+                             integrator.opts.controller, integrator.alg)
+        end
+        return nothing
+    end
+
+    function load_integrator!(restart_file::String, integrator, controller::PIController,
+                              alg)
+        if mpi_isroot()
+            h5open(restart_file, "r") do file
+                integrator.qold = read(attributes(file)["qold"])
+                integrator.dt = read(attributes(file)["dtpropose"])
+            end
+        end
+        if mpi_isparallel()
+            recv_buf = [integrator.qold, integrator.dt]
+            MPI.Bcast!(recv_buf, mpi_root(), mpi_comm())
+            integrator.qold, integrator.dt = recv_buf
+        end
+    end
+
+    function load_integrator!(restart_file::String, integrator, controller::PIDController,
+                              alg)
+        if mpi_isroot()
+            h5open(restart_file, "r") do file
+                integrator.qold = read(attributes(file)["qold"])
+                integrator.dt = read(attributes(file)["dtpropose"])
+                err = read(file["controller_err"])
+                controller.err[1] = err[1]
+                controller.err[2] = err[2]
+                controller.err[3] = err[3]
+            end
+        end
+        if mpi_isparallel()
+            recv_buf = [integrator.qold, integrator.dt]
+            append!(recv_buf, controller.err)
+            MPI.Bcast!(recv_buf, mpi_root(), mpi_comm())
+            integrator.qold = recv_buf[1]
+            integrator.dt = recv_buf[2]
+            controller.err[1] = recv_buf[3]
+            controller.err[2] = recv_buf[4]
+            controller.err[3] = recv_buf[5]
+        end
+    end
+
+    function load_integrator!(restart_file::String, integrator, controller, alg)
+        h5open(restart_file, "r") do file
+            integrator.dt = read(attributes(file)["dtpropose"])
+        end
+        if mpi_isroot()
+            h5open(restart_file, "r") do file
+                integrator.dt = read(attributes(file)["dtpropose"])
+            end
+        end
+        if mpi_isparallel()
+            recv_buf = [integrator.dt]
+            MPI.Bcast!(recv_buf, mpi_root(), mpi_comm())
+            integrator.dt = recv_buf[1]
+        end
+    end
+
+    # this method is called to determine whether the callback should be activated
+    (load_callback::LoadRestartCallback)(u, t, integrator) = false
+
+    # this method is called when the callback is activated
+    (load_callback::LoadRestartCallback)(integrator) = nothing
+end # @muladd

src/callbacks_step/save_restart.jl

@@ -105,6 +105,10 @@ function (restart_callback::SaveRestartCallback)(integrator)
         end
 
         save_restart_file(u_ode, t, dt, iter, semi, restart_callback)
+        if integrator.opts.adaptive
+            save_restart_controller(mesh, integrator, integrator.opts.controller,
+                                    restart_callback)
+        end
     end
 
     # avoid re-evaluating possible FSAL stages

src/callbacks_step/save_restart_dg.jl

@@ -326,4 +326,48 @@ function load_restart_file_on_root(mesh::Union{ParallelTreeMesh, ParallelP4estMe
 
     return u_ode
 end
+
+function save_restart_controller(mesh, integrator,
+                                 controller, restart_callback)
+    if mpi_isroot()
+        save_restart_controller(integrator, controller, restart_callback)
+    end
+    if mpi_isparallel()
+        MPI.Barrier(mpi_comm())
+    end
+end
+
+function save_restart_controller(integrator, controller::PIController, restart_callback)
+    @unpack output_directory = restart_callback
+    timestep = integrator.stats.naccept
+    filename = joinpath(output_directory, @sprintf("restart_%06d.h5", timestep))
+    # Open file (preserve existing content)
+    h5open(filename, "cw") do file
+        attributes(file)["qold"] = integrator.qold
+        attributes(file)["dtpropose"] = integrator.dtpropose
+    end
+end
+
+function save_restart_controller(integrator, controller::PIDController,
+                                 restart_callback)
+    @unpack output_directory = restart_callback
+    timestep = integrator.stats.naccept
+    filename = joinpath(output_directory, @sprintf("restart_%06d.h5", timestep))
+    # Open file (preserve existing content)
+    h5open(filename, "cw") do file
+        attributes(file)["qold"] = integrator.qold
+        attributes(file)["dtpropose"] = integrator.dtpropose
+        file["controller_err"] = controller.err
+    end
+end
+
+function save_restart_controller(integrator, controller, restart_callback)
+    @unpack output_directory = restart_callback
+    timestep = integrator.stats.naccept
+    filename = joinpath(output_directory, @sprintf("restart_%06d.h5", timestep))
+    # Open file (preserve existing content)
+    h5open(filename, "cw") do file
+        attributes(file)["dtpropose"] = integrator.dtpropose
+    end
+end
 end # @muladd