Use contiguous backend for FullGridCellList

src/cell_lists/full_grid.jl

@@ -1,5 +1,7 @@
 """
-    FullGridCellList(; min_corner, max_corner, search_radius = 0.0, periodicity = false)
+    FullGridCellList(; min_corner, max_corner, search_radius = 0.0,
+                     periodicity = false, backend = DynamicVectorOfVectors{Int32},
+                     max_points_per_cell = 100)
 
 A simple cell list implementation where each (empty or non-empty) cell of a rectangular
 (axis-aligned) domain is assigned a list of points.
@@ -19,6 +21,13 @@ See [`copy_neighborhood_search`](@ref) for more details.
                          neighborhood search. When using [`copy_neighborhood_search`](@ref),
                          this option can be ignored an will be set automatically depending
                          on the periodicity of the neighborhood search.
+- `backend = DynamicVectorOfVectors{Int32}`: Type of the data structure to store the actual
+    cell lists. Can be
+    - `Vector{Vector{Int32}}`: Scattered memory, but very memory-efficient.
+    - `DynamicVectorOfVectors{Int32}`: Contiguous memory, optimizing cache-hits.
+- `max_points_per_cell = 100`: Maximum number of points per cell. This will be used to
+                               allocate the `DynamicVectorOfVectors`. It is not used with
+                               the `Vector{Vector{Int32}}` backend.
 """
 struct FullGridCellList{C, LI, MC} <: AbstractCellList
     cells          :: C
@@ -31,7 +40,8 @@ struct FullGridCellList{C, LI, MC} <: AbstractCellList
 end
 
 function FullGridCellList(; min_corner, max_corner, search_radius = 0.0,
-                          periodicity = false)
+                          periodicity = false, backend = DynamicVectorOfVectors{Int32},
+                          max_points_per_cell = 100)
     if search_radius < eps()
         # Create an empty "template" cell list to be used with `copy_cell_list`
         cells = nothing
