Refactor neighborhood search and update functions for improved clarity and performance. Remove redundant checks and streamline system interactions. Update function signatures for consistency across modules.

Author: svchb

src/callbacks/split_integration.jl

@@ -218,10 +218,10 @@ function update_systems_split!(semi, v_ode, u_ode, t)
     # a lot of the speedup that we can gain with split integration.
     # We assume that the TLSPH particles move so little during the substeps
     # that the extrapolated pressure/density values can be treated as constant.
-    update_systems!(v_ode, u_ode, semi, t;
-                    update_nhs=false,
-                    update_boundary_interpolation=false,
-                    update_inter_system=false)
+    return update_systems!(v_ode, u_ode, semi, t;
+                           update_nhs=false,
+                           update_implicit_sph=false,
+                           update_boundary_interpolation=false)
 end
 
 function system_interaction_split!(dv_ode_split, v_ode, u_ode, semi,

src/general/neighborhood_search.jl

@@ -1,4 +1,3 @@
-# === PointNeighbors integration ===
 # Loop over all pairs of particles and neighbors within the kernel cutoff.
 # `f(particle, neighbor, pos_diff, distance)` is called for every particle-neighbor pair.
 # By default, loop over `eachparticle(system)`.
@@ -11,15 +10,44 @@ function PointNeighbors.foreach_point_neighbor(f, system, neighbor_system,
                            points, parallelization_backend)
 end
 
-# === Compact support selection ===
-# -- Generic
+# For non-TLSPH systems, do nothing
+function initialize_self_interaction_nhs(system, neighborhood_search,
+                                         parallelization_backend)
+    return system
+end
+
+function create_neighborhood_search(::Nothing, system, neighbor)
+    nhs = TrivialNeighborhoodSearch{ndims(system)}()
+
+    return create_neighborhood_search(nhs, system, neighbor)
+end
+
+# Avoid method ambiguity
+function create_neighborhood_search(::Nothing, system::TotalLagrangianSPHSystem,
+                                    neighbor::TotalLagrangianSPHSystem)
+    nhs = TrivialNeighborhoodSearch{ndims(system)}()
+
+    return create_neighborhood_search(nhs, system, neighbor)
+end
+
+function create_neighborhood_search(neighborhood_search, system, neighbor)
+    return copy_neighborhood_search(neighborhood_search, compact_support(system, neighbor),
+                                    nparticles(neighbor))
+end
+
+function create_neighborhood_search(neighborhood_search, system::TotalLagrangianSPHSystem,
+                                    neighbor::TotalLagrangianSPHSystem)
+    # TLSPH self-interaction is using a specialized neighborhood search
+    return copy_neighborhood_search(neighborhood_search, zero(eltype(system)),
+                                    nparticles(neighbor))
+end
+
 @inline function compact_support(system, neighbor)
     (; smoothing_kernel) = system
     # TODO: Variable search radius for NHS?
     return compact_support(smoothing_kernel, initial_smoothing_length(system))
 end
 
-# -- Open boundary systems
 @inline function compact_support(system::OpenBoundarySystem,
                                  neighbor::OpenBoundarySystem)
     # This NHS is never used
@@ -32,7 +60,6 @@ end
     return compact_support(system.fluid_system, neighbor.fluid_system)
 end
 
-# -- DEM boundaries
 @inline function compact_support(system::BoundaryDEMSystem, neighbor::BoundaryDEMSystem)
     # This NHS is never used
     return zero(eltype(system))
@@ -43,14 +70,12 @@ end
     return compact_support(neighbor, system)
 end
 
-# -- TLSPH systems
 @inline function compact_support(system::TotalLagrangianSPHSystem,
                                  neighbor::TotalLagrangianSPHSystem)
     (; smoothing_kernel, smoothing_length) = system
     return compact_support(smoothing_kernel, smoothing_length)
 end
 
-# -- Boundary models
 @inline function compact_support(system::Union{TotalLagrangianSPHSystem,
                                                WallBoundarySystem},
                                  neighbor)
@@ -76,79 +101,43 @@ end
     return compact_support(smoothing_kernel, smoothing_length)
 end
 
-# === Neighborhood search creation ===
-function create_neighborhood_search(::Nothing, system, neighbor)
-    nhs = TrivialNeighborhoodSearch{ndims(system)}()
-
-    return create_neighborhood_search(nhs, system, neighbor)
-end
-
-# Avoid method ambiguity
-function create_neighborhood_search(::Nothing, system::TotalLagrangianSPHSystem,
-                                    neighbor::TotalLagrangianSPHSystem)
-    nhs = TrivialNeighborhoodSearch{ndims(system)}()
-
-    return create_neighborhood_search(nhs, system, neighbor)
-end
-
-function create_neighborhood_search(neighborhood_search, system, neighbor)
-    return copy_neighborhood_search(neighborhood_search, compact_support(system, neighbor),
-                                    nparticles(neighbor))
-end
-
-function create_neighborhood_search(neighborhood_search, system::TotalLagrangianSPHSystem,
-                                    neighbor::TotalLagrangianSPHSystem)
-    # TLSPH self-interaction is using a specialized neighborhood search
-    return copy_neighborhood_search(neighborhood_search, zero(eltype(system)),
-                                    nparticles(neighbor))
-end
-
-# === Neighborhood search lookup ===
 @inline function get_neighborhood_search(system, semi)
-    system_index = system_indices(system, semi)
+    (; neighborhood_searches) = 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)
+    return neighborhood_searches[system_index][system_index]
 end
 
