Nhs handler

src/TrixiParticles.jl

@@ -40,7 +40,8 @@ using TrixiBase: @trixi_timeit, timer, timeit_debug_enabled,
                                 SerialIncrementalUpdate,
                                 FullGridCellList, DictionaryCellList,
                                 SerialBackend, PolyesterBackend, ThreadsStaticBackend,
-                                ThreadsDynamicBackend, default_backend
+                                ThreadsDynamicBackend, default_backend,
+                                VariableSearchRadiusNHS
 using PointNeighbors: PointNeighbors, foreach_point_neighbor, copy_neighborhood_search,
                       @threaded
 using WriteVTK: vtk_grid, MeshCell, VTKCellTypes, paraview_collection, vtk_save
@@ -112,5 +113,7 @@ export SurfaceTensionAkinci, CohesionForceAkinci, SurfaceTensionMorris,
        SurfaceTensionMomentumMorris
 export ColorfieldSurfaceNormal
 export SymplecticPositionVerlet
+export GridNHSHandler, PairsNHSHandler, VariableNHSHandler
+export VariableSearchRadiusNHS
 
 end # module

src/general/gpu.jl

@@ -45,3 +45,31 @@ end
 function allocate(backend, ELTYPE, size)
     return Array{ELTYPE, length(size)}(undef, size)
 end
+
+function Adapt.adapt_structure(to, nhs::VariableSearchRadiusNHS)
+    (; search_radii, neighborhood_searches) = nhs
+    radii = Adapt.adapt_structure(to, search_radii)
+    searches = Adapt.adapt_structure(to, neighborhood_searches)
+
+    return VariableSearchRadiusNHS(radii, searches)
+end
+
+function Adapt.adapt_structure(to, nhs_handler::GridNHSHandler)
+    (; search_radii, neighborhood_searches) = nhs_handler
+    radii = Adapt.adapt_structure(to, search_radii)
+    searches = Adapt.adapt_structure(to, neighborhood_searches)
+
+    return GridNHSHandler(radii, searches)
+end
+
+function Adapt.adapt_structure(to, nhs_handler::PairsNHSHandler)
+    searches = Adapt.adapt_structure(to, nhs_handler.neighborhood_searches)
+
+    return PairsNHSHandler(searches)
+end
+
+function Adapt.adapt_structure(to, nhs_handler::VariableNHSHandler)
+    searches = Adapt.adapt_structure(to, nhs_handler.neighborhood_searches)
+
+    return VariableNHSHandler(searches)
+end

src/general/neighborhood_search.jl

@@ -5,10 +5,11 @@
 function PointNeighbors.foreach_point_neighbor(f, system, neighbor_system,
                                                system_coords, neighbor_coords, semi;
                                                points=eachparticle(system),
-                                               parallelization_backend=semi.parallelization_backend)
+                                               parallelization_backend=semi.parallelization_backend,
+                                               search_radius=nothing)
     neighborhood_search = get_neighborhood_search(system, neighbor_system, semi)
     foreach_point_neighbor(f, system_coords, neighbor_coords, neighborhood_search;
-                           points, parallelization_backend)
+                           points, parallelization_backend, search_radius)
 end
 
 # === Compact support selection ===
@@ -103,13 +104,91 @@ function create_neighborhood_search(neighborhood_search, system::TotalLagrangian
                                     nparticles(neighbor))
 end
 
-# === Neighborhood search lookup ===
-@inline function get_neighborhood_search(system, semi)
-    (; neighborhood_searches) = semi
+# === Neighborhood search handlers ===
+# Base type for all neighborhood search handlers. Handlers manage how neighborhood 
+# searches are stored and retrieved for different interacting systems.
+abstract type AbstractNHSHandler end
 
-    system_index = system_indices(system, semi)
+# Handler that stores a distinct neighborhood search for every possible pair of systems.
+# Best suited for standard neighborhood searches where each pair has a fixed search radius.
+struct PairsNHSHandler{NHS} <: AbstractNHSHandler
+    neighborhood_searches::NHS
+end
+
+function PairsNHSHandler(nhs, systems)
+    # Create a tuple of n neighborhood searches for each of the n systems.
+    # We will need one neighborhood search for each pair of systems.
+    searches = Tuple(Tuple(create_neighborhood_search(nhs, system, neighbor)
+                           for neighbor in systems)
+                     for system in systems)
+
+    return PairsNHSHandler(searches)
+end
+
+function get_neighborhood_search(handler::PairsNHSHandler, system_index, neighbor_index, _)
+    return handler.neighborhood_searches[system_index][neighbor_index]
+end
+
+# Handler specifically designed for grid-based neighborhood searches that need to operate
+# with multiple, pre-defined unique search radii per system.
+struct GridNHSHandler{SI, NHS} <: AbstractNHSHandler
+    search_radii::SI
+    neighborhood_searches::NHS
+end
+
+function GridNHSHandler(nhs::PointNeighbors.AbstractNeighborhoodSearch, systems)
+    # Find a list of search radii that will be requested for each system
+    search_radii = Tuple([compact_support(system, neighbor) for neighbor in systems]
+                         for system in systems)
+    search_radii = Tuple.(sort.(unique.(search_radii)))
+
+    # For each system, create a neighborhood search for each unique search radius
+    searches = Tuple(Tuple(copy_neighborhood_search(nhs, search_radius, nparticles(system))
+                           for search_radius in
+                               search_radii[system_indices(system, systems)])
+                     for system in systems)
+
+    return GridNHSHandler{typeof(search_radii), typeof(searches)}(search_radii, searches)
+end
 
