Implement dynamic vector of vectors data structure

src/PointNeighbors.jl

@@ -7,6 +7,7 @@ using Polyester: @batch
 @reexport using StaticArrays: SVector
 
 include("util.jl")
+include("vector_of_vectors.jl")
 include("neighborhood_search.jl")
 include("nhs_trivial.jl")
 include("cell_lists/cell_lists.jl")

src/nhs_grid.jl

@@ -51,24 +51,25 @@ 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, 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         :: CB                    # Multithreaded buffer for `update!`
-    cell_buffer_indices :: Vector{Int} # Store which entries of `cell_buffer` are initialized
-    threaded_update     :: Bool
+struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, PB, UB} <: 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
+    update_buffer   :: UB # Multithreaded buffer for `update!`
+    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_buffer = Array{index_type(cell_list), 2}(undef, n_points, Threads.nthreads())
-        cell_buffer_indices = zeros(Int, Threads.nthreads())
+        # Create update buffer and initialize it with empty vectors
+        update_buffer = DynamicVectorOfVectors{index_type(cell_list)}(max_outer_length = Threads.nthreads(),
+                                                                      max_inner_length = n_points)
+        push!(update_buffer, (NTuple{NDIMS, Int}[] for _ in 1:Threads.nthreads())...)
 
         if search_radius < eps() || isnothing(periodic_box)
             # No periodicity
@@ -90,37 +91,28 @@ struct GridNeighborhoodSearch{NDIMS, ELTYPE, CL, CB, PB} <: AbstractNeighborhood
             end
         end
 
-        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)
+        new{NDIMS, ELTYPE, typeof(cell_list), typeof(periodic_box),
+            typeof(update_buffer)}(cell_list, search_radius, periodic_box, n_cells,
+                                   cell_size, update_buffer, threaded_update)
     end
 end
 
 @inline Base.ndims(::GridNeighborhoodSearch{NDIMS}) where {NDIMS} = NDIMS
 
-@inline function npoints(neighborhood_search::GridNeighborhoodSearch)
-    return size(neighborhood_search.cell_buffer, 1)
-end
-
 function initialize!(neighborhood_search::GridNeighborhoodSearch,
                      x::AbstractMatrix, y::AbstractMatrix)
     initialize_grid!(neighborhood_search, y)
 end
 
-function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{NDIMS},
-                          y::AbstractMatrix) where {NDIMS}
-    initialize_grid!(neighborhood_search, i -> extract_svector(y, Val(NDIMS), i))
-end
-
-function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
+function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::AbstractMatrix)
     (; cell_list) = neighborhood_search
 
     empty!(cell_list)
 
-    for point in 1:npoints(neighborhood_search)
+    for point in axes(y, 2)
         # Get cell index of the point's cell
-        cell = cell_coords(coords_fun(point), neighborhood_search)
+        point_coords = 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)
@@ -147,20 +139,19 @@ end
 
 # Modify the existing hash table by moving points into their new cells
 function update_grid!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
-    (; cell_list, cell_buffer, cell_buffer_indices, threaded_update) = neighborhood_search
+    (; cell_list, update_buffer, threaded_update) = neighborhood_search
 
-    # Reset `cell_buffer` by moving all pointers to the beginning
-    cell_buffer_indices .= 0
+    # Empty each thread's list
+    for i in eachindex(update_buffer)
+        emptyat!(update_buffer, i)
+    end
 
     # Find all cells containing points that now belong to another cell
-    mark_changed_cell!(neighborhood_search, cell_list, coords_fun,
-                       Val(threaded_update))
-
-    # Iterate over all marked cells and move the points into their new cells.
-    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_index = cell_buffer[i, thread]
+    mark_changed_cell!(neighborhood_search, cell_list, coords_fun, Val(threaded_update))
+
+    # Iterate over all marked cells and move the points into their new cells
+    for j in eachindex(update_buffer)
+        for cell_index in update_buffer[j]
             points = cell_list[cell_index]
 
             # Find all points whose coordinates do not match this cell
