Introduce more index-based functions to enable more flexibility in handling system pair specific interactions

Author: svchb

src/general/neighborhood_search.jl

@@ -102,42 +102,43 @@ end
 end
 
 @inline function get_neighborhood_search(system, semi)
-    (; neighborhood_searches) = semi
+    system_index = system_indices(system, semi)
+
+    return get_neighborhood_search(system, semi, system_index)
+end
 
+@inline function get_neighborhood_search(system, neighbor_system, semi)
     system_index = system_indices(system, semi)
+    neighbor_index = system_indices(neighbor_system, semi)
+
+    return get_neighborhood_search(system, neighbor_system, semi, system_index, neighbor_index)
+end
 
-    return neighborhood_searches[system_index][system_index]
+@inline function get_neighborhood_search(system, semi, system_index::Integer)
+    return semi.neighborhood_searches[system_index][system_index]
 end
 
-@inline function get_neighborhood_search(system::TotalLagrangianSPHSystem, semi)
+@inline function get_neighborhood_search(system::TotalLagrangianSPHSystem, semi,
+                                         system_index::Integer)
     # For TLSPH, use the specialized self-interaction neighborhood search
     # for finding neighbors in the initial configuration.
     return system.self_interaction_nhs
 end
 
+@inline function get_neighborhood_search(system, neighbor_system, semi,
+                                         system_index::Integer, neighbor_index::Integer)
+    return semi.neighborhood_searches[system_index][neighbor_index]
+end
 @inline function get_neighborhood_search(system::TotalLagrangianSPHSystem,
-                                         neighbor_system::TotalLagrangianSPHSystem, semi)
-    (; neighborhood_searches) = semi
-
-    system_index = system_indices(system, semi)
-    neighbor_index = system_indices(neighbor_system, semi)
-
+                                         neighbor_system::TotalLagrangianSPHSystem, semi,
+                                         system_index::Integer, neighbor_index::Integer)
     if system_index == neighbor_index
         # For TLSPH, use the specialized self-interaction neighborhood search
         # for finding neighbors in the initial configuration.
         return system.self_interaction_nhs
     end
 
-    return neighborhood_searches[system_index][neighbor_index]
-end
-
-@inline function get_neighborhood_search(system, neighbor_system, semi)
-    (; neighborhood_searches) = semi
-
-    system_index = system_indices(system, semi)
-    neighbor_index = system_indices(neighbor_system, semi)
-
-    return neighborhood_searches[system_index][neighbor_index]
+    return semi.neighborhood_searches[system_index][neighbor_index]
 end
 
 function initialize_neighborhood_searches!(semi)
@@ -150,11 +151,13 @@ function initialize_neighborhood_searches!(semi)
     return semi
 end
 
-function initialize_neighborhood_search!(semi, system, neighbor)
+function initialize_neighborhood_search!(semi, system, neighbor,
+                                         system_index::Integer, neighbor_index::Integer)
     # TODO Initialize after adapting to the GPU.
     # Currently, this cannot use `semi.parallelization_backend`
     # because data is still on the CPU.