@@ -53,15 +63,33 @@ function FullGridCellList(; min_corner, max_corner, search_radius = 0.0,
         linear_indices = LinearIndices(Tuple(n_cells_per_dimension))
         min_cell = Tuple(floor_to_int.(min_corner ./ search_radius))
 
-        cells = [Int32[] for _ in 1:prod(n_cells_per_dimension)]
+        cells = construct_backend(backend, n_cells_per_dimension, max_points_per_cell)
     end
 
     return FullGridCellList{typeof(cells), typeof(linear_indices),
                             typeof(min_cell)}(cells, linear_indices, min_cell)
 end
 
+function construct_backend(::Type{Vector{Vector{T}}}, size, max_points_per_cell) where {T}
+    return [T[] for _ in 1:prod(size)]
+end
+
+function construct_backend(cells::Type{DynamicVectorOfVectors{T}}, size,
+                           max_points_per_cell) where {T}
+    cells = DynamicVectorOfVectors{T}(max_outer_length = prod(size),
+                                      max_inner_length = max_points_per_cell)
+    resize!(cells, prod(size))
+
+    return cells
+end
+
 function Base.empty!(cell_list::FullGridCellList)
-    Base.empty!.(cell_list.cells)
+    (; cells) = cell_list
+
+    # `Base.empty!.(cells)`, but for all backends
+    for i in eachindex(cells)
+        emptyat!(cells, i)
+    end
 
     return cell_list
 end
@@ -72,7 +100,12 @@ function Base.empty!(cell_list::FullGridCellList{Nothing})
 end
 
 function push_cell!(cell_list::FullGridCellList, cell, particle)
-    push!(cell_list[cell], particle)
+    (; cells) = cell_list
+
+    # `push!(cell_list[cell], particle)`, but for all backends
+    pushat!(cells, cell_index(cell_list, cell), particle)
+
+    return cell_list
 end
 
 function push_cell!(cell_list::FullGridCellList{Nothing}, cell, particle)
@@ -81,7 +114,10 @@ function push_cell!(cell_list::FullGridCellList{Nothing}, cell, particle)
 end
 
 function deleteat_cell!(cell_list::FullGridCellList, cell, i)
-    deleteat!(cell_list[cell], i)
+    (; cells) = cell_list
+
+    # `deleteat!(cell_list[cell], i)`, but for all backends
+    deleteatat!(cells, cell_index(cell_list, cell), i)
 end
 
 @inline each_cell_index(cell_list::FullGridCellList) = eachindex(cell_list.cells)
@@ -91,20 +127,22 @@ function each_cell_index(cell_list::FullGridCellList{Nothing})
     throw(UndefRefError("`search_radius` is not defined for this cell list"))
 end
 
-@inline function Base.getindex(cell_list::FullGridCellList, cell::Tuple)
-    (; cells, linear_indices, min_cell) = cell_list
+@inline function cell_index(cell_list::FullGridCellList, cell::Tuple)
+    (; linear_indices, min_cell) = cell_list
 
-    return cells[linear_indices[(cell .- min_cell .+ 1)...]]
+    return linear_indices[(cell .- min_cell .+ 1)...]
 end
 
-@inline function Base.getindex(cell_list::FullGridCellList, i::Integer)
-    return cell_list.cells[i]
-end
+@inline cell_index(::FullGridCellList, cell::Integer) = cell
 
-@inline function is_correct_cell(cell_list::FullGridCellList, cell_coords, cell_index)
-    (; linear_indices, min_cell) = cell_list
+@inline function Base.getindex(cell_list::FullGridCellList, cell)
+    (; cells) = cell_list
+
+    return cells[cell_index(cell_list, cell)]
+end
 
-    return linear_indices[(cell_coords .- min_cell .+ 1)...] == cell_index
+@inline function is_correct_cell(cell_list::FullGridCellList, cell_coords, cell_index_)
+    return cell_index(cell_list, cell_coords) == cell_index_
 end
 
 @inline index_type(::FullGridCellList) = Int32

src/nhs_grid.jl

@@ -137,12 +137,12 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch{NDIMS},
     update_grid!(neighborhood_search, i -> extract_svector(y, Val(NDIMS), i))
 end
 
-# Modify the existing hash table by moving points into their new cells
+# Modify the existing cell lists by moving points into their new cells
 function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
     (; cell_list, update_buffer, threaded_update) = neighborhood_search
 
     # Empty each thread's list
-    for i in eachindex(update_buffer)
+    @threaded for i in eachindex(update_buffer)
         emptyat!(update_buffer, i)
     end
 
@@ -154,40 +154,45 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
         for cell_index in update_buffer[j]
             points = cell_list[cell_index]
 
-            # Find all points whose coordinates do not match this cell
-            moved_point_indices = (i for i in eachindex(points)
-                                   if !is_correct_cell(cell_list,
-                                                       cell_coords(coords_fun(points[i]),
-                                                                   neighborhood_search),
-                                                       cell_index))
-
-            # Add moved points to new cell
-            for i in moved_point_indices
+            # Find all points whose coordinates do not match this cell.
+            #
+            # WARNING!!!
+            # The `DynamicVectorOfVectors` requires this loop to be **in descending order**.
+            # `deleteat_cell!(..., i)` will change the order of points that come after `i`.
+            for i in reverse(eachindex(points))
                 point = points[i]
-                new_cell_coords = cell_coords(coords_fun(point), neighborhood_search)
+                cell_coords_ = cell_coords(coords_fun(point), neighborhood_search)
 
-                # Add point to corresponding cell or create cell if it does not exist
-                push_cell!(cell_list, new_cell_coords, point)
-            end
+                if !is_correct_cell(cell_list, cell_coords_, cell_index)
+                    new_cell_coords = cell_coords(coords_fun(point), neighborhood_search)
 
-            # Remove moved points from this cell
-            deleteat_cell!(cell_list, cell_index, moved_point_indices)
+                    # Add point to new cell or create cell if it does not exist
+                    push_cell!(cell_list, new_cell_coords, point)
+
+                    # Remove moved point from this cell
+                    deleteat_cell!(cell_list, cell_index, i)
+                end
+            end
         end
     end
 
     return neighborhood_search
 end
 
-@inline function mark_changed_cell!(neighborhood_search, cell_list, coords_fun,
-                                    threaded_update::Val{true})
-    # `collect` the keyset to be able to loop over it with `@threaded`
+# 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_cell!(neighborhood_search, cell_list, coords_fun::T,
+                                    threaded_update::Val{true}) where {T}
+    # `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 for cell_index in each_cell_index_threadable(cell_list)
         mark_changed_cell!(neighborhood_search, cell_index, coords_fun)
     end
 end
 
-@inline function mark_changed_cell!(neighborhood_search, cell_list, coords_fun,
-                                    threaded_update::Val{false})
+@inline function mark_changed_cell!(neighborhood_search, cell_list, coords_fun::T,
+                                    threaded_update::Val{false}) where {T}
     for cell_index in each_cell_index(cell_list)
         mark_changed_cell!(neighborhood_search, cell_index, coords_fun)
     end

src/vector_of_vectors.jl

@@ -60,6 +60,12 @@ end
     return vov
 end
 