@@ -207,7 +198,7 @@ end
 # 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_index, coords_fun)
-    (; cell_list, cell_buffer, cell_buffer_indices) = neighborhood_search
+    (; cell_list, update_buffer) = neighborhood_search
 
     for point in cell_list[cell_index]
         cell = cell_coords(coords_fun(point), neighborhood_search)
@@ -216,12 +207,8 @@ end
         # 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_index
+            # Mark this cell and continue with the next one
+            pushat!(update_buffer, Threads.threadid(), cell_index)
             break
         end
     end

src/vector_of_vectors.jl

@@ -0,0 +1,96 @@
+# Data structure that behaves like a `Vector{Vector}`, but uses a contiguous memory layout.
+# Similar to `VectorOfVectors` of ArraysOfArrays.jl, but allows to resize the inner vectors.
+struct DynamicVectorOfVectors{T, ARRAY2D, ARRAY1D} <: AbstractVector{Array{T, 1}}
+    backend::ARRAY2D # Array{T, 2}, where each column represents a vector
+    length_::Base.RefValue{Int32} # Number of vectors
+    lengths::ARRAY1D # Array{Int32, 1} storing the lengths of the vectors
+end
+
+function DynamicVectorOfVectors{T}(; max_outer_length, max_inner_length) where {T}
+    backend = Array{T, 2}(undef, max_inner_length, max_outer_length)
+    length_ = Ref(zero(Int32))
+    lengths = zeros(Int32, max_outer_length)
+
+    return DynamicVectorOfVectors{T, typeof(backend), typeof(lengths)}(backend, length_,
+                                                                       lengths)
+end
+
+@inline Base.size(vov::DynamicVectorOfVectors) = (vov.length_[],)
+
+@inline function Base.getindex(vov::DynamicVectorOfVectors, i)
+    (; backend, lengths) = vov
+
+    @boundscheck checkbounds(vov, i)
+
+    return view(backend, 1:lengths[i], i)
+end
+
+@inline function Base.push!(vov::DynamicVectorOfVectors, vector::AbstractVector)
+    (; backend, length_, lengths) = vov
+
+    # This data structure only supports one-based indexing
+    Base.require_one_based_indexing(vector)
+
+    # Activate a new column of `backend`
+    j = length_[] += 1
+    lengths[j] = length(vector)
+
+    # Fill the new column
+    for i in eachindex(vector)
+        backend[i, j] = vector[i]
+    end
+
+    return vov
+end
+
+@inline function Base.push!(vov::DynamicVectorOfVectors, vector::AbstractVector, vectors...)
+    push!(vov, vector)
+    push!(vov, vectors...)
+end
+
+# `push!(vov[i], value)`
+@inline function pushat!(vov::DynamicVectorOfVectors, i, value)
+    (; backend, lengths) = vov
+
+    @boundscheck checkbounds(vov, i)
+
+    # Activate new entry in column `i`
+    backend[lengths[i] += 1, i] = value
+
+    return vov
+end
+
+# `deleteat!(vov[i], j)`
+@inline function deleteatat!(vov::DynamicVectorOfVectors, i, j)
+    (; backend, lengths) = vov
+
+    # Outer bounds check
+    @boundscheck checkbounds(vov, i)
+    # Inner bounds check
+    @boundscheck checkbounds(1:lengths[i], j)
+
+    # Replace value to delete by the last value in this column
+    last_value = backend[lengths[i], i]
+    backend[j, i] = last_value
+
+    # Remove the last value in this column
+    lengths[i] -= 1
+
+    return vov
+end
+
+@inline function Base.empty!(vov::DynamicVectorOfVectors)
+    # Move all pointers to the beginning
+    vov.lengths .= zero(Int32)
+    vov.length_[] = zero(Int32)
+
+    return vov
+end
+
+# `empty!(vov[i])`
+@inline function emptyat!(vov::DynamicVectorOfVectors, i)
+    # Move length pointer to the beginning
+    vov.lengths[i] = zero(Int32)
+
+    return vov
+end

test/nhs_grid.jl

