Merge branch 'main' into SpatialHashing

Author: RubberLanding

Project.toml

@@ -3,7 +3,6 @@ uuid = "1c4d5385-0a27-49de-8e2c-43b175c8985c"
 authors = ["Erik Faulhaber <erik.faulhaber@uni-koeln.de>"]
 version = "0.6.2"
 
-
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"

src/neighborhood_search.jl

@@ -14,7 +14,8 @@ end
 
 """
     initialize!(search::AbstractNeighborhoodSearch, x, y;
-                parallelization_backend = default_backend(x))
+                parallelization_backend = default_backend(x),
+                eachindex_y = axes(y, 2))
 
 Initialize a neighborhood search with the two coordinate arrays `x` and `y`.
 
@@ -27,16 +28,21 @@ If the neighborhood search type supports parallelization, the keyword argument
 `parallelization_backend` can be used to specify a parallelization backend.
 See [`@threaded`](@ref) for a list of available backends.
 
+Optionally, when points in `y` are to be ignored, the keyword argument `eachindex_y` can be
+passed to specify the indices of the points in `y` that are to be used.
+
 See also [`update!`](@ref).
 """
 @inline function initialize!(search::AbstractNeighborhoodSearch, x, y;
-                             parallelization_backend = default_backend(x))
+                             parallelization_backend = default_backend(x),
+                             eachindex_y = axes(y, 2))
     return search
 end
 
 """
     update!(search::AbstractNeighborhoodSearch, x, y; points_moving = (true, true),
-            parallelization_backend = default_backend(x))
+            parallelization_backend = default_backend(x),
+            eachindex_y = axes(y, 2))
 
 Update an already initialized neighborhood search with the two coordinate arrays `x` and `y`.
 
@@ -59,11 +65,15 @@ If the neighborhood search type supports parallelization, the keyword argument
 `parallelization_backend` can be used to specify a parallelization backend.
 See [`@threaded`](@ref) for a list of available backends.
 
+Optionally, when points in `y` are to be ignored, the keyword argument `eachindex_y` can be
+passed to specify the indices of the points in `y` that are to be used.
+
 See also [`initialize!`](@ref).
 """
 @inline function update!(search::AbstractNeighborhoodSearch, x, y;
                          points_moving = (true, true),
-                         parallelization_backend = default_backend(x))
+                         parallelization_backend = default_backend(x),
+                         eachindex_y = axes(y, 2))
     return search
 end
 

src/nhs_grid.jl

@@ -194,12 +194,14 @@ end
 
 function initialize!(neighborhood_search::GridNeighborhoodSearch,
                      x::AbstractMatrix, y::AbstractMatrix;
-                     parallelization_backend = default_backend(x))
-    initialize_grid!(neighborhood_search, y; parallelization_backend)
+                     parallelization_backend = default_backend(x),
+                     eachindex_y = axes(y, 2))
+    initialize_grid!(neighborhood_search, y; parallelization_backend, eachindex_y)
 end
 
 function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::AbstractMatrix;
-                          parallelization_backend = default_backend(y))
+                          parallelization_backend = default_backend(y),
+                          eachindex_y = axes(y, 2))
     (; cell_list) = neighborhood_search
 
     empty!(cell_list)
@@ -210,8 +212,10 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::Abstra
         return neighborhood_search
     end
 
+    @boundscheck checkbounds(y, eachindex_y)
+
     # Ignore the parallelization backend here. This cannot be parallelized.
-    for point in axes(y, 2)
+    for point in eachindex_y
         # Get cell index of the point's cell
         point_coords = @inbounds extract_svector(y, Val(ndims(neighborhood_search)), point)
         cell = cell_coords(point_coords, neighborhood_search)
@@ -225,7 +229,8 @@ end
 
 function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                       ParallelUpdate},
