Add "full grid" cell list with a separate list for each cell of a rectangular domain (trixi-framework#25)

* Add `FullGridCellList`

* Fix tests

* Reformat

* Fix copying

* Use kwargs

* Add docs

* Fix docs

* Fix typo

* Fix variable names

* Fix `is_correct_cell`

* Avoid allocations from unnecessary `collect`

* Explicitly mention that rectangular domains have to be axis-aligned

* Implement suggestions

---------

Co-authored-by: Sven Berger <[email protected]>

Author: efaulhaber

src/PointNeighbors.jl

@@ -15,6 +15,7 @@ include("nhs_precomputed.jl")
 
 export foreach_point_neighbor, foreach_neighbor
 export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch
+export DictionaryCellList, FullGridCellList
 export initialize!, update!, initialize_grid!, update_grid!
 export PeriodicBox, copy_neighborhood_search
 

src/cell_lists/cell_lists.jl

@@ -1 +1,8 @@
+abstract type AbstractCellList end
+
+# For the `DictionaryCellList`, this is a `KeySet`, which has to be `collect`ed first to be
+# able to be used in a threaded loop.
+@inline each_cell_index_threadable(cell_list::AbstractCellList) = each_cell_index(cell_list)
+
 include("dictionary.jl")
+include("full_grid.jl")

src/cell_lists/dictionary.jl

@@ -1,4 +1,19 @@
-struct DictionaryCellList{NDIMS}
+"""
+    DictionaryCellList{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`.
+By using a dictionary, which only stores non-empty cells, the domain is
+potentially infinite.
+
+This implementation is very simple, but it neither uses an optimized hash function
+for integer tuples, nor does it use a contiguous memory layout.
+Consequently, this cell list is not GPU-compatible.
+
+# Arguments
+- `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`)
 
@@ -43,7 +58,13 @@ function delete_cell!(cell_list, cell)
     delete!(cell_list.hashtable, cell)
 end
 
-@inline eachcell(cell_list::DictionaryCellList) = keys(cell_list.hashtable)
+@inline each_cell_index(cell_list::DictionaryCellList) = keys(cell_list.hashtable)
+
+# For this cell list, this is a `KeySet`, which has to be `collect`ed first to be
+# able to be used in a threaded loop.
+@inline function each_cell_index_threadable(cell_list::DictionaryCellList)
+    return collect(each_cell_index(cell_list))
+end
 
 @inline function Base.getindex(cell_list::DictionaryCellList, cell)
     (; hashtable, empty_vector) = cell_list
@@ -52,3 +73,14 @@ end
     # reuse the empty vector to avoid allocations.
     return get(hashtable, cell, empty_vector)
 end
+
+@inline function is_correct_cell(::DictionaryCellList, cell_coords, cell_index)
+    return cell_coords == cell_index
+end
+
+@inline index_type(::DictionaryCellList{NDIMS}) where {NDIMS} = NTuple{NDIMS, Int}
+
+function copy_cell_list(::DictionaryCellList{NDIMS}, search_radius,
+                        periodic_box) where {NDIMS}
+    return DictionaryCellList{NDIMS}()
+end

src/cell_lists/full_grid.jl

@@ -0,0 +1,118 @@
+"""
+    FullGridCellList(; min_corner, max_corner, search_radius = 0.0, periodicity = false)
+
+A simple cell list implementation where each (empty or non-empty) cell of a rectangular
+(axis-aligned) domain is assigned a list of points.
+This cell list only works when all points are inside the specified domain at all times.
+
+Only set `min_corner` and `max_corner` and use the default values for the other arguments
+to create an empty "template" cell list that can be used to create an empty "template"
+neighborhood search.
+See [`copy_neighborhood_search`](@ref) for more details.
+
+# Keywords
+- `min_corner`: Coordinates of the domain corner in negative coordinate directions.
+- `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.
+"""
+struct FullGridCellList{C, LI, MC} <: AbstractCellList
+    cells          :: C
+    linear_indices :: LI
+    min_cell       :: MC
+
+    function FullGridCellList{C, LI, MC}(cells, linear_indices, min_cell) where {C, LI, MC}
+        new{C, LI, MC}(cells, linear_indices, min_cell)
+    end
+end
+
+function FullGridCellList(; min_corner, max_corner, search_radius = 0.0,
+                          periodicity = false)
+    if search_radius < eps()
+        # Create an empty "template" cell list to be used with `copy_cell_list`
+        cells = nothing
+        linear_indices = nothing
+
+        # Misuse `min_cell` to store min and max corner for copying
+        min_cell = (min_corner, max_corner)
+    else
+        if periodicity
+            # Subtract `min_corner` because that's how the grid NHS works with periodicity
+            max_corner = max_corner .- min_corner
+            min_corner = min_corner .- min_corner
+        end
+
+        # Note that we don't shift everything so that the first cell starts at `min_corner`.
+        # The first cell is the cell containing `min_corner`, so we need to add one layer
+        # in order for `max_corner` to be inside a cell.
+        n_cells_per_dimension = ceil.(Int, (max_corner .- min_corner) ./ search_radius) .+ 1
+        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)]
+    end
+
+    return FullGridCellList{typeof(cells), typeof(linear_indices),
+                            typeof(min_cell)}(cells, linear_indices, min_cell)
+end
+
+function Base.empty!(cell_list::FullGridCellList)
+    Base.empty!.(cell_list.cells)
+
+    return cell_list
+end
+
+function Base.empty!(cell_list::FullGridCellList{Nothing})
+    # This is an empty "template" cell list to be used with `copy_cell_list`
+    throw(UndefRefError("`search_radius` is not defined for this cell list"))
+end
+
+function push_cell!(cell_list::FullGridCellList, cell, particle)
+    push!(cell_list[cell], particle)
+end
+
+function push_cell!(cell_list::FullGridCellList{Nothing}, cell, particle)
+    # This is an empty "template" cell list to be used with `copy_cell_list`
+    throw(UndefRefError("`search_radius` is not defined for this cell list"))
+end
+
+function deleteat_cell!(cell_list::FullGridCellList, cell, i)
+    deleteat!(cell_list[cell], i)
+end
+
+@inline each_cell_index(cell_list::FullGridCellList) = eachindex(cell_list.cells)
+
+function each_cell_index(cell_list::FullGridCellList{Nothing})
+    # This is an empty "template" cell list to be used with `copy_cell_list`
+    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
+
+    return cells[linear_indices[(cell .- min_cell .+ 1)...]]
+end
+
+@inline function Base.getindex(cell_list::FullGridCellList, i::Integer)
+    return cell_list.cells[i]
+end
+
+@inline function is_correct_cell(cell_list::FullGridCellList, cell_coords, cell_index)
+    (; linear_indices, min_cell) = cell_list
+
+    return linear_indices[(cell_coords .- min_cell .+ 1)...] == cell_index
+end
+
+@inline index_type(::FullGridCellList) = Int32
+
+function copy_cell_list(cell_list::FullGridCellList, search_radius, periodic_box)
+    # Misuse `min_cell` to store min and max corner for copying
+    min_corner, max_corner = cell_list.min_cell
+
+    return FullGridCellList(; min_corner, max_corner, search_radius,
+                            periodicity = !isnothing(periodic_box))
+end

src/neighborhood_search.jl

@@ -75,7 +75,7 @@ end
 """
     PeriodicBox(; min_corner, max_corner)
 
-Define a rectangular periodic domain.
+Define a rectangular (axis-aligned) periodic domain.
 
 # Keywords
 - `min_corner`: Coordinates of the domain corner in negative coordinate directions.

src/nhs_grid.jl

@@ -1,6 +1,8 @@
 @doc raw"""
     GridNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
-                                  periodic_box = nothing, threaded_update = true)
+                                  periodic_box = nothing,
+                                  cell_list = DictionaryCellList{NDIMS}(),
+                                  threaded_update = true)
 
 Simple grid-based neighborhood search with uniform search radius.
 The domain is divided into a regular grid.
@@ -32,6 +34,8 @@ since not sorting makes our implementation a lot faster (although less paralleli
                             with [`copy_neighborhood_search`](@ref).
 - `periodic_box = nothing`: In order to use a (rectangular) periodic domain, pass a
                             [`PeriodicBox`](@ref).
+- `cell_list`:              The cell list that maps a cell index to a list of points inside
+                            the cell. By default, a [`DictionaryCellList`](@ref) is used.
 - `threaded_update = true`: Can be used to deactivate thread parallelization in the
                             neighborhood search update. This can be one of the largest
                             sources of variations between simulations with different
@@ -47,23 +51,23 @@ since not sorting makes our implementation a lot faster (although less paralleli
   In: Computer Graphics Forum 30.1 (2011), pages 99–112.
   [doi: 10.1111/J.1467-8659.2010.01832.X](https://doi.org/10.1111/J.1467-8659.2010.01832.X)
 """
-struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, PB} <: AbstractNeighborhoodSearch
+struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, CB, PB} <: AbstractNeighborhoodSearch
     cell_list           :: CL
     search_radius       :: ELTYPE
     periodic_box        :: PB
     n_cells             :: NTuple{NDIMS, Int}    # Required to calculate periodic cell index
     cell_size           :: NTuple{NDIMS, ELTYPE} # Required to calculate cell index
-    cell_buffer         :: Array{NTuple{NDIMS, Int}, 2} # Multithreaded buffer for `update!`
+    cell_buffer         :: CB                    # Multithreaded buffer for `update!`
     cell_buffer_indices :: Vector{Int} # Store which entries of `cell_buffer` are initialized
     threaded_update     :: Bool
 
     function GridNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
                                            periodic_box = nothing,
+                                           cell_list = DictionaryCellList{NDIMS}(),
                                            threaded_update = true) where {NDIMS}
         ELTYPE = typeof(search_radius)
-        cell_list = DictionaryCellList{NDIMS}()
 
-        cell_buffer = Array{NTuple{NDIMS, Int}, 2}(undef, n_points, Threads.nthreads())
+        cell_buffer = Array{index_type(cell_list), 2}(undef, n_points, Threads.nthreads())
         cell_buffer_indices = zeros(Int, Threads.nthreads())
 
         if search_radius < eps() || isnothing(periodic_box)
@@ -86,7 +90,7 @@ struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, PB} <: AbstractNeighborhoodSear
             end
         end
 
-        new{NDIMS, ELTYPE, typeof(cell_list),
+        new{NDIMS, ELTYPE, typeof(cell_list), typeof(cell_buffer),
             typeof(periodic_box)}(cell_list, search_radius, periodic_box, n_cells,
                                   cell_size, cell_buffer, cell_buffer_indices,
                                   threaded_update)
@@ -156,13 +160,15 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
     for thread in 1:Threads.nthreads()
         # Only the entries `1:cell_buffer_indices[thread]` are initialized for `thread`.
         for i in 1:cell_buffer_indices[thread]
-            cell = cell_buffer[i, thread]
-            points = cell_list[cell]
+            cell_index = cell_buffer[i, thread]
+            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 cell_coords(coords_fun(points[i]),
-                                                  neighborhood_search) != cell)
+                                   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
@@ -174,7 +180,7 @@ function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
             end
 
             # Remove moved points from this cell
-            deleteat_cell!(cell_list, cell, moved_point_indices)
+            deleteat_cell!(cell_list, cell_index, moved_point_indices)
         end
     end
 
@@ -184,33 +190,38 @@ 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`
-    @threaded for cell in collect(eachcell(cell_list))
-        mark_changed_cell!(neighborhood_search, cell, coords_fun)
+    @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})
-    for cell in eachcell(cell_list)
-        mark_changed_cell!(neighborhood_search, cell, coords_fun)
+    for cell_index in each_cell_index(cell_list)
+        mark_changed_cell!(neighborhood_search, cell_index, coords_fun)
     end
 end
 
 # Use this function barrier and unpack inside to avoid passing closures to Polyester.jl
 # with `@batch` (`@threaded`).
 # Otherwise, `@threaded` does not work here with Julia ARM on macOS.
 # See https://github.com/JuliaSIMD/Polyester.jl/issues/88.
