Skip to content

Commit 591827b

Browse files
efaulhaberefaulhaber
authored
Implement neighborhood search based on static neighbor lists (#9)
* Implement NHS based on static neighbor lists * Reformat code * Export `NeighborListsNeighborhoodSearch` * Add `NeighborListsNeighborhoodSearch` to tests * Rename variable * Fix `foreach_neighbor` * Improve comments * Add docs for `NeighborListsNeighborhoodSearch` * Add neighbor lists NHS to tests * Rename neighborhood search * Reformat
1 parent e503620 commit 591827b

File tree

3 files changed

+130
-5
lines changed

3 files changed

+130
-5
lines changed

‎src/PointNeighbors.jl‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,12 @@ using Polyester: @batch
99
include("util.jl")
1010
include("neighborhood_search.jl")
1111
include("nhs_trivial.jl")
12-
include("nhs_grid.jl")
1312
include("cell_lists/cell_lists.jl")
13+
include("nhs_grid.jl")
14+
include("nhs_neighbor_lists.jl")
1415

1516
export for_particle_neighbor, foreach_neighbor
16-
export TrivialNeighborhoodSearch, GridNeighborhoodSearch
17+
export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch
1718
export initialize!, update!, initialize_grid!, update_grid!
1819

1920
end # module PointNeighbors

‎src/nhs_neighbor_lists.jl‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
@doc raw"""
2+
PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_particles;
3+
periodic_box_min_corner = nothing,
4+
periodic_box_max_corner = nothing)
5+
6+
Neighborhood search with precomputed neighbor lists. A list of all neighbors is computed
7+
for each particle during initialization and update.
8+
This neighborhood search maximizes the performance of neighbor loops at the cost of a much
9+
slower [`update!`](@ref).
10+
11+
A [`GridNeighborhoodSearch`](@ref) is used internally to compute the neighbor lists during
12+
initialization and update.
13+
14+
# Arguments
15+
- `NDIMS`: Number of dimensions.
16+
- `search_radius`: The uniform search radius.
17+
- `n_particles`: Total number of particles.
18+
19+
# Keywords
20+
- `periodic_box_min_corner`: In order to use a (rectangular) periodic domain, pass the
21+
coordinates of the domain corner in negative coordinate
22+
directions.
23+
- `periodic_box_max_corner`: In order to use a (rectangular) periodic domain, pass the
24+
coordinates of the domain corner in positive coordinate
25+
directions.
26+
"""
27+
struct PrecomputedNeighborhoodSearch{NDIMS, NHS, NL, PB}
28+
neighborhood_search :: NHS
29+
neighbor_lists :: NL
30+
periodic_box :: PB
31+
32+
function PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_particles;
33+
periodic_box_min_corner = nothing,
34+
periodic_box_max_corner = nothing) where {
35+
NDIMS
36+
}
37+
nhs = GridNeighborhoodSearch{NDIMS}(search_radius, n_particles,
38+
periodic_box_min_corner = periodic_box_min_corner,
39+
periodic_box_max_corner = periodic_box_max_corner)
40+
41+
neighbor_lists = Vector{Vector{Int}}()
42+
43+
new{NDIMS, typeof(nhs),
44+
typeof(neighbor_lists),
45+
typeof(nhs.periodic_box)}(nhs, neighbor_lists, nhs.periodic_box)
46+
end
47+
end
48+
49+
@inline function Base.ndims(::PrecomputedNeighborhoodSearch{NDIMS}) where {NDIMS}
50+
return NDIMS
51+
end
52+
53+
function initialize!(search::PrecomputedNeighborhoodSearch,
54+
x::AbstractMatrix, y::AbstractMatrix)
55+
(; neighborhood_search, neighbor_lists) = search
56+
57+
# Initialize grid NHS
58+
initialize!(neighborhood_search, x, y)
59+
60+
initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y)
61+
end
62+
63+
function update!(search::PrecomputedNeighborhoodSearch,
64+
x::AbstractMatrix, y::AbstractMatrix;
65+
particles_moving = (true, true))
66+
(; neighborhood_search, neighbor_lists) = search
67+
68+
# Update grid NHS
69+
update!(neighborhood_search, x, y, particles_moving = particles_moving)
70+
71+
# Skip update if both point sets are static
72+
if any(particles_moving)
73+
initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y)
74+
end
75+
end
76+
77+
function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y)
78+
# Initialize neighbor lists
79+
empty!(neighbor_lists)
80+
resize!(neighbor_lists, size(x, 2))
81+
for i in eachindex(neighbor_lists)
82+
neighbor_lists[i] = Int[]
83+
end
84+
85+
# Fill neighbor lists
86+
for_particle_neighbor(x, y, neighborhood_search) do particle, neighbor, _, _
87+
push!(neighbor_lists[particle], neighbor)
88+
end
89+
end
90+
91+
@inline function foreach_neighbor(f, system_coords, neighbor_system_coords,
92+
neighborhood_search::PrecomputedNeighborhoodSearch,
93+
particle; search_radius = nothing)
94+
(; periodic_box, neighbor_lists) = neighborhood_search
95+
(; search_radius) = neighborhood_search.neighborhood_search
96+
97+
particle_coords = extract_svector(system_coords, Val(ndims(neighborhood_search)),
98+
particle)
99+
for neighbor in neighbor_lists[particle]
100+
neighbor_coords = extract_svector(neighbor_system_coords,
101+
Val(ndims(neighborhood_search)), neighbor)
102+
103+
pos_diff = particle_coords - neighbor_coords
104+
distance2 = dot(pos_diff, pos_diff)
105+
106+
pos_diff, distance2 = compute_periodic_distance(pos_diff, distance2, search_radius,
107+
periodic_box)
108+
109+
distance = sqrt(distance2)
110+
111+
# Inline to avoid loss of performance
112+
# compared to not using `for_particle_neighbor`.
113+
@inline f(particle, neighbor, pos_diff, distance)
114+
end
115+
end