-                          y::AbstractMatrix; parallelization_backend = default_backend(y))
+                          y::AbstractMatrix; parallelization_backend = default_backend(y),
+                          eachindex_y = axes(y, 2))
     (; cell_list) = neighborhood_search
 
     empty!(cell_list)
@@ -236,7 +241,9 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
         return neighborhood_search
     end
 
-    @threaded parallelization_backend for point in axes(y, 2)
+    @boundscheck checkbounds(y, eachindex_y)
+
+    @threaded parallelization_backend for point in eachindex_y
         # Get cell index of the point's cell
         point_coords = @inbounds extract_svector(y, Val(ndims(neighborhood_search)), point)
         cell = cell_coords(point_coords, neighborhood_search)
@@ -250,19 +257,30 @@ end
 
 function update!(neighborhood_search::GridNeighborhoodSearch,
                  x::AbstractMatrix, y::AbstractMatrix;
-                 points_moving = (true, true), parallelization_backend = default_backend(x))
+                 points_moving = (true, true), parallelization_backend = default_backend(x),
+                 eachindex_y = axes(y, 2))
     # The coordinates of the first set of points are irrelevant for this NHS.
     # Only update when the second set is moving.
     points_moving[2] || return neighborhood_search
 
-    update_grid!(neighborhood_search, y; parallelization_backend)
+    update_grid!(neighborhood_search, y; eachindex_y, parallelization_backend)
 end
 
 # Update only with neighbor coordinates
-function update_grid!(neighborhood_search::GridNeighborhoodSearch{NDIMS}, y::AbstractMatrix;
-                      parallelization_backend = default_backend(y)) where {NDIMS}
+function update_grid!(neighborhood_search::Union{GridNeighborhoodSearch{<:Any,
+                                                                        SerialIncrementalUpdate},
+                                                 GridNeighborhoodSearch{<:Any,
+                                                                        SemiParallelUpdate}},
+                      y::AbstractMatrix;
+                      parallelization_backend = default_backend(y),
+                      eachindex_y = axes(y, 2))
     (; cell_list, update_buffer) = neighborhood_search
 
+    if eachindex_y != axes(y, 2)
+        # Incremental update doesn't support inactive points
+        error("this neighborhood search/update strategy does not support inactive points")
+    end
+
     # Empty each thread's list
     @threaded parallelization_backend for i in eachindex(update_buffer)
         emptyat!(update_buffer, i)
@@ -302,33 +320,46 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch{NDIMS}, y::Abs
     return neighborhood_search
 end
 
-# The type annotation is to make Julia specialize on the type of the function.
-# Otherwise, unspecialized code will cause a lot of allocations and heavily impact performance.
-# See https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing
 @inline function mark_changed_cells!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                                  SemiParallelUpdate},
-                                     y, parallelization_backend) where {T}
-    (; cell_list) = neighborhood_search
+                                     y, parallelization_backend)
+    (; cell_list, update_buffer) = neighborhood_search
 
     # `each_cell_index(cell_list)` might return a `KeySet`, which has to be `collect`ed
-    # first to be able to be used in a threaded loop. This function takes care of that.
-    @threaded parallelization_backend for cell_index in each_cell_index_threadable(cell_list)
-        mark_changed_cell!(neighborhood_search, cell_index, y)
+    # first to support indexing.
+    eachcell = each_cell_index_threadable(cell_list)
+
+    # Use chunks (usually one per thread) to index into the update buffer.
+    # We cannot use `Iterators.partition` here, since the resulting iterator does not
+    # support indexing and therefore cannot be used in a threaded loop.
+    chunk_length = div(length(eachcell), length(update_buffer), RoundUp)
+
+    @threaded parallelization_backend for chunk_id in 1:length(update_buffer)
+        # Manual partitioning of `eachcell`
+        start = (chunk_length * (chunk_id - 1)) + 1
+        end_ = min(chunk_length * chunk_id, length(eachcell))
+
+        for i in start:end_
+            cell_index = eachcell[i]
+
+            mark_changed_cell!(neighborhood_search, cell_index, y, chunk_id)
+        end
     end
 end
 
 @inline function mark_changed_cells!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                                  SerialIncrementalUpdate},
