Store coordinates together with the points in a cell list

src/cell_lists/cell_lists.jl

@@ -1,4 +1,6 @@
-abstract type AbstractCellList end
+abstract type AbstractCellList{T} end
+
+Base.eltype(::AbstractCellList{T}) where {T} = T
 
 # For the `DictionaryCellList`, this is a `KeySet`, which has to be `collect`ed first to be
 # able to be used in a threaded loop.

src/cell_lists/dictionary.jl

@@ -1,8 +1,8 @@
 """
-    DictionaryCellList{NDIMS}()
+    DictionaryCellList{T, NDIMS}()
 
 A simple cell list implementation where a cell index `(i, j)` or `(i, j, k)` is mapped to a
-`Vector{Int}` by a `Dict`.
+`Vector{T}` by a `Dict`.
 By using a dictionary, which only stores non-empty cells, the domain is
 potentially infinite.
 
@@ -11,17 +11,19 @@ for integer tuples, nor does it use a contiguous memory layout.
 Consequently, this cell list is not GPU-compatible.
 
 # Arguments
+- `T`: Either an `Integer` type, or `PointWithCoordinates{NDIMS, ELTYPE}`.
+       See [`PointWithCoordinates`](@ref) for more details.
 - `NDIMS`: Number of dimensions.
 """
-struct DictionaryCellList{NDIMS} <: AbstractCellList
-    hashtable    :: Dict{NTuple{NDIMS, Int}, Vector{Int}}
-    empty_vector :: Vector{Int} # Just an empty vector (used in `eachneighbor`)
+struct DictionaryCellList{T, NDIMS} <: AbstractCellList{T}
+    hashtable    :: Dict{NTuple{NDIMS, Int}, Vector{T}}
+    empty_vector :: Vector{T} # Just an empty vector (used in `getindex`)
 
-    function DictionaryCellList{NDIMS}() where {NDIMS}
-        hashtable = Dict{NTuple{NDIMS, Int}, Vector{Int}}()
-        empty_vector = Int[]
+    function DictionaryCellList{T, NDIMS}() where {T, NDIMS}
+        hashtable = Dict{NTuple{NDIMS, Int}, Vector{T}}()
+        empty_vector = T[]
 
-        new{NDIMS}(hashtable, empty_vector)
+        new{T, NDIMS}(hashtable, empty_vector)
     end
 end
 
@@ -37,7 +39,7 @@ function push_cell!(cell_list::DictionaryCellList, cell, point)
     (; hashtable) = cell_list
 
     if haskey(hashtable, cell)
-        append!(hashtable[cell], point)
+        push!(hashtable[cell], point)
     else
         hashtable[cell] = [point]
     end
@@ -80,9 +82,9 @@ end
     return cell_coords == cell_index
 end
 
-@inline index_type(::DictionaryCellList{NDIMS}) where {NDIMS} = NTuple{NDIMS, Int}
+@inline index_type(::DictionaryCellList{<:Any, NDIMS}) where {NDIMS} = NTuple{NDIMS, Int}
 
-function copy_cell_list(::DictionaryCellList{NDIMS}, search_radius,
-                        periodic_box) where {NDIMS}
-    return DictionaryCellList{NDIMS}()
+function copy_cell_list(::DictionaryCellList{T, NDIMS}, search_radius,
+                        periodic_box) where {T, NDIMS}
+    return DictionaryCellList{T, NDIMS}()
 end

src/cell_lists/full_grid.jl

@@ -17,33 +17,43 @@ See [`copy_neighborhood_search`](@ref) for more details.
 - `max_corner`: Coordinates of the domain corner in positive coordinate directions.
 - `search_radius = 0.0`: Search radius of the neighborhood search, which will determine the
                          cell size. Use the default of `0.0` to create a template (see above).
