Added IMEX for KdV eq.

examples/kdv_1d/kdv_1d_IMEX.jl

@@ -0,0 +1,46 @@
+using OrdinaryDiffEqTsit5
+using DispersiveShallowWater
+using SummationByPartsOperators: upwind_operators, periodic_derivative_operator
+
+###############################################################################
+# Semidiscretization of the KdV equation 
+
+equations = KdVEquation1D(gravity = 9.81, D = 1.0)
+initial_condition = initial_condition_convergence_test
+boundary_conditions = boundary_condition_periodic
+
+# create homogeneous mesh
+coordinates_min = -50.0
+coordinates_max = 50.0
+N = 512
+mesh = Mesh1D(coordinates_min, coordinates_max, N)
+
+# Create solver with periodic SBP operators of accuracy order 3,
+# which results in a 4th order accurate semi discretizations.
+# We can set the accuracy order of the upwind operators to 3 since
+# we only use central versions/combinations of the upwind operators.
+D1_upwind = upwind_operators(periodic_derivative_operator;
+                             derivative_order = 1, accuracy_order = 3,
+                             xmin = xmin(mesh), xmax = xmax(mesh),
+                             N = nnodes(mesh))
+solver = Solver(D1_upwind)
+
+semi = Semidiscretization(mesh, equations, initial_condition, solver,
+                          boundary_conditions = boundary_conditions)
+
+tspan = (0.0, 5.0)
+ode = semidiscretize(semi, tspan, no_splitform = false)
+
+summary_callback = SummaryCallback()
+analysis_callback = AnalysisCallback(semi; interval = 100,
+                                     extra_analysis_errors = (:conservation_error,),
+                                     extra_analysis_integrals = (waterheight_total,
+                                                                 waterheight))
+callbacks = CallbackSet(analysis_callback, summary_callback)
+saveat = range(tspan..., length = 100)
+
+# alg = Rodas5() # not working because of https://github.com/SciML/OrdinaryDiffEq.jl/issues/2719
+# alg = KenCarp4() # would need to add OrdinaryDiffEqSDIRK - can to that, but I dont want to bloat DSW.jl with package s
+alg = Tsit5()
+sol = solve(ode, alg, abstol = 1e-7, reltol = 1e-7,
+            save_everystep = false, callback = callbacks, saveat = saveat)

src/DispersiveShallowWater.jl

@@ -30,7 +30,7 @@ using RecursiveArrayTools: ArrayPartition
 using Reexport: @reexport
 using Roots: AlefeldPotraShi, find_zero
 
-using SciMLBase: SciMLBase, DiscreteCallback, ODEProblem, ODESolution
+using SciMLBase: SciMLBase, DiscreteCallback, ODEProblem, ODESolution, SplitFunction
 import SciMLBase: u_modified!
 
 @reexport using StaticArrays: SVector

src/equations/kdv_1d.jl