-                                     y, _) where {T}
+                                     y, _)
     (; cell_list) = neighborhood_search
 
     # Ignore the parallelization backend here for `SerialIncrementalUpdate`.
     for cell_index in each_cell_index(cell_list)
-        mark_changed_cell!(neighborhood_search, cell_index, y)
+        # `chunk_id` is always `1` for `SerialIncrementalUpdate`
+        mark_changed_cell!(neighborhood_search, cell_index, y, 1)
     end
 end
 
-@inline function mark_changed_cell!(neighborhood_search, cell_index, y)
+@inline function mark_changed_cell!(neighborhood_search, cell_index, y, chunk_id)
     (; cell_list, update_buffer) = neighborhood_search
 
     for point in cell_list[cell_index]
@@ -340,7 +371,7 @@ end
         # These can be identical (see `DictionaryCellList`).
         if !is_correct_cell(cell_list, cell, cell_index)
             # Mark this cell and continue with the next one
-            pushat!(update_buffer, Threads.threadid(), cell_index)
+            pushat!(update_buffer, chunk_id, cell_index)
             break
         end
     end
@@ -349,9 +380,15 @@ end
 # Fully parallel incremental update with atomic push.
 function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                   ParallelIncrementalUpdate},
-                      y::AbstractMatrix; parallelization_backend = default_backend(y))
+                      y::AbstractMatrix; parallelization_backend = default_backend(y),
+                      eachindex_y = axes(y, 2))
     (; cell_list, update_buffer) = neighborhood_search
 
+    if eachindex_y != axes(y, 2)
+        # Incremental update doesn't support inactive points
+        error("this neighborhood search/update strategy does not support inactive points")
+    end
+
     # Note that we need two separate loops for adding and removing points.
     # `push_cell_atomic!` only guarantees thread-safety when different threads push
     # simultaneously, but it does not work when `deleteat_cell!` is called at the same time.
@@ -405,8 +442,9 @@ function update_grid!(neighborhood_search::Union{GridNeighborhoodSearch{<:Any,
                                                                         ParallelUpdate},
                                                  GridNeighborhoodSearch{<:Any,
                                                                         SerialUpdate}},
-                      y::AbstractMatrix; parallelization_backend = default_backend(y))
-    initialize_grid!(neighborhood_search, y; parallelization_backend)
+                      y::AbstractMatrix; parallelization_backend = default_backend(y),
+                      eachindex_y = axes(y, 2))
+    initialize_grid!(neighborhood_search, y; parallelization_backend, eachindex_y)
 end
 
 function check_collision(neighbor_cell_, neighbor_coords, cell_list, nhs)

src/nhs_precomputed.jl

@@ -53,14 +53,15 @@ end
 
 function initialize!(search::PrecomputedNeighborhoodSearch,
                      x::AbstractMatrix, y::AbstractMatrix;
-                     parallelization_backend = default_backend(x))
+                     parallelization_backend = default_backend(x),
+                     eachindex_y = axes(y, 2))
     (; neighborhood_search, neighbor_lists) = search
 
     # Initialize grid NHS
-    initialize!(neighborhood_search, x, y; parallelization_backend)
+    initialize!(neighborhood_search, x, y; eachindex_y, parallelization_backend)
 
     initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                               parallelization_backend)
+                               parallelization_backend, eachindex_y)
 end
 
 # WARNING! Experimental feature:
@@ -70,21 +71,22 @@ end
 # `parallelization_backend = KernelAbstractions.CPU()`, even though `x isa Array`.
 function update!(search::PrecomputedNeighborhoodSearch,
                  x::AbstractMatrix, y::AbstractMatrix;
-                 points_moving = (true, true), parallelization_backend = default_backend(x))
+                 points_moving = (true, true), parallelization_backend = default_backend(x),
+                 eachindex_y = axes(y, 2))
     (; neighborhood_search, neighbor_lists) = search
 
     # Update grid NHS