-- `periodicity = false`: Set to `true` when using a [`PeriodicBox`](@ref) with the
-                         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.
+    - `Vector{Vector{T}}`: Scattered memory, but very memory-efficient.
+    - `DynamicVectorOfVectors{T}`: Contiguous memory, optimizing cache-hits.
+    `T` can be either an `Integer` type, or `PointWithCoordinates{NDIMS, ELTYPE}`.
+    See [`PointWithCoordinates`](@ref) for more details.
 - `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.
+                               the `Vector{Vector{T}}` backend.
 """
-struct FullGridCellList{C, LI, MINC, MAXC} <: AbstractCellList
+struct FullGridCellList{T, C, LI, MINC, MAXC} <: AbstractCellList{T}
     cells          :: C
     linear_indices :: LI
     min_corner     :: MINC
     max_corner     :: MAXC
+
+    function FullGridCellList(cells::Vector{Vector{T}}, linear_indices,
+                              min_corner, max_corner) where {T}
+        new{T, typeof(cells), typeof(linear_indices), typeof(min_corner),
+            typeof(max_corner)}(cells, linear_indices, min_corner, max_corner)
+    end
+
+    function FullGridCellList(cells::DynamicVectorOfVectors{T}, linear_indices,
+                              min_corner, max_corner) where {T}
+        new{T, typeof(cells), typeof(linear_indices), typeof(min_corner),
+            typeof(max_corner)}(cells, linear_indices, min_corner, max_corner)
+    end
 end
 
-function supported_update_strategies(::FullGridCellList{<:DynamicVectorOfVectors})
+function supported_update_strategies(::FullGridCellList{<:Any, <:DynamicVectorOfVectors})
     return (ParallelUpdate, SemiParallelUpdate, SerialUpdate)
 end
 
 supported_update_strategies(::FullGridCellList) = (SemiParallelUpdate, SerialUpdate)
 
 function FullGridCellList(; min_corner, max_corner, search_radius = 0.0,
-                          periodicity = false, backend = DynamicVectorOfVectors{Int32},
+                          backend = DynamicVectorOfVectors{Int32},
                           max_points_per_cell = 100)
     # Pad domain to avoid 0 in cell indices due to rounding errors.
     # We can't just use `eps()`, as one might use lower precision types.
@@ -112,7 +122,7 @@ function Base.empty!(cell_list::FullGridCellList)
     return cell_list
 end
 
-function Base.empty!(cell_list::FullGridCellList{Nothing})
+function Base.empty!(::FullGridCellList{<:Any, Nothing})
     # This is an empty "template" cell list to be used with `copy_cell_list`
     error("`search_radius` is not defined for this cell list")
 end
@@ -126,7 +136,7 @@ function push_cell!(cell_list::FullGridCellList, cell, particle)
     return cell_list
 end
 
-function push_cell!(cell_list::FullGridCellList{Nothing}, cell, particle)
+function push_cell!(::FullGridCellList{<:Any, Nothing}, cell, particle)
     # This is an empty "template" cell list to be used with `copy_cell_list`
     error("`search_radius` is not defined for this cell list")
 end
@@ -151,7 +161,7 @@ end
 
 @inline each_cell_index(cell_list::FullGridCellList) = eachindex(cell_list.cells)
 
-function each_cell_index(cell_list::FullGridCellList{Nothing})
+function each_cell_index(::FullGridCellList{<:Any, Nothing})
     # This is an empty "template" cell list to be used with `copy_cell_list`
     error("`search_radius` is not defined for this cell list")
 end

src/nhs_grid.jl

@@ -66,7 +66,7 @@ end
 
 function GridNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
                                        periodic_box = nothing,
-                                       cell_list = DictionaryCellList{NDIMS}(),
+                                       cell_list = DictionaryCellList{Int, NDIMS}(),
                                        update_strategy = nothing) where {NDIMS}
     if isnothing(update_strategy)
         # Automatically choose best available update option for this cell list
@@ -103,6 +103,29 @@ function GridNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
                                   cell_size, update_buffer, update_strategy)
 end
 
+struct PointWithCoordinates{NDIMS, ELTYPE}
+    index       :: Int32
+    coordinates :: SVector{NDIMS, ELTYPE}
+end
+
+@inline coordinates(point::Integer, coords_fun::Function) = coords_fun(point)
+
+@inline function coordinates(point::Integer, x::AbstractMatrix, nhs::GridNeighborhoodSearch)
+    return coordinates(point, x, Val(ndims(nhs)))
+end
+
+@inline function coordinates(point::Integer, x::AbstractMatrix, ::Val{NDIMS}) where {NDIMS}
+    return extract_svector(x, Val(NDIMS), point)
+end
+
+@inline coordinates(point::PointWithCoordinates, args...) = point.coordinates
+
+@inline assemble_point(::Type{<:Integer}, point_index, coordinates) = point_index
+@inline assemble_point(type, point_index, coordinates) = type(point_index, coordinates)
+
+@inline index(i::Integer) = i
+@inline index(point::PointWithCoordinates) = point.index
+
 """
     ParallelUpdate()
 
@@ -167,11 +190,14 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::Abstra
 
     empty!(cell_list)
 
-    for point in axes(y, 2)
+    for point_index in axes(y, 2)
         # Get cell index of the point's cell
-        point_coords = extract_svector(y, Val(ndims(neighborhood_search)), point)
+        point_coords = extract_svector(y, Val(ndims(neighborhood_search)), point_index)
         cell = cell_coords(point_coords, neighborhood_search)
 
+        # Get either point index or store point together with its coordinates
+        point = assemble_point(eltype(cell_list), point_index, point_coords)
+
         # Add point to corresponding cell
         push_cell!(cell_list, cell, point)
     end
@@ -232,11 +258,10 @@ function update_grid!(neighborhood_search::Union{GridNeighborhoodSearch{<:Any,
             # `deleteat_cell!(..., i)` will change the order of points that come after `i`.
             for i in reverse(eachindex(points))
                 point = points[i]
-                cell_coords_ = cell_coords(coords_fun(point), neighborhood_search)
-
-                if !is_correct_cell(cell_list, cell_coords_, cell_index)
-                    new_cell_coords = cell_coords(coords_fun(point), neighborhood_search)
+                new_cell_coords = cell_coords(coordinates(point, coords_fun),
+                                              neighborhood_search)
 
+                if !is_correct_cell(cell_list, new_cell_coords, cell_index)
                     # Add point to new cell or create cell if it does not exist
                     push_cell!(cell_list, new_cell_coords, point)
 
@@ -279,7 +304,7 @@ end
     (; cell_list, update_buffer) = neighborhood_search
 
     for point in cell_list[cell_index]
-        cell = cell_coords(coords_fun(point), neighborhood_search)
+        cell = cell_coords(coordinates(point, coords_fun), neighborhood_search)
 
         # `cell` is a tuple, `cell_index` is the linear index used internally by the
         # cell list to store cells inside `cell`.
@@ -305,11 +330,10 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any, Paralle
     # Add points to new cells
     @threaded parallelization_backend for cell_index in each_cell_index_threadable(cell_list)
         for point in cell_list[cell_index]
-            cell_coords_ = cell_coords(coords_fun(point), neighborhood_search)
-
-            if !is_correct_cell(cell_list, cell_coords_, cell_index)
-                new_cell_coords = cell_coords(coords_fun(point), neighborhood_search)
+            new_cell_coords = cell_coords(coordinates(point, coords_fun),
+                                          neighborhood_search)
 
+            if !is_correct_cell(cell_list, new_cell_coords, cell_index)
                 # Add point to new cell or create cell if it does not exist
                 push_cell_atomic!(cell_list, new_cell_coords, point)
             end
@@ -325,9 +349,10 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any, Paralle
         # `deleteat_cell!(..., i)` will change the order of points that come after `i`.
         for i in reverse(eachindex(points))
             point = points[i]
-            cell_coords_ = cell_coords(coords_fun(point), neighborhood_search)
+            new_cell_coords = cell_coords(coordinates(point, coords_fun),
+                                          neighborhood_search)
 
-            if !is_correct_cell(cell_list, cell_coords_, cell_index)
+            if !is_correct_cell(cell_list, new_cell_coords, cell_index)
                 # Remove moved point from this cell
                 deleteat_cell!(cell_list, cell_index, i)
             end
@@ -338,19 +363,20 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any, Paralle
 end
 
 @inline function foreach_neighbor(f, system_coords, neighbor_system_coords,
-                                  neighborhood_search::GridNeighborhoodSearch, point;
+                                  neighborhood_search::GridNeighborhoodSearch, point_index;
                                   search_radius = search_radius(neighborhood_search))
     (; periodic_box) = neighborhood_search
 
-    point_coords = extract_svector(system_coords, Val(ndims(neighborhood_search)), point)
+    point_coords = extract_svector(system_coords, Val(ndims(neighborhood_search)),
+                                   point_index)
     cell = cell_coords(point_coords, neighborhood_search)
 
     for neighbor_cell_ in neighboring_cells(cell, neighborhood_search)
         neighbor_cell = Tuple(neighbor_cell_)
 
         for neighbor in points_in_cell(neighbor_cell, neighborhood_search)
-            neighbor_coords = extract_svector(neighbor_system_coords,
-                                              Val(ndims(neighborhood_search)), neighbor)
+            neighbor_coords = coordinates(neighbor, neighbor_system_coords,
+                                          neighborhood_search)
 
             pos_diff = point_coords - neighbor_coords
             distance2 = dot(pos_diff, pos_diff)
@@ -364,7 +390,7 @@ end
 
                 # Inline to avoid loss of performance
                 # compared to not using `foreach_point_neighbor`.
-                @inline f(point, neighbor, pos_diff, distance)
+                @inline f(point_index, index(neighbor), pos_diff, distance)
             end
         end
     end