-@inline function mark_changed_cell!(neighborhood_search, cell, coords_fun)
+@inline function mark_changed_cell!(neighborhood_search, cell_index, coords_fun)
     (; cell_list, cell_buffer, cell_buffer_indices) = neighborhood_search
 
-    for point in cell_list[cell]
-        if cell_coords(coords_fun(point), neighborhood_search) != cell
+    for point in cell_list[cell_index]
+        cell = cell_coords(coords_fun(point), neighborhood_search)
+
+        # `cell` is a tuple, `cell_index` is the linear index used internally by the
+        # cell list to store cells inside `cell`.
+        # These can be identical (see `DictionaryCellList`).
+        if !is_correct_cell(cell_list, cell, cell_index)
             # Mark this cell and continue with the next one.
             #
             # `cell_buffer` is preallocated,
             # but only the entries 1:i are used for this thread.
             i = cell_buffer_indices[Threads.threadid()] += 1
-            cell_buffer[i, Threads.threadid()] = cell
+            cell_buffer[i, Threads.threadid()] = cell_index
             break
         end
     end
@@ -290,28 +301,12 @@ end
     return Tuple(floor_to_int.(offset_coords ./ cell_size))
 end
 
-# When points end up with coordinates so big that the cell coordinates
-# exceed the range of Int, then `floor(Int, i)` will fail with an InexactError.
-# In this case, we can just use typemax(Int), since we can assume that points
-# that far away will not interact with anything, anyway.
-# This usually indicates an instability, but we don't want the simulation to crash,
-# since adaptive time integration methods may detect the instability and reject the
-# time step.
-# If we threw an error here, we would prevent the time integration method from
-# retrying with a smaller time step, and we would thus crash perfectly fine simulations.
-@inline function floor_to_int(i)
-    if isnan(i) || i > typemax(Int)
-        return typemax(Int)
-    elseif i < typemin(Int)
-        return typemin(Int)
-    end
-
-    return floor(Int, i)
-end
-
 function copy_neighborhood_search(nhs::GridNeighborhoodSearch, search_radius, n_points;
                                   eachpoint = 1:n_points)