-@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,
-                                         system_index::Integer)
+@inline function get_neighborhood_search(system::TotalLagrangianSPHSystem, semi)
     # 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,
-                                         system_index::Integer, neighbor_index::Integer)
+                                         neighbor_system::TotalLagrangianSPHSystem, semi)
+    (; neighborhood_searches) = semi
+
+    system_index = system_indices(system, semi)
+    neighbor_index = system_indices(neighbor_system, semi)
+
     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 semi.neighborhood_searches[system_index][neighbor_index]
+    return neighborhood_searches[system_index][neighbor_index]
 end
-# === Initialization ===
-# For non-TLSPH systems, do nothing
-function initialize_self_interaction_nhs(system, neighborhood_search,
-                                         parallelization_backend)
-    return system
+
+@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]
 end
 
 function initialize_neighborhood_searches!(semi)
@@ -161,13 +150,11 @@ function initialize_neighborhood_searches!(semi)
     return semi
 end
 
-function initialize_neighborhood_search!(semi, system, neighbor,
-                                         system_index::Integer, neighbor_index::Integer)
+function initialize_neighborhood_search!(semi, system, neighbor)
     # 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,
-                                                       system_index, neighbor_index),
+    PointNeighbors.initialize!(get_neighborhood_search(system, neighbor, semi),
                                initial_coordinates(system),
                                initial_coordinates(neighbor),
                                eachindex_y=each_active_particle(neighbor),
@@ -177,39 +164,27 @@ function initialize_neighborhood_search!(semi, system, neighbor,
 end
 
 function initialize_neighborhood_search!(semi, system::TotalLagrangianSPHSystem,
-                                         neighbor::TotalLagrangianSPHSystem,
-                                         system_index::Integer, neighbor_index::Integer)
+                                         neighbor::TotalLagrangianSPHSystem)
     # 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_indexed(semi) do system_index, system
-        u_system = wrap_u(u_ode, system, semi, system_index)
+    foreach_system(semi) do system
+        u_system = wrap_u(u_ode, system, 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)
+        foreach_system(semi) do neighbor
+            u_neighbor = wrap_u(u_ode, neighbor, semi)
+            neighborhood_search = get_neighborhood_search(system, neighbor, semi)
 
             update_nhs!(neighborhood_search, system, neighbor, u_system, u_neighbor, semi)
         end
     end
 end
 
-# === Neighborhood search updates (per-pair dispatch) ===
+# NHS updates
 # To prevent hard-to-find bugs, there is no default version
-# -- Fluid / structure interactions
 function update_nhs!(neighborhood_search,
                      system::AbstractFluidSystem,
                      neighbor::Union{AbstractFluidSystem, TotalLagrangianSPHSystem},
@@ -221,7 +196,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(true, true), eachindex_y=each_active_particle(neighbor))
 end
 
-# -- Fluid / wall interactions
 function update_nhs!(neighborhood_search,
                      system::Union{AbstractFluidSystem,
                                    OpenBoundarySystem{<:BoundaryModelDynamicalPressureZhang}},
@@ -234,7 +208,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(true, neighbor.ismoving[]))
 end
 
-# -- Open boundary interactions
 function update_nhs!(neighborhood_search,
                      system::AbstractFluidSystem, neighbor::OpenBoundarySystem,
                      u_system, u_neighbor, semi)
@@ -291,7 +264,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(true, true), eachindex_y=each_active_particle(neighbor))
 end
 
