Towards partitioned RHS's (1D): TreeMesh only

src/semidiscretization/semidiscretization_hyperbolic.jl

@@ -409,4 +409,25 @@ function rhs!(du_ode, u_ode, semi::SemidiscretizationHyperbolic, t)
 
     return nothing
 end
+
+# `rhs!` for partitioned Runge-Kutta methods, such as the Paired Explicit Runge-Kutta (PERK) methods
+function rhs!(du_ode, u_ode, semi::SemidiscretizationHyperbolic, t,
+              element_indices, interface_indices, boundary_indices, mortar_indices)
+    @unpack mesh, equations, boundary_conditions, source_terms, solver, cache = semi
+
+    u = wrap_array(u_ode, mesh, equations, solver, cache)
+    du = wrap_array(du_ode, mesh, equations, solver, cache)
+
+    # TODO: Taal decide, do we need to pass the mesh?
+    time_start = time_ns()
+    @trixi_timeit timer() "rhs! (part.)" rhs!(du, u, t, mesh, equations,
+                                              boundary_conditions, source_terms,
+                                              solver, cache,
+                                              element_indices, interface_indices,
+                                              boundary_indices, mortar_indices)
+    runtime = time_ns() - time_start
+    put!(semi.performance_counter, runtime)
+
+    return nothing
+end
 end # @muladd

src/solvers/dgsem_tree/dg_1d.jl

@@ -70,33 +70,37 @@ end
 function rhs!(du, u, t,
               mesh::TreeMesh{1}, equations,
               boundary_conditions, source_terms::Source,
-              dg::DG, cache) where {Source}
+              dg::DG, cache,
+              element_indices = eachelement(dg, cache),
+              interface_indices = eachinterface(dg, cache),
+              boundary_indices = eachboundary(dg, cache),
+              mortar_indices = nothing) where {Source}
     # Reset du
-    @trixi_timeit timer() "reset ∂u/∂t" reset_du!(du, dg, cache)
+    @trixi_timeit timer() "reset ∂u/∂t" reset_du!(du, dg, cache, element_indices)
 
     # Calculate volume integral
     @trixi_timeit timer() "volume integral" begin
         calc_volume_integral!(du, u, mesh,
                               have_nonconservative_terms(equations), equations,
-                              dg.volume_integral, dg, cache)
+                              dg.volume_integral, dg, cache, element_indices)
     end
 
     # Prolong solution to interfaces
     @trixi_timeit timer() "prolong2interfaces" begin
-        prolong2interfaces!(cache, u, mesh, equations, dg)
+        prolong2interfaces!(cache, u, mesh, equations, dg, interface_indices)
     end
 
     # Calculate interface fluxes
     @trixi_timeit timer() "interface flux" begin
         calc_interface_flux!(cache.elements.surface_flux_values, mesh,
                              have_nonconservative_terms(equations), equations,
-                             dg.surface_integral, dg, cache)
+                             dg.surface_integral, dg, cache, interface_indices)
     end
 
     # Prolong solution to boundaries
     @trixi_timeit timer() "prolong2boundaries" begin
         prolong2boundaries!(cache, u, mesh, equations,
-                            dg.surface_integral, dg)
+                            dg.surface_integral, dg, boundary_indices)
     end
 
     # Calculate boundary fluxes
@@ -108,15 +112,16 @@ function rhs!(du, u, t,
     # Calculate surface integrals
     @trixi_timeit timer() "surface integral" begin
         calc_surface_integral!(du, u, mesh, equations,
-                               dg.surface_integral, dg, cache)
+                               dg.surface_integral, dg, cache, element_indices)
     end
 
     # Apply Jacobian from mapping to reference element
-    @trixi_timeit timer() "Jacobian" apply_jacobian!(du, mesh, equations, dg, cache)
+    @trixi_timeit timer() "Jacobian" apply_jacobian!(du, mesh, equations, dg, cache,
+                                                     element_indices)
 
     # Calculate source terms
     @trixi_timeit timer() "source terms" begin
-        calc_sources!(du, u, t, source_terms, equations, dg, cache)
+        calc_sources!(du, u, t, source_terms, equations, dg, cache, element_indices)
     end
 
     return nothing
@@ -245,11 +250,12 @@ end
 function calc_volume_integral!(du, u, mesh::Union{TreeMesh{1}, StructuredMesh{1}},
                                have_nonconservative_terms, equations,
                                volume_integral::VolumeIntegralPureLGLFiniteVolumeO2,
-                               dg::DGSEM, cache)
+                               dg::DGSEM, cache,
+                               element_indices = eachelement(dg, cache))
     @unpack x_interfaces, volume_flux_fv, reconstruction_mode, slope_limiter = volume_integral
 
     # Calculate LGL second-order FV volume integral