-    PointNeighbors.initialize!(get_neighborhood_search(system, neighbor, semi),
+    PointNeighbors.initialize!(get_neighborhood_search(system, neighbor, semi,
+                                                       system_index, neighbor_index),
                                initial_coordinates(system),
                                initial_coordinates(neighbor),
                                eachindex_y=each_active_particle(neighbor),
@@ -164,19 +167,30 @@ function initialize_neighborhood_search!(semi, system, neighbor)
 end
 
 function initialize_neighborhood_search!(semi, system::TotalLagrangianSPHSystem,
-                                         neighbor::TotalLagrangianSPHSystem)
+                                         neighbor::TotalLagrangianSPHSystem,
+                                         system_index::Integer, neighbor_index::Integer)
     # For TLSPH, the self-interaction NHS is already initialized in the system constructor
     return semi
 end
 
+function initialize_neighborhood_search!(semi, system, neighbor)
+    system_index = system_indices(system, semi)
+    neighbor_index = system_indices(neighbor, semi)
+
+    return initialize_neighborhood_search!(semi, system, neighbor,
+                                           system_index, neighbor_index)
+end
+
+# === Neighborhood search updates (per-system) ===
 function update_nhs!(semi, u_ode)
     # Update NHS for each pair of systems
-    foreach_system(semi) do system
-        u_system = wrap_u(u_ode, system, semi)
+    foreach_system_indexed(semi) do system_index, system
+        u_system = wrap_u(u_ode, system, semi, system_index)
 
-        foreach_system(semi) do neighbor
-            u_neighbor = wrap_u(u_ode, neighbor, semi)
-            neighborhood_search = get_neighborhood_search(system, neighbor, semi)
+        foreach_system_indexed(semi) do neighbor_index, neighbor
+            u_neighbor = wrap_u(u_ode, neighbor, semi, neighbor_index)
+            neighborhood_search = get_neighborhood_search(system, neighbor, semi,
+                                                          system_index, neighbor_index)
 
             update_nhs!(neighborhood_search, system, neighbor, u_system, u_neighbor, semi)
         end

src/general/semidiscretization.jl

@@ -171,6 +171,38 @@ end
 
 @inline foreach_system(f, systems) = foreach_noalloc(f, systems)
 
+# This is just for readability to loop over all systems with indices without allocations.
+@inline function foreach_system_indexed(f, semi::Union{NamedTuple, Semidiscretization})
+    return foreach_system_indexed(f, semi.systems)
+end
+
+@inline function foreach_system_indexed(f, systems::Tuple)
+    indices = ntuple(identity, Val(length(systems)))
+
+    return foreach_noalloc(indices, systems) do (index, system)
+        f(index, system)
+    end
+end
+
+# This is just for readability to loop over all systems with wrapped arrays.
+@inline function foreach_system_wrapped(f, semi::Union{NamedTuple, Semidiscretization},
+                                        v_ode, u_ode)
+    return foreach_system_indexed(semi) do system_index, system
+        f(system,
+          wrap_v(v_ode, system, semi, system_index),
+          wrap_u(u_ode, system, semi, system_index))
+    end
+end
+
+@inline function foreach_system_wrapped(f, semi::Union{NamedTuple, Semidiscretization},
+                                        dv_ode, v_ode, u_ode)
+    return foreach_system_indexed(semi) do system_index, system
+        f(system,
+          wrap_v(dv_ode, system, semi, system_index),
+          wrap_v(v_ode, system, semi, system_index),
+          wrap_u(u_ode, system, semi, system_index))
+    end
+end
 """
     semidiscretize(semi, tspan; reset_threads=true)
 
@@ -334,24 +366,40 @@ end
 # We have to pass `system` here for type stability,
 # since the type of `system` determines the return type.
 @inline function wrap_v(v_ode, system, semi)
-    (; ranges_v) = semi
+    return wrap_v(v_ode, system, semi, system_indices(system, semi))
+end
 
-    range = ranges_v[system_indices(system, semi)]
+@inline function wrap_v(v_ode, system, semi, system_index::Integer)
+    return wrap_v(v_ode, system, semi.ranges_v[system_index])
+end
 
-    @boundscheck @assert length(range) ==
-                         v_nvariables(system) * n_integrated_particles(system)
+@inline function wrap_v(v_ode, system, range::AbstractUnitRange)
+    @boundscheck begin
+        expected = v_nvariables(system) * n_integrated_particles(system)
+        range_length = length(range)
+        range_length == expected ||
+            throw(DimensionMismatch("v range length $range_length does not match expected $expected"))
+    end
 
     return wrap_array(v_ode, range,
                       (StaticInt(v_nvariables(system)), n_integrated_particles(system)))
 end
 
 @inline function wrap_u(u_ode, system, semi)
-    (; ranges_u) = semi
+    return wrap_u(u_ode, system, semi, system_indices(system, semi))
+end
 
-    range = ranges_u[system_indices(system, semi)]
+@inline function wrap_u(u_ode, system, semi, system_index::Integer)
+    return wrap_u(u_ode, system, semi.ranges_u[system_index])
+end
 
-    @boundscheck @assert length(range) ==
-                         u_nvariables(system) * n_integrated_particles(system)
+@inline function wrap_u(u_ode, system, range::AbstractUnitRange)
+    @boundscheck begin
+        expected = u_nvariables(system) * n_integrated_particles(system)
+        range_length = length(range)
+        range_length == expected ||
+            throw(DimensionMismatch("u range length $range_length does not match expected $expected"))
+    end
 
     return wrap_array(u_ode, range,
                       (StaticInt(u_nvariables(system)), n_integrated_particles(system)))
@@ -639,19 +687,18 @@ end
 
 function system_interaction!(dv_ode, v_ode, u_ode, semi)
     # Call `interact!` for each pair of systems
-    foreach_system(semi) do system
-        foreach_system(semi) do neighbor
+    foreach_system_indexed(semi) do system_index, system
+        foreach_system_indexed(semi) do neighbor_index, neighbor
             # Construct string for the interactions timer.
             # Avoid allocations from string construction when no timers are used.
             if timeit_debug_enabled()
-                system_index = system_indices(system, semi)
-                neighbor_index = system_indices(neighbor, semi)
                 timer_str = "$(timer_name(system))$system_index-$(timer_name(neighbor))$neighbor_index"
             else
                 timer_str = ""
             end
 
-            interact!(dv_ode, v_ode, u_ode, system, neighbor, semi, timer_str=timer_str)
+            interact!(dv_ode, v_ode, u_ode, system, neighbor, semi,
+                      system_index, neighbor_index; timer_str=timer_str)
         end
     end
 
@@ -663,12 +710,21 @@ end
 # dv_ode, du_ode = copy(sol.u[end]).x; v_ode, u_ode = copy(sol.u[end]).x;
 # @btime TrixiParticles.interact!($dv_ode, $v_ode, $u_ode, $fluid_system, $fluid_system, $semi);
 @inline function interact!(dv_ode, v_ode, u_ode, system, neighbor, semi; timer_str="")
-    dv = wrap_v(dv_ode, system, semi)
-    v_system = wrap_v(v_ode, system, semi)
-    u_system = wrap_u(u_ode, system, semi)
+    system_index = system_indices(system, semi)
+    neighbor_index = system_indices(neighbor, semi)
+
+    return interact!(dv_ode, v_ode, u_ode, system, neighbor, semi,
+                     system_index, neighbor_index; timer_str=timer_str)
+end
+
+@inline function interact!(dv_ode, v_ode, u_ode, system, neighbor, semi,
+                           system_index::Integer, neighbor_index::Integer; timer_str="")
+    dv = wrap_v(dv_ode, system, semi, system_index)
+    v_system = wrap_v(v_ode, system, semi, system_index)
+    u_system = wrap_u(u_ode, system, semi, system_index)
 
-    v_neighbor = wrap_v(v_ode, neighbor, semi)
-    u_neighbor = wrap_u(u_ode, neighbor, semi)
+    v_neighbor = wrap_v(v_ode, neighbor, semi, neighbor_index)
+    u_neighbor = wrap_u(u_ode, neighbor, semi, neighbor_index)
 
     @trixi_timeit timer() timer_str begin
         interact!(dv, v_system, u_system, v_neighbor, u_neighbor, system, neighbor, semi)

src/io/io.jl

@@ -36,8 +36,8 @@ function create_meta_data_dict(callback, integrator)
 
     systems = Dict{String, Any}()
     foreach_system(semi) do system
-        idx = system_indices(system, semi)
-        name = add_underscore_to_optional_prefix(prefix) * names[idx]
+        system_index = system_indices(system, semi)
+        name = add_underscore_to_optional_prefix(prefix) * names[system_index]
 
         system_data = Dict{String, Any}()
         add_system_data!(system_data, system)