@@ -91,9 +91,7 @@
 
         # Create neighborhood search
         nhs1 = GridNeighborhoodSearch{3}(; search_radius, n_points)
-
-        coords_fun(i) = coordinates1[:, i]
-        initialize_grid!(nhs1, coords_fun)
+        initialize_grid!(nhs1, coordinates1)
 
         # Get each neighbor for `point_position1`
         neighbors1 = sort(collect(PointNeighbors.eachneighbor(point_position1, nhs1)))

test/test_util.jl

@@ -1,6 +1,6 @@
 # All `using` calls are in this file, so that one can run any test file
 # after running only this file.
-using Test: @test, @testset
+using Test: @test, @testset, @test_throws
 using PointNeighbors
 
 """

test/unittest.jl

@@ -1,6 +1,7 @@
 # Separate file that can be executed to only run unit tests.
 # Include `test_util.jl` first.
 @testset verbose=true "Unit Tests" begin
+    include("vector_of_vectors.jl")
     include("nhs_trivial.jl")
     include("nhs_grid.jl")
     include("neighborhood_search.jl")

test/vector_of_vectors.jl

@@ -0,0 +1,81 @@
+@testset verbose=true "`DynamicVectorOfVectors`" begin
+    # Test different types by defining a function to convert to this type
+    types = [Int32, Float64, i -> (i, i)]
+
+    @testset verbose=true "Eltype $(eltype(type(1)))" for type in types
+        ELTYPE = typeof(type(1))
+        vov_ref = Vector{Vector{ELTYPE}}()
+        vov = PointNeighbors.DynamicVectorOfVectors{ELTYPE}(max_outer_length = 20,
+                                                            max_inner_length = 10)
+
+        # Test internal size
+        @test size(vov.backend) == (10, 20)
+
+        function verify(vov, vov_ref)
+            @test length(vov) == length(vov_ref)
+            @test eachindex(vov) == eachindex(vov_ref)
+            @test axes(vov) == axes(vov_ref)
+
+            @test_throws BoundsError vov[0]
+            @test_throws BoundsError vov[length(vov) + 1]
+
+            for i in eachindex(vov_ref)
+                @test vov[i] == vov_ref[i]
+            end
+        end
+
+        # Initial check
+        verify(vov, vov_ref)
+
+        # First `push!`
+        push!(vov_ref, type.([1, 2, 3]))
+        push!(vov, type.([1, 2, 3]))
+
+        verify(vov, vov_ref)
+
+        # `push!` multiple items
+        push!(vov_ref, type.([4]), type.([5, 6, 7, 8]))
+        push!(vov, type.([4]), type.([5, 6, 7, 8]))
+
+        verify(vov, vov_ref)
+
+        # `push!` to an inner vector
+        push!(vov_ref[1], type(12))
+        PointNeighbors.pushat!(vov, 1, type(12))
+
+        verify(vov, vov_ref)
+
+        # Delete entry of inner vector. Note that this changes the order of the elements.
+        deleteat!(vov_ref[3], 2)
+        PointNeighbors.deleteatat!(vov, 3, 2)
+
+        @test vov_ref[3] == type.([5, 7, 8])
+        @test vov[3] == type.([5, 8, 7])
+
+        # Delete second to last entry
+        deleteat!(vov_ref[3], 2)
+        PointNeighbors.deleteatat!(vov, 3, 2)
+
+        @test vov_ref[3] == type.([5, 8])
+        @test vov[3] == type.([5, 7])
+
+        # Delete last entry
+        deleteat!(vov_ref[3], 2)
+        PointNeighbors.deleteatat!(vov, 3, 2)
+
+        # Now they are identical again
+        verify(vov, vov_ref)
+
+        # Delete the remaining entry of this vector
+        deleteat!(vov_ref[3], 1)
+        PointNeighbors.deleteatat!(vov, 3, 1)
+
+        verify(vov, vov_ref)
+
+        # `empty!`
+        empty!(vov_ref)
+        empty!(vov)
+
+        verify(vov, vov_ref)
+    end
+end