Implement copy_neighborhood_search to allow empty "template" neighborhood search (trixi-framework#32)

efaulhaber · web-flow · commit 5d52dfd7ed63 · 2024-06-24T10:35:47.000+02:00
* Rename `threaded_nhs_update` to `threaded_update` * Add `threaded_update` option to `PrecomputedNeighborhoodSearch` * Properly implement `copy_neighborhood_search` * Create empty NHS templates by default * Fix tests * Add tests for copied templates * Add doctest * Fix rebase * Break doctests to verify that they are running * Revert "Break doctests to verify that they are running" This reverts commit 7e4275d. * Fix typo * Implement suggestions
diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl
@@ -11,10 +11,11 @@ include("neighborhood_search.jl")
 include("nhs_trivial.jl")
 include("cell_lists/cell_lists.jl")
 include("nhs_grid.jl")
-include("nhs_neighbor_lists.jl")
+include("nhs_precomputed.jl")
 
-export foreach_point_neighbor, foreach_neighbor, PeriodicBox
+export foreach_point_neighbor, foreach_neighbor
 export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch
 export initialize!, update!, initialize_grid!, update_grid!
+export PeriodicBox, copy_neighborhood_search
 
 end # module PointNeighbors
diff --git a/src/neighborhood_search.jl b/src/neighborhood_search.jl
@@ -41,6 +41,37 @@ See also [`initialize!`](@ref).
     return search
 end
 
+"""
+    copy_neighborhood_search(search::AbstractNeighborhoodSearch, search_radius, n_points;
+                             eachpoint = 1:n_points)
+
+Create a new **uninitialized** neighborhood search of the same type and with the same
+configuration options as `search`, but with a different search radius and number of points.
+
+The [`TrivialNeighborhoodSearch`](@ref) also requires an iterator `eachpoint`, which most
+of the time will be `1:n_points`. If the `TrivialNeighborhoodSearch` is never going to be
+used, the keyword argument `eachpoint` can be ignored.
+
+This is useful when a simulation code requires multiple neighborhood searches of the same
+kind. One can then just pass an empty neighborhood search as a template and use
+this function inside the simulation code to generate similar neighborhood searches with
+different search radii and different numbers of points.
+```jldoctest; filter = r"GridNeighborhoodSearch{2,.*"
+# Template
+nhs = GridNeighborhoodSearch{2}()
+
+# Inside the simulation code, generate similar neighborhood searches
+nhs1 = copy_neighborhood_search(nhs, 1.0, 100)
+
+# output
+GridNeighborhoodSearch{2, Float64, ...}(...)
+```
+"""
+@inline function copy_neighborhood_search(search::AbstractNeighborhoodSearch,
+                                          search_radius, n_points; eachpoint = 1:n_points)
+    return nothing
+end
+
 """
     PeriodicBox(; min_corner, max_corner)
 
diff --git a/src/nhs_grid.jl b/src/nhs_grid.jl
@@ -1,6 +1,6 @@
 @doc raw"""
-    GridNeighborhoodSearch{NDIMS}(search_radius, n_points;
-                                  periodic_box = nothing, threaded_nhs_update = true)
+    GridNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
+                                  periodic_box = nothing, threaded_update = true)
 
 Simple grid-based neighborhood search with uniform search radius.
 The domain is divided into a regular grid.
@@ -23,17 +23,19 @@ As opposed to (Ihmsen et al. 2011), we do not sort the points in any way,
 since not sorting makes our implementation a lot faster (although less parallelizable).
 
 # Arguments
-- `NDIMS`:          Number of dimensions.
-- `search_radius`:  The uniform search radius.
-- `n_points`:    Total number of points.
+- `NDIMS`: Number of dimensions.
 
 # Keywords
+- `search_radius = 0.0`:    The fixed search radius. The default of `0.0` is useful together
+                            with [`copy_neighborhood_search`](@ref).
+- `n_points = 0`:           Total number of points. The default of `0` is useful together
+                            with [`copy_neighborhood_search`](@ref).
 - `periodic_box = nothing`: In order to use a (rectangular) periodic domain, pass a
                             [`PeriodicBox`](@ref).