+@inline function pushat!(vov::Vector{<:Vector{<:Any}}, i, value)
+    push!(vov[i], value)
+
+    return vov
+end
+
 # `deleteat!(vov[i], j)`
 @inline function deleteatat!(vov::DynamicVectorOfVectors, i, j)
     (; backend, lengths) = vov
@@ -79,6 +85,12 @@ end
     return vov
 end
 
+@inline function deleteatat!(vov::Vector{<:Vector{<:Any}}, i, j)
+    deleteat!(vov[i], j)
+
+    return vov
+end
+
 @inline function Base.empty!(vov::DynamicVectorOfVectors)
     # Move all pointers to the beginning
     vov.lengths .= zero(Int32)
@@ -94,3 +106,15 @@ end
 
     return vov
 end
+
+@inline function emptyat!(vov::Vector{<:Vector{<:Any}}, i)
+    Base.empty!(vov[i])
+end
+
+@inline function Base.resize!(vov::DynamicVectorOfVectors, n)
+    # Make sure that all newly added vectors are empty
+    vov.lengths[(length(vov) + 1):n] .= zero(Int32)
+    vov.length_[] = n
+
+    return vov
+end

test/neighborhood_search.jl

@@ -38,17 +38,26 @@
             n_points = size(coords, 2)
             search_radius = 0.1
 
+            min_corner = periodic_boxes[i].min_corner
+            max_corner = max_corner = periodic_boxes[i].max_corner
+
             neighborhood_searches = [
                 TrivialNeighborhoodSearch{NDIMS}(; search_radius, eachpoint = 1:n_points,
                                                  periodic_box = periodic_boxes[i]),
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               periodic_box = periodic_boxes[i]),
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               periodic_box = periodic_boxes[i],
-                                              cell_list = FullGridCellList(;
-                                                                           min_corner = periodic_boxes[i].min_corner,
-                                                                           max_corner = periodic_boxes[i].max_corner,
+                                              cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
+                                                                           search_radius,
+                                                                           periodicity = true)),
+                GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                              periodic_box = periodic_boxes[i],
+                                              cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
                                                                            search_radius,
+                                                                           backend = Vector{Vector{Int32}},
                                                                            periodicity = true)),
                 PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                                      periodic_box = periodic_boxes[i]),
@@ -57,7 +66,8 @@
             names = [
                 "`TrivialNeighborhoodSearch`",
                 "`GridNeighborhoodSearch`",
-                "`GridNeighborhoodSearch` with `FullGridCellList",
+                "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors`",
+                "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
                 "`PrecomputedNeighborhoodSearch`",
             ]
 
@@ -68,6 +78,10 @@
                 GridNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i],
                                               cell_list = FullGridCellList(min_corner = periodic_boxes[i].min_corner,
                                                                            max_corner = periodic_boxes[i].max_corner)),
+                GridNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i],
+                                              cell_list = FullGridCellList(min_corner = periodic_boxes[i].min_corner,
+                                                                           max_corner = periodic_boxes[i].max_corner,
+                                                                           backend = Vector{Vector{Int32}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i]),
             ]
             copied_nhs = copy_neighborhood_search.(template_nhs, search_radius, n_points)
@@ -138,23 +152,29 @@
                 append!(neighbors_expected[point], neighbor)
             end
 
+            # Expand the domain by `search_radius`, as we need the neighboring cells of
+            # the minimum and maximum coordinates as well.
             min_corner = minimum(coords, dims = 2) .- search_radius
             max_corner = maximum(coords, dims = 2) .+ search_radius
 
             neighborhood_searches = [
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points),
-                # Expand the domain by `search_radius`, as we need the neighboring cells of
-                # the minimum and maximum coordinates as well.
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               cell_list = FullGridCellList(; min_corner,
                                                                            max_corner,
                                                                            search_radius)),
+                GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                              cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
+                                                                           search_radius,
+                                                                           backend = Vector{Vector{Int}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points),
             ]
 
             names = [
                 "`GridNeighborhoodSearch`",
-                "`GridNeighborhoodSearch` with `FullGridCellList`",
+                "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors`",
+                "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
                 "`PrecomputedNeighborhoodSearch`",
             ]
 
@@ -163,6 +183,9 @@
                 GridNeighborhoodSearch{NDIMS}(),
                 GridNeighborhoodSearch{NDIMS}(cell_list = FullGridCellList(; min_corner,
                                                                            max_corner)),
+                GridNeighborhoodSearch{NDIMS}(cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
+                                                                           backend = Vector{Vector{Int32}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(),
             ]
             copied_nhs = copy_neighborhood_search.(template_nhs, search_radius, n_points)