-    return neighborhood_searches[system_index][system_index]
+function get_neighborhood_search(handler::GridNHSHandler, system_index, neighbor_index,
+                                 search_radius)
+    # Find the index of the search radius in the list of search radii for this system
+    search_radii = handler.search_radii[neighbor_index]
+    idx = searchsortedfirst(SVector(search_radii), search_radius - eps(search_radius))
+    # Return the neighborhood search at that index
+    return handler.neighborhood_searches[neighbor_index][idx]
+end
+
+# Handler that uses `VariableSearchRadiusNHS` to support different 
+# compact supports for the same system.
+struct VariableNHSHandler{NHS} <: AbstractNHSHandler
+    neighborhood_searches::NHS
+end
+
+function VariableNHSHandler(nhs, systems)
+    # Find a list of search radii that will be requested for each system
+    search_radii = Tuple([compact_support(system, neighbor) for neighbor in systems]
+                         for system in systems)
+    search_radii = Tuple.(sort.(unique.(search_radii)))
+
+    # For each system, create a neighborhood search for each unique search radius
+    searches = Tuple(VariableSearchRadiusNHS(nhs,
+                                             search_radii[system_indices(system,
+                                                                         systems)],
+                                             nparticles(system))
+                     for system in systems)
+
+    return VariableNHSHandler(searches)
+end
+
+function get_neighborhood_search(handler::VariableNHSHandler, system_index, neighbor_index,
+                                 search_radius)
+    handler.neighborhood_searches[neighbor_index]
+end
+
+@inline function get_neighborhood_search(system, semi)
+    return get_neighborhood_search(system, system, semi)
 end
 
 @inline function get_neighborhood_search(system::TotalLagrangianSPHSystem, semi)
@@ -119,17 +198,21 @@ end
 end
 
 @inline function get_neighborhood_search(system, neighbor_system, semi)
-    (; neighborhood_searches) = semi
+    (; neighborhood_search_handler) = semi
 
     system_index = system_indices(system, semi)
     neighbor_index = system_indices(neighbor_system, semi)
 
-    return neighborhood_searches[system_index][neighbor_index]
+    search_radius = compact_support(system, neighbor_system)
+
+    return get_neighborhood_search(neighborhood_search_handler, system_index,
+                                   neighbor_index,
+                                   search_radius)
 end
 
 @inline function get_neighborhood_search(system::TotalLagrangianSPHSystem,
                                          neighbor_system::TotalLagrangianSPHSystem, semi)
-    (; neighborhood_searches) = semi
+    (; neighborhood_search_handler) = semi
 
     system_index = system_indices(system, semi)
     neighbor_index = system_indices(neighbor_system, semi)
@@ -140,7 +223,11 @@ end
         return system.self_interaction_nhs
     end
 
-    return neighborhood_searches[system_index][neighbor_index]
+    search_radius = compact_support(system, neighbor_system)
+
+    return get_neighborhood_search(neighborhood_search_handler, system_index,
+                                   neighbor_index,
+                                   search_radius)
 end
 
 # === Initialization ===

src/general/semidiscretization.jl

