Add kwarg points_active to update! and initialize!

Author: efaulhaber

src/neighborhood_search.jl

@@ -35,7 +35,7 @@ See also [`update!`](@ref).
 """
 @inline function initialize!(search::AbstractNeighborhoodSearch, x, y;
                              parallelization_backend = default_backend(x),
-                             eachindex_y = axes(y, 2))
+                             eachindex_y = axes(y, 2), points_active = nothing)
     return search
 end
 
@@ -73,7 +73,7 @@ See also [`initialize!`](@ref).
 @inline function update!(search::AbstractNeighborhoodSearch, x, y;
                          points_moving = (true, true),
                          parallelization_backend = default_backend(x),
-                         eachindex_y = axes(y, 2))
+                         eachindex_y = axes(y, 2), points_active = nothing)
     return search
 end
 

src/nhs_grid.jl

@@ -201,7 +201,7 @@ end
 
 function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::AbstractMatrix;
                           parallelization_backend = default_backend(y),
-                          eachindex_y = axes(y, 2))
+                          eachindex_y = axes(y, 2), points_active = nothing)
     (; cell_list) = neighborhood_search
 
     empty!(cell_list)
@@ -216,12 +216,15 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::Abstra
 
     # Ignore the parallelization backend here. This cannot be parallelized.
     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)
-
-        # Add point to corresponding cell
-        push_cell!(cell_list, cell, point)
+        # This check is optimized away when `points_active` is `nothing`
+        if point_active(points_active, point)
+            # 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)
+
+            # Add point to corresponding cell
+            push_cell!(cell_list, cell, point)
+        end
     end
 
     return neighborhood_search
@@ -230,7 +233,7 @@ end
 function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                       ParallelUpdate},
                           y::AbstractMatrix; parallelization_backend = default_backend(y),
-                          eachindex_y = axes(y, 2))
+                          eachindex_y = axes(y, 2), points_active = nothing)
     (; cell_list) = neighborhood_search
 
     empty!(cell_list)
@@ -244,12 +247,15 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
     @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)
-
-        # Add point to corresponding cell
-        push_cell_atomic!(cell_list, cell, point)
+        # This check is optimized away when `points_active` is `nothing`
+        if point_active(points_active, point)
+            # 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)
+
+            # Add point to corresponding cell
+            push_cell_atomic!(cell_list, cell, point)
+        end
     end
 
     return neighborhood_search
@@ -258,12 +264,13 @@ end
 function update!(neighborhood_search::GridNeighborhoodSearch,
                  x::AbstractMatrix, y::AbstractMatrix;
                  points_moving = (true, true), parallelization_backend = default_backend(x),
-                 eachindex_y = axes(y, 2))
+                 eachindex_y = axes(y, 2), points_active = nothing)
     # 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; eachindex_y, parallelization_backend)
+    update_grid!(neighborhood_search, y; eachindex_y, points_active,
+                 parallelization_backend)
 end
 
 # Update only with neighbor coordinates
@@ -273,10 +280,10 @@ function update_grid!(neighborhood_search::Union{GridNeighborhoodSearch{<:Any,
                                                                         SemiParallelUpdate}},
                       y::AbstractMatrix;
                       parallelization_backend = default_backend(y),
-                      eachindex_y = axes(y, 2))
+                      eachindex_y = axes(y, 2), points_active = nothing)
     (; cell_list, update_buffer) = neighborhood_search
 
-    if eachindex_y != axes(y, 2)
+    if eachindex_y != axes(y, 2) || points_active !== nothing
         # Incremental update doesn't support inactive points
         error("this neighborhood search/update strategy does not support inactive points")
     end
@@ -381,10 +388,10 @@ end
 function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                   ParallelIncrementalUpdate},
                       y::AbstractMatrix; parallelization_backend = default_backend(y),
-                      eachindex_y = axes(y, 2))
+                      eachindex_y = axes(y, 2), points_active = nothing)
     (; cell_list, update_buffer) = neighborhood_search
 
-    if eachindex_y != axes(y, 2)
+    if eachindex_y != axes(y, 2) || points_active !== nothing
         # Incremental update doesn't support inactive points
         error("this neighborhood search/update strategy does not support inactive points")
     end
@@ -445,8 +452,9 @@ function update_grid!(neighborhood_search::Union{GridNeighborhoodSearch{<:Any,
                                                  GridNeighborhoodSearch{<:Any,
                                                                         SerialUpdate}},
                       y::AbstractMatrix; parallelization_backend = default_backend(y),
-                      eachindex_y = axes(y, 2))
-    initialize_grid!(neighborhood_search, y; parallelization_backend, eachindex_y)
+                      eachindex_y = axes(y, 2), points_active = nothing)
+    initialize_grid!(neighborhood_search, y; parallelization_backend,
+                     eachindex_y, points_active)
 end
 
 # Specialized version of the function in `neighborhood_search.jl`, which is faster

src/nhs_precomputed.jl

@@ -72,21 +72,22 @@ end
 function update!(search::PrecomputedNeighborhoodSearch,
                  x::AbstractMatrix, y::AbstractMatrix;
                  points_moving = (true, true), parallelization_backend = default_backend(x),
-                 eachindex_y = axes(y, 2))
+                 eachindex_y = axes(y, 2), points_active = nothing)
     (; neighborhood_search, neighbor_lists) = search
 
     # Update grid NHS
-    update!(neighborhood_search, x, y; eachindex_y, points_moving, parallelization_backend)
+    update!(neighborhood_search, x, y; eachindex_y, points_moving, points_active,
+            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, eachindex_y)
+                                   parallelization_backend, eachindex_y, points_active)
     end
 end
 
 function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                    parallelization_backend, eachindex_y)
+                                    parallelization_backend, eachindex_y, points_active)
     # Initialize neighbor lists
     empty!(neighbor_lists)
     resize!(neighbor_lists, size(x, 2))
@@ -95,8 +96,9 @@ function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
     end
 
     # Fill neighbor lists
-    foreach_point_neighbor(x, y, neighborhood_search; parallelization_backend,
-                           points = eachindex_y) do point, neighbor, _, _
+    foreach_point_neighbor(x, y, neighborhood_search;
+                           parallelization_backend, points = eachindex_y,
+                           points_active) do point, neighbor, _, _
         push!(neighbor_lists[point], neighbor)
     end
 end

src/nhs_trivial.jl

@@ -34,14 +34,14 @@ end
 
 @inline function initialize!(search::TrivialNeighborhoodSearch, x, y;
                              parallelization_backend = default_backend(x),
-                             eachindex_y = axes(y, 2))
+                             eachindex_y = axes(y, 2), points_active = nothing)
     return search
 end
 
 @inline function update!(search::TrivialNeighborhoodSearch, x, y;
                          points_moving = (true, true),
                          parallelization_backend = default_backend(x),
-                         eachindex_y = axes(y, 2))
+                         eachindex_y = axes(y, 2), points_active = nothing)
     return search
 end