Add proper support for inactive points (#107)

* Remove `requires_resizing` again

* Add proper support for inactive points

* Fix

* Improve error message for domain check on CPUs

* Improve error message

* Fix `update!` on GPUs

* Fix tests

* Update docs

* Fix GPU update

* Fix merge

* Fix tests

* Remove duplicate function and unused type annotation

* Add tests

---------

Co-authored-by: Niklas Neher <[email protected]>

Author: efaulhaber

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, y::AbstractMatrix;
-                      parallelization_backend = default_backend(y))
+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)
@@ -362,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.
@@ -420,8 +444,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
 
 # Specialized version of the function in `neighborhood_search.jl`, which is faster

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