-    update!(neighborhood_search, x, y; points_moving, parallelization_backend)
+    update!(neighborhood_search, x, y; eachindex_y, points_moving, parallelization_backend)
 
     # Skip update if both point sets are static
     if any(points_moving)
         initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                   parallelization_backend)
+                                   parallelization_backend, eachindex_y)
     end
 end
 
 function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                    parallelization_backend)
+                                    parallelization_backend, eachindex_y)
     # Initialize neighbor lists
     empty!(neighbor_lists)
     resize!(neighbor_lists, size(x, 2))
@@ -93,8 +95,8 @@ function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
     end
 
     # Fill neighbor lists
-    foreach_point_neighbor(x, y, neighborhood_search;
-                           parallelization_backend) do point, neighbor, _, _
+    foreach_point_neighbor(x, y, neighborhood_search; parallelization_backend,
+                           points = eachindex_y) do point, neighbor, _, _
         push!(neighbor_lists[point], neighbor)
     end
 end

src/nhs_trivial.jl

@@ -33,13 +33,15 @@ end
 @inline requires_update(::TrivialNeighborhoodSearch) = (false, false)
 
 @inline function initialize!(search::TrivialNeighborhoodSearch, x, y;
-                             parallelization_backend = default_backend(x))
+                             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))
+                         parallelization_backend = default_backend(x),
+                         eachindex_y = axes(y, 2))
     return search
 end
 

test/nhs_grid.jl

@@ -145,7 +145,7 @@
         @test neighbors5 == [36, 37, 38, 43, 44, 45]
     end
 
-    @testset "Rectangular Point Cloud 3D" begin
+    @testset verbose=true "Rectangular Point Cloud 3D" begin
         #### Setup
         # Rectangle of equidistantly spaced points
         # from (x, y, z) = (-0.25, -0.25, -0.25) to (x, y, z) = (0.35, 0.35, 0.35).
@@ -188,6 +188,36 @@
         @test neighbors3 ==
               [115, 116, 117, 122, 123, 124, 129, 130, 131, 164, 165, 166, 171, 172,
             173, 178, 179, 180, 213, 214, 215, 220, 221, 222, 227, 228, 229]
+
+        update_strategies = (SerialUpdate(), ParallelUpdate())
+        @testset verbose=true "eachindex_y $update_strategy" for update_strategy in
+                                                                 update_strategies
+            # Test that `eachindex_y` is passed correctly to the neighborhood search.
+            # This requires `SerialUpdate` or `ParallelUpdate`.
+            min_corner = min.(minimum(coordinates1, dims = 2),
+                              minimum(coordinates2, dims = 2))
+            max_corner = max.(maximum(coordinates1, dims = 2),
+                              maximum(coordinates2, dims = 2))
+            cell_list = FullGridCellList(; min_corner, max_corner, search_radius)
+            nhs2 = GridNeighborhoodSearch{3}(; search_radius, n_points, update_strategy,
+                                             cell_list)
+
+            # Initialize with all points
+            initialize!(nhs2, coordinates1, coordinates1)
+
+            # Update with a subset of points
+            update!(nhs2, coordinates2, coordinates2; eachindex_y = 120:220)
+
+            neighbors2 = sort(collect(PointNeighbors.eachneighbor(point_position1, nhs2)))
+            neighbors3 = sort(collect(PointNeighbors.eachneighbor(point_position2, nhs2)))
+
+            # Check that the neighbors are the intersection of the previous neighbors
+            # with the `eachindex_y` range.
+            @test neighbors2 == Int[]
+            @test neighbors3 ==
+                  [122, 123, 124, 129, 130, 131, 164, 165, 166, 171, 172, 173,
+                178, 179, 180, 213, 214, 215, 220]
+        end
     end
 
     @testset verbose=true "Periodicity" begin