-    @threaded for element in eachelement(dg, cache)
+    @threaded for element in element_indices
         fvO2_kernel!(du, u, mesh,
                      have_nonconservative_terms, equations,
                      volume_flux_fv, dg, cache, element,
@@ -408,12 +414,13 @@ end
     return nothing
 end
 
-function prolong2interfaces!(cache, u, mesh::TreeMesh{1}, equations, dg::DG)
+function prolong2interfaces!(cache, u, mesh::TreeMesh{1}, equations, dg::DG,
+                             interface_indices = eachinterface(dg, cache))
     @unpack interfaces = cache
     @unpack neighbor_ids = interfaces
     interfaces_u = interfaces.u
 
-    @threaded for interface in eachinterface(dg, cache)
+    @threaded for interface in interface_indices
         left_element = neighbor_ids[1, interface]
         right_element = neighbor_ids[2, interface]
 
@@ -430,11 +437,12 @@ end
 function calc_interface_flux!(surface_flux_values,
                               mesh::TreeMesh{1},
                               have_nonconservative_terms::False, equations,
-                              surface_integral, dg::DG, cache)
+                              surface_integral, dg::DG, cache,
+                              interface_indices = eachinterface(dg, cache))
     @unpack surface_flux = surface_integral
     @unpack u, neighbor_ids, orientations = cache.interfaces
 
-    @threaded for interface in eachinterface(dg, cache)
+    @threaded for interface in interface_indices
         # Get neighboring elements
         left_id = neighbor_ids[1, interface]
         right_id = neighbor_ids[2, interface]
@@ -450,6 +458,8 @@ function calc_interface_flux!(surface_flux_values,
 
         # Copy flux to left and right element storage
         for v in eachvariable(equations)
+            # TODO: Accumulate fluxes, i.e., `+=`? 
+            # Reset would then be handled by `reset_du!`
             surface_flux_values[v, left_direction, left_id] = flux[v]
             surface_flux_values[v, right_direction, right_id] = flux[v]
         end
@@ -461,11 +471,12 @@ end
 function calc_interface_flux!(surface_flux_values,
                               mesh::TreeMesh{1},
                               have_nonconservative_terms::True, equations,
-                              surface_integral, dg::DG, cache)
+                              surface_integral, dg::DG, cache,
+                              interface_indices = eachinterface(dg, cache))
     surface_flux, nonconservative_flux = surface_integral.surface_flux
     @unpack u, neighbor_ids, orientations = cache.interfaces
 
-    @threaded for interface in eachinterface(dg, cache)
+    @threaded for interface in interface_indices
         # Get neighboring elements
         left_id = neighbor_ids[1, interface]
         right_id = neighbor_ids[2, interface]
@@ -501,11 +512,12 @@ function calc_interface_flux!(surface_flux_values,
 end
 
 function prolong2boundaries!(cache, u,
-                             mesh::TreeMesh{1}, equations, surface_integral, dg::DG)
+                             mesh::TreeMesh{1}, equations, surface_integral, dg::DG,
+                             boundary_indices = eachboundary(dg, cache))
     @unpack boundaries = cache
     @unpack neighbor_sides = boundaries
 
-    @threaded for boundary in eachboundary(dg, cache)
+    @threaded for boundary in boundary_indices
         element = boundaries.neighbor_ids[boundary]
 
         # boundary in x-direction
@@ -616,7 +628,8 @@ function calc_boundary_flux_by_direction!(surface_flux_values::AbstractArray{<:A
 end
 
 function calc_surface_integral!(du, u, mesh::Union{TreeMesh{1}, StructuredMesh{1}},
-                                equations, surface_integral, dg::DGSEM, cache)
+                                equations, surface_integral, dg::DGSEM, cache,
+                                element_indices = eachelement(dg, cache))
     @unpack boundary_interpolation = dg.basis
     @unpack surface_flux_values = cache.elements
 
@@ -626,7 +639,7 @@ function calc_surface_integral!(du, u, mesh::Union{TreeMesh{1}, StructuredMesh{1
     # into FMAs (see comment at the top of the file).
     factor_1 = boundary_interpolation[1, 1]
     factor_2 = boundary_interpolation[nnodes(dg), 2]
-    @threaded for element in eachelement(dg, cache)
+    @threaded for element in element_indices
         for v in eachvariable(equations)
             # surface at -x
             du[v, 1, element] = (du[v, 1, element] -
@@ -642,10 +655,11 @@ function calc_surface_integral!(du, u, mesh::Union{TreeMesh{1}, StructuredMesh{1
 end
 
 function apply_jacobian!(du, mesh::TreeMesh{1},
-                         equations, dg::DG, cache)
+                         equations, dg::DG, cache,
+                         element_indices = eachelement(dg, cache))
     @unpack inverse_jacobian = cache.elements
 
-    @threaded for element in eachelement(dg, cache)
+    @threaded for element in element_indices
         factor = -inverse_jacobian[element]
 
         for i in eachnode(dg)
@@ -660,15 +674,17 @@ end
 
 # Need dimension specific version to avoid error at dispatching
 function calc_sources!(du, u, t, source_terms::Nothing,
-                       equations::AbstractEquations{1}, dg::DG, cache)
+                       equations::AbstractEquations{1}, dg::DG, cache,
+                       element_indices = eachelement(dg, cache))
     return nothing
 end
 
 function calc_sources!(du, u, t, source_terms,
-                       equations::AbstractEquations{1}, dg::DG, cache)
+                       equations::AbstractEquations{1}, dg::DG, cache,
+                       element_indices = eachelement(dg, cache))
     @unpack node_coordinates = cache.elements
 
-    @threaded for element in eachelement(dg, cache)
+    @threaded for element in element_indices
         for i in eachnode(dg)
             u_local = get_node_vars(u, equations, dg, i, element)
             x_local = get_node_coords(node_coordinates, equations, dg,