-    return GridNeighborhoodSearch{ndims(nhs)}(; search_radius, n_points,
-                                              periodic_box = nhs.periodic_box,
-                                              threaded_update = nhs.threaded_update)
+    (; periodic_box, threaded_update) = nhs
+
+    cell_list = copy_cell_list(nhs.cell_list, search_radius, periodic_box)
+
+    return GridNeighborhoodSearch{ndims(nhs)}(; search_radius, n_points, periodic_box,
+                                              cell_list, threaded_update)
 end

src/util.jl

@@ -3,6 +3,25 @@
     return SVector(ntuple(@inline(dim->A[dim, i]), NDIMS))
 end
 
+# When particles end up with coordinates so big that the cell coordinates
+# exceed the range of Int, then `floor(Int, i)` will fail with an InexactError.
+# In this case, we can just use typemax(Int), since we can assume that particles
+# that far away will not interact with anything, anyway.
+# This usually indicates an instability, but we don't want the simulation to crash,
+# since adaptive time integration methods may detect the instability and reject the
+# time step.
+# If we threw an error here, we would prevent the time integration method from
+# retrying with a smaller time step, and we would thus crash perfectly fine simulations.
+@inline function floor_to_int(i)
+    if isnan(i) || i > typemax(Int)
+        return typemax(Int)
+    elseif i < typemin(Int)
+        return typemin(Int)
+    end
+
+    return floor(Int, i)
+end
+
 """
     @threaded for ... end