‎test/neighborhood_search.jl‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,10 +45,14 @@
4545
GridNeighborhoodSearch{NDIMS}(search_radius, n_particles,
4646
periodic_box_min_corner = periodic_boxes[i][1],
4747
periodic_box_max_corner = periodic_boxes[i][2]),
48+
PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_particles,
49+
periodic_box_min_corner = periodic_boxes[i][1],
50+
periodic_box_max_corner = periodic_boxes[i][2]),
4851
]
4952
neighborhood_searches_names = [
5053
"`TrivialNeighborhoodSearch`",
5154
"`GridNeighborhoodSearch`",
55+
"`PrecomputedNeighborhoodSearch`",
5256
]
5357

5458
# Run this for every neighborhood search
@@ -86,11 +90,14 @@
8690
]
8791

8892
seeds = [1, 2]
89-
@testset verbose=true "$(length(cloud_size))D with $(prod(cloud_size)) Particles ($(seed == 1 ? "`initialize!`" : "`update!`"))" for cloud_size in cloud_sizes,
90-
seed in seeds
93+
name(size, seed) = "$(length(size))D with $(prod(size)) Particles " *
94+
"($(seed == 1 ? "`initialize!`" : "`update!`"))"
95+
@testset verbose=true "$(name(cloud_size, seed)))" for cloud_size in cloud_sizes,
96+
seed in seeds
9197

9298
coords = point_cloud(cloud_size, seed = seed)
9399
NDIMS = length(cloud_size)
100+
n_particles = size(coords, 2)
94101
search_radius = 2.5
95102

96103
# Use different coordinates for `initialize!` and then `update!` with the
@@ -110,11 +117,13 @@
110117
end
111118

112119
neighborhood_searches = [
113-
GridNeighborhoodSearch{NDIMS}(search_radius, size(coords, 2)),
120+
GridNeighborhoodSearch{NDIMS}(search_radius, n_particles),
121+
PrecomputedNeighborhoodSearch{NDIMS}(search_radius, n_particles),
114122
]
115123

116124
neighborhood_searches_names = [
117125
"`GridNeighborhoodSearch`",
126+
"`PrecomputedNeighborhoodSearch`",
118127
]
119128

120129
@testset "$(neighborhood_searches_names[i])" for i in eachindex(neighborhood_searches_names)

0 commit comments

Comments
 (0)