diff --git a/Project.toml b/Project.toml index a327029..ff05601 100644 --- a/Project.toml +++ b/Project.toml @@ -4,12 +4,14 @@ authors = ["Erik Faulhaber "] version = "0.2.4-dev" [deps] +ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [compat] +ArraysOfArrays = "0.6" LinearAlgebra = "1" Polyester = "0.7.5" Reexport = "1" diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl index 3db27cf..4bf29e9 100644 --- a/src/PointNeighbors.jl +++ b/src/PointNeighbors.jl @@ -2,6 +2,7 @@ module PointNeighbors using Reexport: @reexport +using ArraysOfArrays: VectorOfVectors using LinearAlgebra: dot using Polyester: @batch @reexport using StaticArrays: SVector @@ -10,9 +11,10 @@ include("util.jl") include("neighborhood_search.jl") include("nhs_trivial.jl") include("nhs_grid.jl") +include("nhs_neighbor_lists.jl") export for_particle_neighbor, foreach_neighbor -export TrivialNeighborhoodSearch, GridNeighborhoodSearch +export TrivialNeighborhoodSearch, GridNeighborhoodSearch, NeighborListsNeighborhoodSearch export initialize!, update!, initialize_grid!, update_grid! end # module PointNeighbors diff --git a/src/nhs_neighbor_lists.jl b/src/nhs_neighbor_lists.jl new file mode 100644 index 0000000..e3bfda5 --- /dev/null +++ b/src/nhs_neighbor_lists.jl @@ -0,0 +1,134 @@ +@doc raw""" + NeighborListsNeighborhoodSearch{NDIMS}(search_radius, n_particles; + periodic_box_min_corner = nothing, + periodic_box_max_corner = nothing, + backend = VectorOfVectors{Int32}) + +Neighborhood search with precomputed neighbor lists. A list of all neighbors is computed +for each particle during initialization and update. +This neighborhood search maximizes the performance of neighbor loops at the cost of a much +slower [`update!`](@ref). + +A [`GridNeighborhoodSearch`](@ref) is used internally to compute the neighbor lists during +initialization and update. + +# Arguments +- `NDIMS`: Number of dimensions. +- `search_radius`: The uniform search radius. +- `n_particles`: Total number of particles. + +# Keywords +- `periodic_box_min_corner`: In order to use a (rectangular) periodic domain, pass the + coordinates of the domain corner in negative coordinate + directions. +- `periodic_box_max_corner`: In order to use a (rectangular) periodic domain, pass the + coordinates of the domain corner in positive coordinate + directions. +- `backend=VectorOfVectors{Int32}`: Data structure to store the neighbor lists. The default + `VectorOfVectors` is a data structure from + [ArraysOfArrays.jl](https://github.com/JuliaArrays/ArraysOfArrays.jl), + which behaves like a `Vector` of `Vector`s, but uses + a single `Vector` for a contiguous memory layout. + Use `backend=Vector{Vector{Int32}}` to benchmark + the benefits of this representation. +""" +struct NeighborListsNeighborhoodSearch{NDIMS, NHS, NL, PB} + neighborhood_search :: NHS + neighbor_lists :: NL + periodic_box :: PB + + function NeighborListsNeighborhoodSearch{NDIMS}(search_radius, n_particles; + periodic_box_min_corner = nothing, + periodic_box_max_corner = nothing, + backend = VectorOfVectors{Int32}) where { + NDIMS + } + nhs = GridNeighborhoodSearch{NDIMS}(search_radius, n_particles, + periodic_box_min_corner = periodic_box_min_corner, + periodic_box_max_corner = periodic_box_max_corner) + + neighbor_lists = backend() + + new{NDIMS, typeof(nhs), typeof(neighbor_lists), + typeof(nhs.periodic_box)}(nhs, neighbor_lists, nhs.periodic_box) + end +end + +@inline function Base.ndims(::NeighborListsNeighborhoodSearch{NDIMS}) where {NDIMS} + return NDIMS +end + +function initialize!(search::NeighborListsNeighborhoodSearch, + x::AbstractMatrix, y::AbstractMatrix) + (; neighborhood_search, neighbor_lists) = search + + # Initialize grid NHS + initialize!(neighborhood_search, x, y) + + initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y) +end + +function update!(search::NeighborListsNeighborhoodSearch, + x::AbstractMatrix, y::AbstractMatrix; + particles_moving = (true, true)) + (; neighborhood_search, neighbor_lists) = search + + # Update grid NHS + update!(neighborhood_search, x, y, particles_moving = particles_moving) + + # Skip update if both point sets are static + if any(particles_moving) + initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y) + end +end + +function initialize_neighbor_lists!(neighbor_lists::Vector{<:Vector}, neighborhood_search, + x, y) + # Initialize neighbor lists + empty!(neighbor_lists) + resize!(neighbor_lists, size(x, 2)) + for i in eachindex(neighbor_lists) + neighbor_lists[i] = Int[] + end + + # Fill neighbor lists + for_particle_neighbor(x, y, neighborhood_search) do particle, neighbor, _, _ + push!(neighbor_lists[particle], neighbor) + end +end + +function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y) + neighbor_lists_ = Vector{Vector{Int32}}() + initialize_neighbor_lists!(neighbor_lists_, neighborhood_search, x, y) + + empty!(neighbor_lists) + for i in eachindex(neighbor_lists_) + push!(neighbor_lists, neighbor_lists_[i]) + end +end + +@inline function foreach_neighbor(f, system_coords, neighbor_system_coords, + neighborhood_search::NeighborListsNeighborhoodSearch, + particle; search_radius = nothing) + (; periodic_box, neighbor_lists) = neighborhood_search + (; search_radius) = neighborhood_search.neighborhood_search + + particle_coords = extract_svector(system_coords, Val(ndims(neighborhood_search)), + particle) + for neighbor in neighbor_lists[particle] + neighbor_coords = extract_svector(neighbor_system_coords, + Val(ndims(neighborhood_search)), neighbor) + + pos_diff = particle_coords - neighbor_coords + distance2 = dot(pos_diff, pos_diff) + + pos_diff, distance2 = compute_periodic_distance(pos_diff, distance2, search_radius, + periodic_box) + + distance = sqrt(distance2) + + # Inline to avoid loss of performance + # compared to not using `for_particle_neighbor`. + @inline f(particle, neighbor, pos_diff, distance) + end +end diff --git a/test/neighborhood_search.jl b/test/neighborhood_search.jl index 2045d1d..2483f78 100644 --- a/test/neighborhood_search.jl +++ b/test/neighborhood_search.jl @@ -45,10 +45,19 @@ GridNeighborhoodSearch{NDIMS}(search_radius, n_particles, periodic_box_min_corner = periodic_boxes[i][1], periodic_box_max_corner = periodic_boxes[i][2]), + NeighborListsNeighborhoodSearch{NDIMS}(search_radius, n_particles, + periodic_box_min_corner = periodic_boxes[i][1], + periodic_box_max_corner = periodic_boxes[i][2]), + NeighborListsNeighborhoodSearch{NDIMS}(search_radius, n_particles, + periodic_box_min_corner = periodic_boxes[i][1], + periodic_box_max_corner = periodic_boxes[i][2], + backend = Vector{Vector{Int32}}), ] neighborhood_searches_names = [ "`TrivialNeighborhoodSearch`", "`GridNeighborhoodSearch`", + "`NeighborListsNeighborhoodSearch`", + "`NeighborListsNeighborhoodSearch` with `Vector{Vector}` backend", ] # Run this for every neighborhood search @@ -86,11 +95,14 @@ ] seeds = [1, 2] - @testset verbose=true "$(length(cloud_size))D with $(prod(cloud_size)) Particles ($(seed == 1 ? "`initialize!`" : "`update!`"))" for cloud_size in cloud_sizes, - seed in seeds + name(size, seed) = "$(length(size))D with $(prod(size)) Particles " * + "($(seed == 1 ? "`initialize!`" : "`update!`"))" + @testset verbose=true "$(name(cloud_size, seed)))" for cloud_size in cloud_sizes, + seed in seeds coords = point_cloud(cloud_size, seed = seed) NDIMS = length(cloud_size) + n_particles = size(coords, 2) search_radius = 2.5 # Use different coordinates for `initialize!` and then `update!` with the @@ -110,11 +122,16 @@ end neighborhood_searches = [ - GridNeighborhoodSearch{NDIMS}(search_radius, size(coords, 2)), + GridNeighborhoodSearch{NDIMS}(search_radius, n_particles), + NeighborListsNeighborhoodSearch{NDIMS}(search_radius, n_particles), + NeighborListsNeighborhoodSearch{NDIMS}(search_radius, n_particles, + backend = Vector{Vector{Int32}}), ] neighborhood_searches_names = [ "`GridNeighborhoodSearch`", + "`NeighborListsNeighborhoodSearch`", + "`NeighborListsNeighborhoodSearch` with `Vector{Vector}` backend", ] @testset "$(neighborhood_searches_names[i])" for i in eachindex(neighborhood_searches_names)