VariableNHSHandler and GridNHSHandler are running correctly with the 2d dam break example.

Author: RubberLanding

src/general/neighborhood_search.jl

@@ -114,85 +114,135 @@ function create_neighborhood_search_handler(::Type{<:PairsNHSHandler},
                                             neighborhood_search, 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.
-    return Tuple(Tuple(create_neighborhood_search(neighborhood_search,
-                                                  system, neighbor)
+    searches = Tuple(Tuple(create_neighborhood_search(neighborhood_search,
+                                                      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
 
-struct VariableRadiusNHS{SR, NHS} <: PointNeighbors.AbstractNeighborhoodSearch
+struct VariableNHSHandler{NHS} <: AbstractNHSHandler
+    neighborhood_searches::NHS
+end
+
+function create_neighborhood_search_handler(::Type{<: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.(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)
+
+    # For each system, create a neighborhood search for each unique search radius
+    return VariableNHSHandler(searches)
+end
+
+function get_neighborhood_search(handler::VariableNHSHandler, system_index, neighbor_index,
+                                 search_radius)
+    handler.neighborhood_searches[neighbor_index]
+end
+
+struct VariableSearchRadiusNHS{SR, NHS} <: PointNeighbors.AbstractNeighborhoodSearch
     search_radii          :: SR
     neighborhood_searches :: NHS
 end
 
 function VariableSearchRadiusNHS(nhs_implementation, search_radii, n_particles)
-    searches = Tuple(copy_neighborhood_search(nhs_implementation, search_radius, n_particles)
+    searches = Tuple(PointNeighbors.copy_neighborhood_search(nhs_implementation,
+                                                             search_radius, n_particles)
                      for search_radius in search_radii)
-    return VariableRadiusNHS(search_radii, searches)
+    return VariableSearchRadiusNHS(search_radii, searches)
 end
 
-@inline Base.ndims(::TrivialNeighborhoodSearch{NDIMS}) where {NDIMS} = NDIMS
+@inline Base.ndims(search::VariableSearchRadiusNHS) = ndims(first(search.neighborhood_searches))
 
-@inline requires_update(::TrivialNeighborhoodSearch) = (false, true)
+@inline requires_update(::VariableSearchRadiusNHS) = (false, true)
 
-@inline function initialize!(search::TrivialNeighborhoodSearch, x, y;
-                             parallelization_backend = default_backend(x),
-                             eachindex_y = axes(y, 2))
+# TODO
+@inline function initialize!(search::VariableSearchRadiusNHS, x, y;
+                             parallelization_backend=default_backend(x),
+                             eachindex_y=axes(y, 2))
     return search
 end
 
-@inline function update!(search::TrivialNeighborhoodSearch, x, y;
-                         points_moving = (true, true),
-                         parallelization_backend = default_backend(x),
-                         eachindex_y = axes(y, 2))
+# TODO
+@inline function update!(search::VariableSearchRadiusNHS, x, y;
+                         points_moving=(true, true),
+                         parallelization_backend=default_backend(x),
+                         eachindex_y=axes(y, 2))
     return search
 end
 
 # Create a copy of a neighborhood search but with a different search radius
-function copy_neighborhood_search(nhs::TrivialNeighborhoodSearch, search_radius, x, y)
-    return nhs
+function PointNeighbors.copy_neighborhood_search(nhs::VariableSearchRadiusNHS,
+                                                 search_radius, x, y)
+    search_radii = Tuple(PointNeighbors.search_radius(search)
+                         for search in nhs.neighborhood_searches)
+    searches = Tuple(PointNeighbors.copy_neighborhood_search(search)
+                     for search in nhs.neighborhood_searches)
+
+    return VariableSearchRadiusNHS(search_radii, searches)
 end
 
-function copy_neighborhood_search(nhs::TrivialNeighborhoodSearch,
-                                  search_radius, n_points; eachpoint = 1:n_points)
+function PointNeighbors.copy_neighborhood_search(nhs::VariableSearchRadiusNHS,
+                                                 search_radius, n_points;
+                                                 eachpoint=1:n_points)
     return TrivialNeighborhoodSearch{ndims(nhs)}(; search_radius, eachpoint,
-                                                 periodic_box = nhs.periodic_box)
+                                                 periodic_box=nhs.periodic_box)
 end
 
 @inline function foreach_neighbor(f, neighbor_system_coords,
-                                  neighborhood_search::VariableRadiusNHS,
+                                  neighborhood_search::VariableSearchRadiusNHS,
                                   point, point_coords, search_radius)
     # Find the index of the search radius in the list of search radii for this system
     search_radii = neighborhood_search.search_radii
-    idx = searchsortedfirst(search_radii, search_radius - eps(search_radius))
+    idx = searchsortedfirst(SVector(search_radii), search_radius - eps(search_radius))
 
     nhs = neighborhood_search.neighborhood_searches[idx]
 
-    PointNeighbors.foreach_neighbor(f, neighbor_system_coords, nhs, point, point_coords, search_radius)
+    PointNeighbors.foreach_neighbor(f, neighbor_system_coords, nhs, point, point_coords,
+                                    search_radius)
 end
 
-struct VariableNHSHandler{NHS} <: AbstractNHSHandler
+struct GridNHSHandler{SI, NHS} <: AbstractNHSHandler
+    search_radii::SI
     neighborhood_searches::NHS
 end
 
-function create_neighborhood_search_handler(::Type{<:VariableNHSHandler}, nhs, systems)
+function create_neighborhood_search_handler(::Type{<:GridNHSHandler}, 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([compact_support(system, neighbor) for neighbor in systems]
+                         for system in systems)
     search_radii = Tuple.(unique.(search_radii))
 
     # For each system, create a neighborhood search for each unique search radius
-    return Tuple(VariableSearchRadiusNHS(nhs, search_radii[system_indices(system, systems)], nparticles(system))
-                 for system in systems)
+    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(search_radii, searches)
 end
 
-function get_neighborhood_search(handler::VariableNHSHandler, system_index, neighbor_index,
+function get_neighborhood_search(handler::GridNHSHandler, system_index, neighbor_index,
                                  search_radius)
-    handler.neighborhood_searches[neighbor_index]
+    # 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
+# ==================================
 
 @inline function get_neighborhood_search(system, semi)
     return get_neighborhood_search(system, system, semi)
@@ -205,20 +255,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)
 
     search_radius = compact_support(system, neighbor_system)
 
-    return get_neighborhood_search(neighborhood_searches, system_index, neighbor_index,
+    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)
@@ -231,7 +282,8 @@ end
 
     search_radius = compact_support(system, neighbor_system)
 
-    return get_neighborhood_search(neighborhood_searches, system_index, neighbor_index,
+    return get_neighborhood_search(neighborhood_search_handler, system_index,
+                                   neighbor_index,
                                    search_radius)
 end
 

src/general/semidiscretization.jl

@@ -49,32 +49,34 @@ 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}...;
+                            nhs_handler=GridNHSHandler,
                             neighborhood_search=GridNeighborhoodSearch{ndims(first(systems))}(),
                             parallelization_backend=PolyesterBackend())
     systems = filter(system -> !isnothing(system), systems)
@@ -98,7 +100,9 @@ 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))
 
-    searches = create_neighborhood_search_handler(PairsNHSHandler, neighborhood_search, systems)
+    neighborhood_search_handler = create_neighborhood_search_handler(nhs_handler,
+                                                                     neighborhood_search,
+                                                                     systems)
 
     # These will be set to true inside the `UpdateCallback`.
     # Some techniques require the use of this callback, and this flag can be used
@@ -110,7 +114,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
@@ -139,15 +144,16 @@ 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 system_indices(system, semi::Semidiscretization) = system_indices(system, semi.systems)
+@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`