@@ -49,33 +49,36 @@ semi = Semidiscretization(fluid_system, boundary_system,
 └──────────────────────────────────────────────────────────────────────────────────────────────────┘
 ```
 """
-struct Semidiscretization{BACKEND, S, RU, RV, NS, UCU, IT}
-    systems                 :: S
-    ranges_u                :: RU
-    ranges_v                :: RV
-    neighborhood_searches   :: NS
-    parallelization_backend :: BACKEND
-    update_callback_used    :: UCU
-    integrate_tlsph         :: IT # `false` if TLSPH integration is decoupled
+struct Semidiscretization{BACKEND, S, RU, RV, NSH, UCU, IT}
+    systems                     :: S
+    ranges_u                    :: RU
+    ranges_v                    :: RV
+    neighborhood_search_handler :: NSH
+    parallelization_backend     :: BACKEND
+    update_callback_used        :: UCU
+    integrate_tlsph             :: IT # `false` if TLSPH integration is decoupled
 
     # Dispatch at `systems` to distinguish this constructor from the one below when
     # 4 systems are passed.
     # This is an internal constructor only used in `test/count_allocations.jl`
     # and by Adapt.jl.
-    function Semidiscretization(systems::Tuple, ranges_u, ranges_v, neighborhood_searches,
+    function Semidiscretization(systems::Tuple, ranges_u, ranges_v,
+                                neighborhood_search_handler,
                                 parallelization_backend::PointNeighbors.ParallelizationBackend,
                                 update_callback_used, integrate_tlsph)
         new{typeof(parallelization_backend), typeof(systems), typeof(ranges_u),
-            typeof(ranges_v), typeof(neighborhood_searches),
+            typeof(ranges_v), typeof(neighborhood_search_handler),
             typeof(update_callback_used),
             typeof(integrate_tlsph)}(systems, ranges_u, ranges_v,
-                                     neighborhood_searches, parallelization_backend,
+                                     neighborhood_search_handler, parallelization_backend,
                                      update_callback_used, integrate_tlsph)
     end
 end
 
 function Semidiscretization(systems::Union{AbstractSystem, Nothing}...;
                             neighborhood_search=GridNeighborhoodSearch{ndims(first(systems))}(),
+                            neighborhood_search_handler=PairsNHSHandler(neighborhood_search,
+                                                                        systems),
                             parallelization_backend=PolyesterBackend())
     systems = filter(system -> !isnothing(system), systems)
 
@@ -98,13 +101,6 @@ function Semidiscretization(systems::Union{AbstractSystem, Nothing}...;
     ranges_v = Tuple((sum(sizes_v[1:(i - 1)]) + 1):sum(sizes_v[1:i])
                      for i in eachindex(sizes_v))
 
-    # Create a tuple of n neighborhood searches for each of the n systems.
-    # We will need one neighborhood search for each pair of systems.
-    searches = Tuple(Tuple(create_neighborhood_search(neighborhood_search,
-                                                      system, neighbor)
-                           for neighbor in systems)
-                     for system in systems)
-
     # These will be set to true inside the `UpdateCallback`.
     # Some techniques require the use of this callback, and this flag can be used
     # to determine if the callback is used in a simulation.
@@ -115,7 +111,8 @@ function Semidiscretization(systems::Union{AbstractSystem, Nothing}...;
     # with this set to false.
     integrate_tlsph = Ref(true)
 
-    return Semidiscretization(systems, ranges_u, ranges_v, searches,
+    return Semidiscretization(systems, ranges_u, ranges_v,
+                              neighborhood_search_handler,
                               parallelization_backend, update_callback_used,
                               integrate_tlsph)
 end
@@ -144,18 +141,21 @@ function Base.show(io::IO, ::MIME"text/plain", semi::Semidiscretization)
         summary_line(io, "#spatial dimensions", ndims(semi.systems[1]))
         summary_line(io, "#systems", length(semi.systems))
         summary_line(io, "neighborhood search",
-                     semi.neighborhood_searches |> eltype |> eltype |> nameof)
+                     semi.neighborhood_search_handler |> eltype |> eltype |> nameof)
         summary_line(io, "total #particles", sum(nparticles.(semi.systems)))
         summary_line(io, "eltype", eltype(semi.systems[1]))
         summary_line(io, "coordinates eltype", coordinates_eltype(semi.systems[1]))
         summary_footer(io)
     end
 end
 
-@inline function system_indices(system, semi)
+@inline system_indices(system,
+                       semi::Semidiscretization) = system_indices(system, semi.systems)
+
+@inline function system_indices(system, systems)
     # Note that this takes only about 5 ns, while mapping systems to indices with a `Dict`
     # is ~30x slower because `hash(::System)` is very slow.
-    index = findfirst(==(system), semi.systems)
+    index = findfirst(==(system), systems)
 
     if isnothing(index)
         throw(ArgumentError("system is not in the semidiscretization"))

test/general/nhs_handler.jl

@@ -0,0 +1,38 @@
+using OrdinaryDiffEq
+
+@testset verbose=true "Test NHS Handlers" begin
+    nhs_handlers = (
+                    # PairsNHSHandler, 
+                    # GridNHSHandler, 
+                    VariableNHSHandler,)
+
+    @testset verbose=true "$nhs_handler" for nhs_handler in nhs_handlers
+
+        # Run the `dam_break_gate_2d.jl` file for 0 timesteps to instantiate all systems
+        sol = trixi_include(@__MODULE__,
+                            joinpath(examples_dir(), "fsi", "dam_break_gate_2d.jl"),
+                            tspan=(0.0, 0.0),
+                            callbacks=OrdinaryDiffEq.CallbackSet())
+
+        # sol.prob.p is the original semidiscretization
+        semi = sol.prob.p
+        systems = semi.systems
+        nhs = TrixiParticles.get_neighborhood_search(first(systems), semi)
+        neighborhood_search_handler = nhs_handler(nhs, systems)
+
+        new_semi = Semidiscretization(systems...,
+                                      neighborhood_search_handler=neighborhood_search_handler,
+                                      parallelization_backend=PolyesterBackend())
+
+        tspan_test = (0.0, 0.1)
+        ode = semidiscretize(new_semi, tspan_test)
+
+        new_sol = Base.invokelatest(solve, ode, RDPK3SpFSAL49();
+                                    abstol=1e-6,
+                                    reltol=1e-4,
+                                    dtmax=1e-3,
+                                    save_everystep=false)
+
+        @test new_sol.retcode == ReturnCode.Success
+    end
+end