-- `threaded_nhs_update = true`: Can be used to deactivate thread parallelization in the
-                                neighborhood search update. This can be one of the largest
-                                sources of variations between simulations with different
-                                thread numbers due to neighbor ordering changes.
+- `threaded_update = true`: Can be used to deactivate thread parallelization in the
+                            neighborhood search update. This can be one of the largest
+                            sources of variations between simulations with different
+                            thread numbers due to neighbor ordering changes.
 
 ## References
 - M. Chalela, E. Sillero, L. Pereyra, M.A. Garcia, J.B. Cabral, M. Lares, M. Merchán.
@@ -53,11 +55,11 @@ struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, PB} <: AbstractNeighborhoodSear
     cell_size           :: NTuple{NDIMS, ELTYPE} # Required to calculate cell index
     cell_buffer         :: Array{NTuple{NDIMS, Int}, 2} # Multithreaded buffer for `update!`
     cell_buffer_indices :: Vector{Int} # Store which entries of `cell_buffer` are initialized
-    threaded_nhs_update :: Bool
+    threaded_update     :: Bool
 
-    function GridNeighborhoodSearch{NDIMS}(search_radius, n_points;
+    function GridNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
                                            periodic_box = nothing,
-                                           threaded_nhs_update = true) where {NDIMS}
+                                           threaded_update = true) where {NDIMS}
         ELTYPE = typeof(search_radius)
         cell_list = DictionaryCellList{NDIMS}()
 
@@ -87,7 +89,7 @@ struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, PB} <: AbstractNeighborhoodSear
         new{NDIMS, ELTYPE, typeof(cell_list),
             typeof(periodic_box)}(cell_list, search_radius, periodic_box, n_cells,
                                   cell_size, cell_buffer, cell_buffer_indices,
-                                  threaded_nhs_update)
+                                  threaded_update)
     end
 end
 
@@ -141,14 +143,14 @@ end
 
 # Modify the existing hash table by moving points into their new cells
 function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
-    (; cell_list, cell_buffer, cell_buffer_indices, threaded_nhs_update) = neighborhood_search
+    (; cell_list, cell_buffer, cell_buffer_indices, threaded_update) = neighborhood_search
 
     # Reset `cell_buffer` by moving all pointers to the beginning
     cell_buffer_indices .= 0
 
     # Find all cells containing points that now belong to another cell
     mark_changed_cell!(neighborhood_search, cell_list, coords_fun,
-                       Val(threaded_nhs_update))
+                       Val(threaded_update))
 
     # Iterate over all marked cells and move the points into their new cells.
     for thread in 1:Threads.nthreads()
@@ -180,15 +182,15 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
 end
 
 @inline function mark_changed_cell!(neighborhood_search, cell_list, coords_fun,
-                                    threaded_nhs_update::Val{true})
+                                    threaded_update::Val{true})
     # `collect` the keyset to be able to loop over it with `@threaded`
     @threaded for cell in collect(eachcell(cell_list))
         mark_changed_cell!(neighborhood_search, cell, coords_fun)
     end
 end
 
 @inline function mark_changed_cell!(neighborhood_search, cell_list, coords_fun,
-                                    threaded_nhs_update::Val{false})
+                                    threaded_update::Val{false})
     for cell in eachcell(cell_list)
         mark_changed_cell!(neighborhood_search, cell, coords_fun)
     end
@@ -307,13 +309,9 @@ end
     return floor(Int, i)
 end
 
-# Create a copy of a neighborhood search but with a different search radius
-function copy_neighborhood_search(nhs::GridNeighborhoodSearch, search_radius, x, y)
-    search = GridNeighborhoodSearch{ndims(nhs)}(search_radius, npoints(nhs),
-                                                periodic_box = nhs.periodic_box)
-
-    # Initialize neighborhood search
-    initialize!(search, x, y)
-
-    return search
+function copy_neighborhood_search(nhs::GridNeighborhoodSearch, search_radius, n_points;
+                                  eachpoint = 1:n_points)
+    return GridNeighborhoodSearch{ndims(nhs)}(; search_radius, n_points,
+                                              periodic_box = nhs.periodic_box,
+                                              threaded_update = nhs.threaded_update)
 end