-# -- Open boundary combinations that are never used
 function update_nhs!(neighborhood_search,
                      system::OpenBoundarySystem, neighbor::TotalLagrangianSPHSystem,
                      u_system, u_neighbor, semi)
@@ -306,7 +278,6 @@ function update_nhs!(neighborhood_search,
     return neighborhood_search
 end
 
-# -- TLSPH interactions
 function update_nhs!(neighborhood_search,
                      system::TotalLagrangianSPHSystem, neighbor::AbstractFluidSystem,
                      u_system, u_neighbor, semi)
@@ -317,6 +288,14 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(true, true), eachindex_y=each_active_particle(neighbor))
 end
 
+function update_nhs!(neighborhood_search,
+                     system::TotalLagrangianSPHSystem, neighbor::TotalLagrangianSPHSystem,
+                     u_system, u_neighbor, semi)
+    # Don't update. This NHS is never used.
+    # TLSPH systems have their own self-interaction NHS.
+    return neighborhood_search
+end
+
 function update_nhs!(neighborhood_search,
                      system::TotalLagrangianSPHSystem, neighbor::WallBoundarySystem,
                      u_system, u_neighbor, semi)
@@ -328,15 +307,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(true, neighbor.ismoving[]))
 end
 
-function update_nhs!(neighborhood_search,
-                     system::TotalLagrangianSPHSystem, neighbor::TotalLagrangianSPHSystem,
-                     u_system, u_neighbor, semi)
-    # Don't update. This NHS is never used.
-    # TLSPH systems have their own self-interaction NHS.
-    return neighborhood_search
-end
-
-# -- Wall dummy particle interactions
 # This function is the same as the one below to avoid ambiguous dispatch when using `Union`
 function update_nhs!(neighborhood_search,
                      system::WallBoundarySystem{<:BoundaryModelDummyParticles},
@@ -391,7 +361,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(system.ismoving[], true))
 end
 
-# -- Wall / wall interactions
 function update_nhs!(neighborhood_search,
                      system::WallBoundarySystem{<:BoundaryModelDummyParticles},
                      neighbor::WallBoundarySystem,
@@ -404,7 +373,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(system.ismoving[], neighbor.ismoving[]))
 end
 
-# -- DEM interactions
 function update_nhs!(neighborhood_search,
                      system::DEMSystem, neighbor::DEMSystem,
                      u_system, u_neighbor, semi)
@@ -425,7 +393,6 @@ function update_nhs!(neighborhood_search,
             semi, points_moving=(true, false))
 end
 
-# -- Combinations that are never used
 function update_nhs!(neighborhood_search,
                      system::WallBoundarySystem,
                      neighbor::AbstractFluidSystem,
@@ -450,7 +417,7 @@ function update_nhs!(neighborhood_search,
     return neighborhood_search
 end
 
-# === PointNeighbors forwarding ===
+# Forward to PointNeighbors.jl
 function update!(neighborhood_search, x, y, semi; points_moving=(true, true),
                  eachindex_y=axes(y, 2))
     PointNeighbors.update!(neighborhood_search, x, y; points_moving, eachindex_y,