@@ -181,3 +181,72 @@ function rhs!(dq, q, t, mesh, equations::KdVEquation1D, initial_condition,
 
     return nothing
 end
+
+function rhs_split_1!(dq, q, t, mesh, equations::KdVEquation1D, initial_condition,
+                      ::BoundaryConditionPeriodic, source_terms, solver, cache)
+    eta, = q.x
+    deta, = dq.x
+
+    (; c_0, DD) = cache
+    # In order to use automatic differentiation, we need to extract
+    # the storage vectors using `get_tmp` from PreallocationTools.jl
+    # so they can also hold dual numbers when needed.
+    tmp_1 = get_tmp(cache.tmp_1, eta)
+    tmp_2 = get_tmp(cache.tmp_2, eta)
+
+    @trixi_timeit timer() "third-order derivatives" begin
+        if solver.D1 isa PeriodicUpwindOperators && isnothing(solver.D3)
+            # eta_xxx = Dp * Dc * Dm * eta
+            mul!(tmp_1, solver.D1.minus, eta)
+            mul!(tmp_2, solver.D1.central, tmp_1)
+            mul!(tmp_1, solver.D1.plus, tmp_2)
+        else
+            # eta_xxx = D3 * eta
+            mul!(tmp_1, solver.D3, eta)
+        end
+
+        # deta = 1 / 6 sqrt(g * D) D^2 eta_xxx 
+
+        @.. deta = -1 / 6 * c_0 * DD * tmp_1
+    end
+
+    return nothing
+end
+
+function rhs_split_2!(dq, q, t, mesh, equations::KdVEquation1D, initial_condition,
+                      ::BoundaryConditionPeriodic, source_terms, solver, cache)
+    eta, = q.x
+    deta, = dq.x
+
+    (; c_0, c_1) = cache
+    # In order to use automatic differentiation, we need to extract
+    # the storage vectors using `get_tmp` from PreallocationTools.jl
+    # so they can also hold dual numbers when needed.
+    tmp_1 = get_tmp(cache.tmp_1, eta)
+    tmp_2 = get_tmp(cache.tmp_2, eta)
+
+    if solver.D1 isa PeriodicUpwindOperators && isnothing(solver.D3)
+        D1 = solver.D1.central
+    else
+        D1 = solver.D1
+    end
+
+    @trixi_timeit timer() "hyperbolic" begin
+        # eta2 = eta^2
+        @.. tmp_1 = eta^2
+
+        # eta2_x = D1 * eta2
+        mul!(tmp_2, D1, tmp_1)
+
+        # eta_x = D1 * eta
+        mul!(tmp_1, D1, eta)
+
+        # deta -= sqrt(g * D) * eta_x + 1 / 2 * sqrt(g / D) * (eta * eta_x + eta2_x) 
+        @.. deta = -(c_0 * tmp_1 + c_1 * (eta * tmp_1 + tmp_2))
+    end
+
+    @trixi_timeit timer() "source terms" calc_sources!(dq, q, t, source_terms, equations,
+                                                       solver)
+
+    return nothing
+end

src/semidiscretization.jl

@@ -186,6 +186,28 @@ function rhs!(dq, q, semi::Semidiscretization, t)
     return nothing
 end
 
+function rhs_split_1!(dq, q, semi::Semidiscretization, t)
+    @unpack mesh, equations, initial_condition, boundary_conditions, solver, source_terms, cache = semi
+
+    @trixi_timeit timer() "rhs_split_1!" rhs_split_1!(dq, q, t, mesh, equations,
+                                                      initial_condition,
+                                                      boundary_conditions, source_terms,
+                                                      solver,
+                                                      cache)
+    return nothing
+end
+
+function rhs_split_2!(dq, q, semi::Semidiscretization, t)
+    @unpack mesh, equations, initial_condition, boundary_conditions, solver, source_terms, cache = semi
+
+    @trixi_timeit timer() "rhs_split_2!" rhs_split_2!(dq, q, t, mesh, equations,
+                                                      initial_condition,
+                                                      boundary_conditions, source_terms,
+                                                      solver,
+                                                      cache)
+    return nothing
+end
+
 function compute_coefficients(func, t, semi::Semidiscretization)
     @unpack mesh, equations, solver = semi
     q = allocate_coefficients(mesh_equations_solver(semi)...)
@@ -214,16 +236,21 @@ function check_bathymetry(equations::AbstractShallowWaterEquations, q0)
 end
 
 """
-    semidiscretize(semi::Semidiscretization, tspan)
+    semidiscretize(semi::Semidiscretization, tspan; no_splitform = true)
 
 Wrap the semidiscretization `semi` as an ODE problem in the time interval `tspan`
 that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).
 """
-function semidiscretize(semi::Semidiscretization, tspan)
+function semidiscretize(semi::Semidiscretization, tspan; no_splitform = true)
     q0 = compute_coefficients(semi.initial_condition, first(tspan), semi)
     check_bathymetry(semi.equations, q0)
     iip = true # is-inplace, i.e., we modify a vector when calling rhs!
-    return ODEProblem{iip}(rhs!, q0, tspan, semi)
+    if no_splitform
+        ode = ODEProblem{iip}(rhs!, q0, tspan, semi)
+    else
+        ode = ODEProblem{iip}(SplitFunction(rhs_split_1!, rhs_split_2!), q0, tspan, semi)
+    end
+    return ode
 end
 
 """
@@ -241,7 +268,8 @@ of the semidiscretization `semi` at the state `q0`.
 function jacobian(semi::Semidiscretization;
                   t = 0.0,
                   q0 = compute_coefficients(semi.initial_condition, t, semi))
-    J = ForwardDiff.jacobian(similar(q0), q0) do dq, q
+    @unpack tmp_partitioned = semi.cache
+    J = ForwardDiff.jacobian(tmp_partitioned, q0) do dq, q
         DispersiveShallowWater.rhs!(dq, q, semi, t)
     end
     return J

test/test_kdv_1d.jl

@@ -59,3 +59,14 @@ end
 
     @test_allocations(semi, sol, allocs=5_000)
 end
+
+@testitem "kdv_1d_IMEX" setup=[Setup, KdVEquation1D] begin
+    @test_trixi_include(joinpath(EXAMPLES_DIR, "kdv_1d_IMEX.jl"),
+                        tspan=(0.0, 5.0),
+                        l2=[0.0007835879713461127],
+                        linf=[0.0005961613764722262],
+                        cons_error=[4.440892098500626e-16],
+                        change_waterheight=-4.440892098500626e-16)
+
+    @test_allocations_splitform(semi, sol, allocs=5_000)
+end

test/test_util.jl

@@ -176,6 +176,24 @@ macro test_allocations(semi, sol, allocs)
     end
 end
 
+"""
+    @test_allocations_splitform(semi, sol, allocs)
+
+Test that the memory allocations of `DispersiveShallowWater.rhs_split_1!` 
+and `DispersiveShallowWater.rhs_split_2!` are below `allocs`
+(e.g., from type instabilities).
+"""
+macro test_allocations_splitform(semi, sol, allocs)
+    quote
+        t = $sol.t[end]
+        q = $sol.u[end]
+        dq = similar(q)
+        a1 = @allocated DispersiveShallowWater.rhs_split_1!(dq, q, $semi, t)
+        a2 = @allocated DispersiveShallowWater.rhs_split_2!(dq, q, $semi, t)
+        @test (a1 + a2) < $allocs
+    end
+end
+
 macro test_nowarn_mod(expr, additional_ignore_content = [])
     quote
         add_to_additional_ignore_content = [