diff --git a/src/nhs_precomputed.jl b/src/nhs_precomputed.jl
@@ -1,6 +1,6 @@
 @doc raw"""
-    PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_points;
-                                         periodic_box = nothing)
+    PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
+                                         periodic_box = nothing, threaded_update = true)
 
 Neighborhood search with precomputed neighbor lists. A list of all neighbors is computed
 for each point during initialization and update.
@@ -11,23 +11,31 @@ A [`GridNeighborhoodSearch`](@ref) is used internally to compute the neighbor li
 initialization and update.
 
 # Arguments
-- `NDIMS`:          Number of dimensions.
-- `search_radius`:  The uniform search radius.
-- `n_points`:    Total number of points.
+- `NDIMS`: Number of dimensions.
 
 # Keywords
+- `search_radius = 0.0`:    The fixed search radius. The default of `0.0` is useful together
+                            with [`copy_neighborhood_search`](@ref).
+- `n_points = 0`:           Total number of points. The default of `0` is useful together
+                            with [`copy_neighborhood_search`](@ref).
 - `periodic_box = nothing`: In order to use a (rectangular) periodic domain, pass a
                             [`PeriodicBox`](@ref).
+- `threaded_update = true`: Can be used to deactivate thread parallelization in the
+                            neighborhood search update. This can be one of the largest
+                            sources of variations between simulations with different
+                            thread numbers due to neighbor ordering changes.
 """
 struct PrecomputedNeighborhoodSearch{NDIMS, NHS, NL, PB}
     neighborhood_search :: NHS
     neighbor_lists      :: NL
     periodic_box        :: PB
 
-    function PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_points;
-                                                  periodic_box = nothing) where {NDIMS}
-        nhs = GridNeighborhoodSearch{NDIMS}(search_radius, n_points,
-                                            periodic_box = periodic_box)
+    function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
+                                                  periodic_box = nothing,
+                                                  threaded_update = true) where {NDIMS}
+        nhs = GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                            periodic_box = periodic_box,
+                                            threaded_update = threaded_update)
 
         neighbor_lists = Vector{Vector{Int}}()
 
@@ -103,3 +111,11 @@ end
         @inline f(point, neighbor, pos_diff, distance)
     end
 end
+
+function copy_neighborhood_search(nhs::PrecomputedNeighborhoodSearch,
+                                  search_radius, n_points; eachpoint = 1:n_points)
+    threaded_update = nhs.neighborhood_search.threaded_update
+    return PrecomputedNeighborhoodSearch{ndims(nhs)}(; search_radius, n_points,
+                                                     periodic_box = nhs.periodic_box,
+                                                     threaded_update)
+end
diff --git a/src/nhs_trivial.jl b/src/nhs_trivial.jl
@@ -1,14 +1,18 @@
 @doc raw"""
-    TrivialNeighborhoodSearch{NDIMS}(search_radius, eachpoint, periodic_box = nothing)
+    TrivialNeighborhoodSearch{NDIMS}(; search_radius = 0.0, eachpoint = 1:0,
+                                     periodic_box = nothing)
 
 Trivial neighborhood search that simply loops over all points.
 
 # Arguments
-- `NDIMS`:          Number of dimensions.
-- `search_radius`:  The uniform search radius.
-- `eachparticle`:   Iterator for all point indices. Usually just `1:n_points`.
+- `NDIMS`: Number of dimensions.
 
 # Keywords
+- `search_radius = 0.0`:    The fixed search radius. The default of `0.0` is useful together
+                            with [`copy_neighborhood_search`](@ref).
+- `eachpoint = 1:0`:        Iterator for all point indices. Usually just `1:n_points`.
+                            The default of `1:0` is useful together with
+                            [`copy_neighborhood_search`](@ref).
 - `periodic_box = nothing`: In order to use a (rectangular) periodic domain, pass a
                             [`PeriodicBox`](@ref).
 """
@@ -17,7 +21,7 @@ struct TrivialNeighborhoodSearch{NDIMS, ELTYPE, EP, PB} <: AbstractNeighborhoodS
     eachpoint     :: EP
     periodic_box  :: PB
 
