Apply custom source term (that depends on du) with a custom RHS

examples/p4est_2d_dgsem/elixir_euler_cubed_sphere_shell.jl

@@ -67,13 +67,41 @@ function source_terms_convert_to_linear_advection(u, du, x, t,
     return SVector(0.0, s2, s3, s4, s5)
 end
 
+# Custom RHS that applies a source term that depends on du (used to convert the 3D Euler equations into the 3D linear advection
+# equations with position-dependent advection speed)
+function rhs_semi_custom!(du_ode, u_ode, semi, t)
+    # Compute standard Trixi RHS
+    Trixi.rhs!(du_ode, u_ode, semi, t)
+
+    # Now apply the custom source term
+    Trixi.@trixi_timeit Trixi.timer() "custom source term" begin
+        @unpack solver, equations, cache = semi
+        @unpack node_coordinates = cache.elements
+
+        # Wrap the solution and RHS
+        u = Trixi.wrap_array(u_ode, semi)
+        du = Trixi.wrap_array(du_ode, semi)
+
+        Trixi.@threaded for element in eachelement(solver, cache)
+            for j in eachnode(solver), i in eachnode(solver)
+                u_local = Trixi.get_node_vars(u, equations, solver, i, j, element)
+                du_local = Trixi.get_node_vars(du, equations, solver, i, j, element)
+                x_local = Trixi.get_node_coords(node_coordinates, equations, solver,
+                                                i, j, element)
+                source = source_terms_convert_to_linear_advection(u_local, du_local,
+                                                                  x_local, t, equations)
+                Trixi.add_to_node_vars!(du, source, equations, solver, i, j, element)
+            end
+        end
+    end
+end
+
 initial_condition = initial_condition_advection_sphere
 
 mesh = Trixi.P4estMeshCubedSphere2D(5, 1.0, polydeg = polydeg, initial_refinement_level = 0)
 
 # A semidiscretization collects data structures and functions for the spatial discretization
-semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
-                                    source_terms = source_terms_convert_to_linear_advection)
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)
 
 ###############################################################################
 # ODE solvers, callbacks etc.
@@ -82,6 +110,12 @@ semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
 tspan = (0.0, pi)
 ode = semidiscretize(semi, tspan)
 
+# Create custom discretization that runs with the custom RHS
+ode_semi_custom = ODEProblem(rhs_semi_custom!,
+                             ode.u0,
+                             ode.tspan,
+                             semi)
+
 # At the beginning of the main loop, the SummaryCallback prints a summary of the simulation setup
 # and resets the timers
 summary_callback = SummaryCallback()
@@ -106,7 +140,7 @@ callbacks = CallbackSet(summary_callback, analysis_callback, save_solution,
 # run the simulation
 
 # OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed callbacks
-sol = solve(ode, CarpenterKennedy2N54(williamson_condition = false),
+sol = solve(ode_semi_custom, CarpenterKennedy2N54(williamson_condition = false),
             dt = 1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
             save_everystep = false, callback = callbacks);
 

src/solvers/dgsem_tree/dg_2d.jl

@@ -172,11 +172,7 @@ function rhs!(du, u, t,
 
     # Calculate source terms
     @trixi_timeit timer() "source terms" begin
-        if size(cache.elements.node_coordinates, 1) == 3
-            calc_extended_sources!(du, u, t, source_terms, equations, dg, cache)
-        else
-            calc_sources!(du, u, t, source_terms, equations, dg, cache)
-        end
+        calc_sources!(du, u, t, source_terms, equations, dg, cache)
     end
 
     return nothing
@@ -1169,29 +1165,4 @@ function calc_sources!(du, u, t, source_terms,
 
     return nothing
 end
-
-function calc_extended_sources!(du, u, t, source_terms::Nothing,
-                                equations::AbstractEquations{3}, dg::DG, cache)
-    return nothing
-end
-
-# Source term that depends on du for 3D equations
-# TODO: Is there a better way to do this?
-function calc_extended_sources!(du, u, t, source_terms,
-                                equations::AbstractEquations{3}, dg::DG, cache)
-    @unpack node_coordinates = cache.elements
-
-    @threaded for element in eachelement(dg, cache)
-        for j in eachnode(dg), i in eachnode(dg)
-            u_local = get_node_vars(u, equations, dg, i, j, element)
-            du_local = get_node_vars(du, equations, dg, i, j, element)
-            x_local = get_node_coords(node_coordinates, equations, dg,
-                                      i, j, element)
-            source = source_terms(u_local, du_local, x_local, t, equations)
-            add_to_node_vars!(du, source, equations, dg, i, j, element)
-        end
-    end
-
-    return nothing
-end
 end # @muladd