Add automatic padding for the FullGridCellList

Project.toml

@@ -1,7 +1,7 @@
 name = "PointNeighbors"
 uuid = "1c4d5385-0a27-49de-8e2c-43b175c8985c"
 authors = ["Erik Faulhaber <[email protected]>"]
-version = "0.4.7"
+version = "0.4.8"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/cell_lists/full_grid.jl

@@ -42,24 +42,24 @@ end
 
 supported_update_strategies(::FullGridCellList) = (SemiParallelUpdate, SerialUpdate)
 
-function FullGridCellList(; min_corner, max_corner, search_radius = 0.0,
+function FullGridCellList(; min_corner, max_corner,
+                          search_radius = zero(eltype(min_corner)),
                           backend = DynamicVectorOfVectors{Int32},
                           max_points_per_cell = 100)
-    # Pad domain to avoid 0 in cell indices due to rounding errors.
+    # Add one layer in each direction to make sure neighbor cells exist.
+    # Also pad domain a little more to avoid 0 in cell indices due to rounding errors.
     # We can't just use `eps()`, as one might use lower precision types.
     # This padding is safe, and will give us one more layer of cells in the worst case.
-    min_corner = SVector(Tuple(min_corner .- 1.0f-3 * search_radius))
-    max_corner = SVector(Tuple(max_corner .+ 1.0f-3 * search_radius))
+    # `1001 // 1000` is 1.001 without forcing a float type.
+    min_corner = SVector(Tuple(min_corner .- 1001 // 1000 * search_radius))
+    max_corner = SVector(Tuple(max_corner .+ 1001 // 1000 * search_radius))
 
     if search_radius < eps()
         # Create an empty "template" cell list to be used with `copy_cell_list`
         cells = construct_backend(backend, 0, max_points_per_cell)
         linear_indices = nothing
     else
-        # 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
+        n_cells_per_dimension = ceil.(Int, (max_corner .- min_corner) ./ search_radius)
         linear_indices = LinearIndices(Tuple(n_cells_per_dimension))
 
         cells = construct_backend(backend, n_cells_per_dimension, max_points_per_cell)
@@ -96,11 +96,29 @@ end
 
     # Subtract `min_corner` to offset coordinates so that the min corner of the grid
     # corresponds to the (1, 1, 1) cell.
-    # Note that we use `min_corner == periodic_box.min_corner`, so we don't have to handle
-    # periodic boxes differently, as they also use 1-based indexing.
     return Tuple(floor_to_int.((coords .- min_corner) ./ cell_size)) .+ 1
 end
 
+@inline function cell_coords(coords, periodic_box::PeriodicBox, cell_list::FullGridCellList,
+                             cell_size)
+    # Subtract `periodic_box.min_corner` to offset coordinates so that the min corner
+    # of the grid corresponds to the (0, 0, 0) cell.
+    offset_coords = periodic_coords(coords, periodic_box) .- periodic_box.min_corner
+
+    # Add 2, so that the min corner will be the (2, 2, 2)-cell.
+    # With this, we still have one padding layer in each direction around the periodic box,
+    # just like without using a periodic box.
+    # This is not needed for finding neighbor cells, but to make the bounds check
+    # work the same way as without a periodic box.
+    return Tuple(floor_to_int.(offset_coords ./ cell_size)) .+ 2
+end
+
+@inline function periodic_cell_index(cell_index, ::PeriodicBox, n_cells,
+                                     cell_list::FullGridCellList)
+    # 2-based modulo to match the indexing of the periodic box explained above.
+    return mod.(cell_index .- 2, n_cells) .+ 2
+end
+
 function Base.empty!(cell_list::FullGridCellList)
     (; cells) = cell_list
 
@@ -120,8 +138,10 @@ end
 function push_cell!(cell_list::FullGridCellList, cell, particle)
     (; cells) = cell_list
 
+    @boundscheck check_cell_bounds(cell_list, cell)
+
     # `push!(cell_list[cell], particle)`, but for all backends
-    pushat!(cells, cell_index(cell_list, cell), particle)
+    @inbounds pushat!(cells, cell_index(cell_list, cell), particle)
 
     return cell_list
 end
@@ -134,17 +154,21 @@ end
 @inline function push_cell_atomic!(cell_list::FullGridCellList, cell, particle)
     (; cells) = cell_list
 
+    @boundscheck check_cell_bounds(cell_list, cell)
+
     # `push!(cell_list[cell], particle)`, but for all backends.
     # The atomic version of `pushat!` uses atomics to avoid race conditions when `pushat!`
     # is used in a parallel loop.
-    pushat_atomic!(cells, cell_index(cell_list, cell), particle)
+    @inbounds pushat_atomic!(cells, cell_index(cell_list, cell), particle)
 
     return cell_list
 end
 
 function deleteat_cell!(cell_list::FullGridCellList, cell, i)
     (; cells) = cell_list
 
+    @boundscheck check_cell_bounds(cell_list, cell)
+
     # `deleteat!(cell_list[cell], i)`, but for all backends
     deleteatat!(cells, cell_index(cell_list, cell), i)
 end
@@ -170,8 +194,10 @@ end
     return cells[cell_index(cell_list, cell)]
 end
 
-@inline function is_correct_cell(cell_list::FullGridCellList, cell_coords, cell_index_)
-    return cell_index(cell_list, cell_coords) == cell_index_
+@inline function is_correct_cell(cell_list::FullGridCellList, cell, cell_index_)
+    @boundscheck check_cell_bounds(cell_list, cell)
+
+    return cell_index(cell_list, cell) == cell_index_
 end
 
 @inline index_type(::FullGridCellList) = Int32
@@ -188,7 +214,23 @@ function max_points_per_cell(cells::DynamicVectorOfVectors)
     return size(cells.backend, 1)
 end
 
-# Fallback when backend is a `Vector{Vector{T}}`
+# Fallback when backend is a `Vector{Vector{T}}`. Only used for copying the cell list.
 function max_points_per_cell(cells)
     return 100
 end
+
+@inline function check_cell_bounds(cell_list::FullGridCellList, cell::Tuple)
+    (; linear_indices) = cell_list
+
+    # Make sure that points are not added to the outer padding layer, which is needed
+    # to ensure that neighboring cells in all directions of all non-empty cells exist.
+    if !all(cell[i] in 2:(size(linear_indices, i) - 1) for i in eachindex(cell))
+        error("particle coordinates are NaN or outside the domain bounds of the cell list")
+    end
+end
+
+@inline function check_cell_bounds(cell_list::FullGridCellList, cell::Integer)
+    (; cells) = cell_list
+
+    checkbounds(cells, cell)
+end

src/nhs_grid.jl

@@ -185,7 +185,7 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::Abstra
     return neighborhood_search
 end
 
-# WARNING! Undocumented, experimental feature:
+# WARNING! Experimental feature:
 # By default, determine the parallelization backend from the type of `x`.
 # Optionally, pass a `KernelAbstractions.Backend` to run the KernelAbstractions.jl code
 # on this backend. This can be useful to run GPU kernels on the CPU by passing
@@ -418,16 +418,16 @@ end
 end
 
 @inline function periodic_cell_index(cell_index, neighborhood_search)
-    (; n_cells, periodic_box) = neighborhood_search
+    (; n_cells, periodic_box, cell_list) = neighborhood_search
 
-    periodic_cell_index(cell_index, periodic_box, n_cells)
+    periodic_cell_index(cell_index, periodic_box, n_cells, cell_list)
 end
 
-@inline periodic_cell_index(cell_index, ::Nothing, n_cells) = cell_index
+@inline periodic_cell_index(cell_index, ::Nothing, n_cells, cell_list) = cell_index
 
-@inline function periodic_cell_index(cell_index, ::PeriodicBox, n_cells)
+@inline function periodic_cell_index(cell_index, ::PeriodicBox, n_cells, cell_list)
     # 1-based modulo
-    return rem.(cell_index .- 1, n_cells, RoundDown) .+ 1
+    return mod1.(cell_index, n_cells)
 end
 
 @inline function cell_coords(coords, neighborhood_search)

src/vector_of_vectors.jl

@@ -59,10 +59,16 @@ end
 @inline function pushat!(vov::DynamicVectorOfVectors, i, value)
     (; backend, lengths) = vov
 
+    # Outer bounds check
     @boundscheck checkbounds(vov, i)
 
+    lengths[i] += 1
+
+    # Inner bounds check
+    @boundscheck check_list_bounds(vov, i)
+
     # Activate new entry in column `i`
-    backend[lengths[i] += 1, i] = value
+    backend[lengths[i], i] = value
 
     return vov
 end
@@ -76,20 +82,33 @@ end
 @inline function pushat_atomic!(vov::DynamicVectorOfVectors, i, value)
     (; backend, lengths) = vov
 
+    # Outer bounds check
     @boundscheck checkbounds(vov, i)
 
     # Increment the column length with an atomic add to avoid race conditions.
     # Store the new value since it might be changed immediately afterwards by another
     # thread.
-    new_length = Atomix.@atomic lengths[i] += 1
+    new_length = @inbounds Atomix.@atomic lengths[i] += 1
+
+    # Inner bounds check
+    @boundscheck check_list_bounds(vov, i)
 
     # We can write here without race conditions, since the atomic add guarantees
     # that `new_length` is different for each thread.
-    backend[new_length, i] = value
+    @inbounds backend[new_length, i] = value
 
     return vov
 end
 
+@inline function check_list_bounds(vov::DynamicVectorOfVectors, i)
+    (; backend, lengths) = vov
+
+    if lengths[i] > size(backend, 1)
+        Atomix.@atomic lengths[i] -= 1
+        error("cell list is full. Use a larger `max_points_per_cell`.")
+    end
+end
+
 # `deleteat!(vov[i], j)`
 @inline function deleteatat!(vov::DynamicVectorOfVectors, i, j)
     (; backend, lengths) = vov
@@ -101,10 +120,10 @@ end
 
     # Replace value to delete by the last value in this column
     last_value = backend[lengths[i], i]
-    backend[j, i] = last_value
+    @inbounds backend[j, i] = last_value
 
     # Remove the last value in this column
-    lengths[i] -= 1
+    @inbounds lengths[i] -= 1
 
     return vov
 end

test/cell_lists/full_grid.jl

@@ -0,0 +1,52 @@
+@testset "`FullGridCellList`" verbose=true begin
+    # Test that `update!` throws an error when a particle is outside the bounding box
+    @testset "`update!` bounds check" begin
+        @testset "$(N)D" for N in 1:3
+            min_corner = fill(0.0, N)
+            max_corner = fill(10.0, N)
+            search_radius = 1.0
+
+            cell_list = FullGridCellList(; search_radius, min_corner, max_corner)
+
+            # Introduce the same rounding errors for this to pass
+            @test cell_list.min_corner == fill(-1.001, N)
+            @test cell_list.max_corner == fill(10.0 + 1.001, N)
+
+            nhs = GridNeighborhoodSearch{N}(; cell_list, search_radius)
+            y = rand(N, 10)
+            error_string = "particle coordinates are NaN or outside the domain bounds of the cell list"
+            error = ErrorException(error_string)
+
+            y[1, 7] = NaN
+            @test_throws error initialize!(nhs, y, y)
+
+            y[1, 7] = min_corner[1] - 0.01
+            @test_throws error initialize!(nhs, y, y)
+
+            # A bit more than max corner might still be inside the grid,
+            # but one search radius more is always outside.
+            # Also accounting for 0.001 extra padding (see above).
+            y[1, 7] = max_corner[1] + 1.01
+            @test_throws error initialize!(nhs, y, y)
+
+            y[1, 7] = 0.0
+            @test_nowarn_mod initialize!(nhs, y, y)
+            @test_nowarn_mod update!(nhs, y, y)
+
+            y[1, 7] = 10.0
+            @test_nowarn_mod update!(nhs, y, y)
+
+            y[1, 7] = NaN
+            @test_throws error update!(nhs, y, y)
+
+            # A bit more than max corner might still be inside the grid,
+            # but one search radius more is always outside.
+            # Also accounting for 0.001 extra padding (see above).
+            y[1, 7] = max_corner[1] + 1.01
+            @test_throws error update!(nhs, y, y)
+
+            y[1, 7] = min_corner[1] - 0.01
+            @test_throws error update!(nhs, y, y)
+        end
+    end
+end

test/neighborhood_search.jl

@@ -39,7 +39,7 @@
             search_radius = 0.1
 
             min_corner = periodic_boxes[i].min_corner
-            max_corner = max_corner = periodic_boxes[i].max_corner
+            max_corner = periodic_boxes[i].max_corner
 
             neighborhood_searches = [
                 TrivialNeighborhoodSearch{NDIMS}(; search_radius, eachpoint = 1:n_points,

test/nhs_grid.jl

@@ -75,7 +75,7 @@
         @test PointNeighbors.cell_coords(coords1, nothing, cell_list, (1.0, 1.0)) ==
               (typemax(Int), typemin(Int)) .+ 1
         @test PointNeighbors.cell_coords(coords2, nothing, cell_list, (1.0, 1.0)) ==
-              (typemax(Int), 0) .+ 1
+              (typemax(Int), 1) .+ 1
         @test PointNeighbors.cell_coords(coords3, nothing, cell_list, (1.0, 1.0)) ==
               (typemax(Int), typemin(Int)) .+ 1
     end

test/unittest.jl

@@ -5,4 +5,5 @@
     include("nhs_trivial.jl")
     include("nhs_grid.jl")
     include("neighborhood_search.jl")
+    include("cell_lists/full_grid.jl")
 end;

test/vector_of_vectors.jl

@@ -6,10 +6,10 @@
         ELTYPE = typeof(type(1))
         vov_ref = Vector{Vector{ELTYPE}}()
         vov = PointNeighbors.DynamicVectorOfVectors{ELTYPE}(max_outer_length = 20,
-                                                            max_inner_length = 10)
+                                                            max_inner_length = 4)
 
         # Test internal size
-        @test size(vov.backend) == (10, 20)
+        @test size(vov.backend) == (4, 20)
 
         function verify(vov, vov_ref)
             @test length(vov) == length(vov_ref)
@@ -43,6 +43,10 @@
         push!(vov_ref[1], type(12))
         PointNeighbors.pushat!(vov, 1, type(12))
 
+        # `push!` overflow
+        error_ = ErrorException("cell list is full. Use a larger `max_points_per_cell`.")
+        @test_throws error_ PointNeighbors.pushat!(vov, 1, type(13))
+
         verify(vov, vov_ref)
 
         # Delete entry of inner vector. Note that this changes the order of the elements.