-    function TrivialNeighborhoodSearch{NDIMS}(search_radius, eachpoint;
+    function TrivialNeighborhoodSearch{NDIMS}(; search_radius = 0.0, eachpoint = 1:0,
                                               periodic_box = nothing) where {NDIMS}
         new{NDIMS, typeof(search_radius),
             typeof(eachpoint), typeof(periodic_box)}(search_radius, eachpoint, periodic_box)
@@ -39,3 +43,9 @@ end
 function copy_neighborhood_search(nhs::TrivialNeighborhoodSearch, search_radius, x, y)
     return nhs
 end
+
+function copy_neighborhood_search(nhs::TrivialNeighborhoodSearch,
+                                  search_radius, n_points; eachpoint = 1:n_points)
+    return TrivialNeighborhoodSearch{ndims(nhs)}(; search_radius, eachpoint,
+                                                 periodic_box = nhs.periodic_box)
+end
diff --git a/test/neighborhood_search.jl b/test/neighborhood_search.jl
@@ -39,21 +39,34 @@
             search_radius = 0.1
 
             neighborhood_searches = [
-                TrivialNeighborhoodSearch{NDIMS}(search_radius, 1:n_points,
+                TrivialNeighborhoodSearch{NDIMS}(; search_radius, eachpoint = 1:n_points,
                                                  periodic_box = periodic_boxes[i]),
-                GridNeighborhoodSearch{NDIMS}(search_radius, n_points,
+                GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               periodic_box = periodic_boxes[i]),
-                PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_points,
+                PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                                      periodic_box = periodic_boxes[i]),
             ]
-            neighborhood_searches_names = [
+
+            names = [
                 "`TrivialNeighborhoodSearch`",
                 "`GridNeighborhoodSearch`",
                 "`PrecomputedNeighborhoodSearch`",
             ]
 
+            # Also test copied templates
+            template_nhs = [
+                TrivialNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i]),
+                GridNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i]),
+                PrecomputedNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i]),
+            ]
+            copied_nhs = copy_neighborhood_search.(template_nhs, search_radius, n_points)
+            append!(neighborhood_searches, copied_nhs)
+
+            names_copied = [name * " copied" for name in names]
+            append!(names, names_copied)
+
             # Run this for every neighborhood search
-            @testset "$(neighborhood_searches_names[j])" for j in eachindex(neighborhood_searches_names)
+            @testset "$(names[j])" for j in eachindex(names)
                 nhs = neighborhood_searches[j]
 
                 initialize!(nhs, coords, coords)
@@ -103,7 +116,8 @@
 
             # Compute expected neighbor lists by brute-force looping over all points
             # as potential neighbors (`TrivialNeighborhoodSearch`).
-            trivial_nhs = TrivialNeighborhoodSearch{NDIMS}(search_radius, axes(coords, 2))
+            trivial_nhs = TrivialNeighborhoodSearch{NDIMS}(; search_radius,
+                                                           eachpoint = axes(coords, 2))
 
             neighbors_expected = [Int[] for _ in axes(coords, 2)]
 
@@ -114,16 +128,27 @@
             end
 
             neighborhood_searches = [
-                GridNeighborhoodSearch{NDIMS}(search_radius, n_points),
-                PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_points),
+                GridNeighborhoodSearch{NDIMS}(; search_radius, n_points),
+                PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points),
             ]
 
-            neighborhood_searches_names = [
+            names = [
                 "`GridNeighborhoodSearch`",
                 "`PrecomputedNeighborhoodSearch`",
             ]
 
-            @testset "$(neighborhood_searches_names[i])" for i in eachindex(neighborhood_searches_names)
+            # Also test copied templates
+            template_nhs = [
+                GridNeighborhoodSearch{NDIMS}(),
+                PrecomputedNeighborhoodSearch{NDIMS}(),
+            ]
+            copied_nhs = copy_neighborhood_search.(template_nhs, search_radius, n_points)
+            append!(neighborhood_searches, copied_nhs)
+
+            names_copied = [name * " copied" for name in names]
+            append!(names, names_copied)
+
+            @testset "$(names[i])" for i in eachindex(names)
                 nhs = neighborhood_searches[i]
 
                 # Initialize with `seed = 1`
diff --git a/test/nhs_grid.jl b/test/nhs_grid.jl
diff --git a/test/nhs_trivial.jl b/test/nhs_trivial.jl
diff --git a/test/runtests.jl b/test/runtests.jl
