Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rewrite NHS to use periodic cell indices instead of ghost cells and enable 3D periodicity #212

Merged
merged 5 commits into from
Aug 28, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 34 additions & 90 deletions src/neighborhood_search/grid_nhs.jl
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,13 @@ since that makes our implementation a lot faster (although less parallelizable).
[doi: 10.1111/J.1467-8659.2010.01832.X](https://doi.org/10.1111/J.1467-8659.2010.01832.X)
"""
struct GridNeighborhoodSearch{NDIMS, ELTYPE, PB}
hashtable :: Dict{NTuple{NDIMS, Int}, Vector{Int}}
search_radius :: ELTYPE
empty_vector :: Vector{Int} # Just an empty vector (used in `eachneighbor`)
cell_buffer :: Array{NTuple{NDIMS, Int}, 2} # Multithreaded buffer for `update!`
cell_buffer_indices :: Vector{Int} # Store which entries of `cell_buffer` are initialized
periodic_box :: PB
bound_and_ghost_cells :: Vector{NTuple{NDIMS, Int}}
hashtable :: Dict{NTuple{NDIMS, Int}, Vector{Int}}
search_radius :: ELTYPE
empty_vector :: Vector{Int} # Just an empty vector (used in `eachneighbor`)
cell_buffer :: Array{NTuple{NDIMS, Int}, 2} # Multithreaded buffer for `update!`
cell_buffer_indices :: Vector{Int} # Store which entries of `cell_buffer` are initialized
periodic_box :: PB
n_cells :: NTuple{NDIMS, Int}

function GridNeighborhoodSearch{NDIMS}(search_radius, n_particles;
min_corner=nothing,
Expand All @@ -53,12 +53,8 @@ struct GridNeighborhoodSearch{NDIMS, ELTYPE, PB}
if (min_corner === nothing && max_corner === nothing) || search_radius < eps()
# No periodicity
periodic_box = nothing
bound_and_ghost_cells = NTuple{NDIMS, Int}[]
n_cells = ntuple(_ -> -1, Val(NDIMS))
elseif min_corner !== nothing && max_corner !== nothing
if NDIMS == 3
throw(ArgumentError("periodic neighborhood search is not yet supported in 3D"))
end

periodic_box = PeriodicBox(min_corner, max_corner)

# If box size is not an integer multiple of search radius
Expand All @@ -68,75 +64,22 @@ struct GridNeighborhoodSearch{NDIMS, ELTYPE, PB}
"of `search_radius`"))
end

# Get cell index of min and max corner cells
min_cell = cell_coords(min_corner .+
0.5 * search_radius * ones(SVector{NDIMS, ELTYPE}),
search_radius, periodic_box)
max_cell = cell_coords(max_corner .-
0.5 * search_radius * ones(SVector{NDIMS, ELTYPE}),
search_radius, periodic_box)

# Initialize boundary cells with empty lists and set ghost cells to pointers
# to these lists.
bound_and_ghost_cells = initialize_boundary_and_ghost_cells(hashtable,
min_cell, max_cell)
n_cells = Tuple(round.(Int, periodic_box.size / search_radius))

if any(i -> i < 3, n_cells)
throw(ArgumentError("the `GridNeighborhoodSearch` needs at least 3 cells " *
"in each dimension when used with periodicity. " *
"Please use no NHS for very small problems."))
end
else
throw(ArgumentError("`min_corner` and `max_corner` must either be " *
"both `nothing` or both an array or tuple"))
end

new{NDIMS, ELTYPE,
typeof(periodic_box)}(hashtable, search_radius, empty_vector,
cell_buffer, cell_buffer_indices,
periodic_box, bound_and_ghost_cells)
end
end

# Initialize boundary cells with empty lists and set ghost cells to pointers to these lists
function initialize_boundary_and_ghost_cells(hashtable, min_cell, max_cell)
# To make sure that the hashtable starts empty.
# Then, we can just return the keyset after we initialized all boundary and ghost cells.
empty!(hashtable)

# Initialize boundary and ghost cells
for y in min_cell[2]:max_cell[2]
# -x boundary cells
hashtable[(min_cell[1], y)] = Int[]

# +x boundary cells
hashtable[(max_cell[1], y)] = Int[]

# -x ghost cells
hashtable[(min_cell[1] - 1, y)] = hashtable[(max_cell[1], y)]

# +x ghost cells
hashtable[(max_cell[1] + 1, y)] = hashtable[(min_cell[1], y)]
end

for x in min_cell[1]:max_cell[1]
# Avoid setting corner boundary cells again
if min_cell[1] < x < max_cell[1]
# -y boundary cells
hashtable[(x, min_cell[2])] = Int[]

# +y boundary cells
hashtable[(x, max_cell[2])] = Int[]
end

# -y ghost cells
hashtable[(x, min_cell[2] - 1)] = hashtable[(x, max_cell[2])]

# +y ghost cells
hashtable[(x, max_cell[2] + 1)] = hashtable[(x, min_cell[2])]
cell_buffer, cell_buffer_indices, periodic_box, n_cells)
end

# Corner ghost cells
hashtable[(min_cell[1] - 1, min_cell[2] - 1)] = hashtable[(max_cell[1], max_cell[2])]
hashtable[(max_cell[1] + 1, min_cell[2] - 1)] = hashtable[(min_cell[1], max_cell[2])]
hashtable[(min_cell[1] - 1, max_cell[2] + 1)] = hashtable[(max_cell[1], min_cell[2])]
hashtable[(max_cell[1] + 1, max_cell[2] + 1)] = hashtable[(min_cell[1], min_cell[2])]

return collect(keys(hashtable))
end

@inline Base.ndims(neighborhood_search::GridNeighborhoodSearch{NDIMS}) where {NDIMS} = NDIMS
Expand All @@ -156,19 +99,9 @@ function initialize!(neighborhood_search::GridNeighborhoodSearch{NDIMS},
end

function initialize!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
@unpack hashtable, bound_and_ghost_cells = neighborhood_search
@unpack hashtable = neighborhood_search

# Delete all cells that are not boundary or ghost cells
for cell in keys(hashtable)
if !(cell in bound_and_ghost_cells)
delete!(hashtable, cell)
end
end

# Empty boundary cells
for cell in bound_and_ghost_cells
empty!(hashtable[cell])
end
empty!(hashtable)

# This is needed to prevent lagging on macOS ARM.
# See https://github.com/JuliaSIMD/Polyester.jl/issues/89
Expand Down Expand Up @@ -201,8 +134,7 @@ end

# Modify the existing hash table by moving particles into their new cells
function update!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
@unpack hashtable, cell_buffer, cell_buffer_indices,
bound_and_ghost_cells = neighborhood_search
@unpack hashtable, cell_buffer, cell_buffer_indices = neighborhood_search

# Reset `cell_buffer` by moving all pointers to the beginning.
cell_buffer_indices .= 0
Expand Down Expand Up @@ -243,8 +175,7 @@ function update!(neighborhood_search::GridNeighborhoodSearch, coords_fun)
end

# Remove moved particles from this cell or delete the cell if it is now empty
if !(cell in bound_and_ghost_cells) &&
count(_ -> true, moved_particle_indices) == length(particles)
if count(_ -> true, moved_particle_indices) == length(particles)
delete!(hashtable, cell)
else
deleteat!(particles, moved_particle_indices)
Expand Down Expand Up @@ -302,7 +233,20 @@ end

# Return an empty vector when `cell_index` is not a key of `hashtable` and
# reuse the empty vector to avoid allocations
return get(hashtable, cell_index, empty_vector)
return get(hashtable, periodic_cell_index(cell_index, neighborhood_search),
empty_vector)
end

@inline function periodic_cell_index(cell_index, neighborhood_search)
@unpack n_cells, periodic_box = neighborhood_search

periodic_cell_index(cell_index, periodic_box, n_cells)
end

@inline periodic_cell_index(cell_index, ::Nothing, n_cells) = cell_index

@inline function periodic_cell_index(cell_index, periodic_box, n_cells)
return rem.(cell_index, n_cells, RoundDown)
end

@inline function cell_coords(coords, neighborhood_search)
Expand Down
40 changes: 38 additions & 2 deletions test/neighborhood_search/grid_nhs.jl
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@

neighbors_loop = [Int[] for _ in axes(coords, 2)]

TrixiParticles.for_particle_neighbor(Val(2), Val(2),
TrixiParticles.for_particle_neighbor(nothing, nothing,
coords, coords, nhs,
particles=axes(coords, 2)) do particle,
neighbor,
Expand Down Expand Up @@ -181,7 +181,7 @@

# 5 x 1 cells
nhs = GridNeighborhoodSearch{2}(1.0, size(coords, 2),
min_corner=[-1.5, 0.0], max_corner=[2.5, 1.0])
min_corner=[-1.5, 0.0], max_corner=[2.5, 3.0])

TrixiParticles.initialize!(nhs, coords)

Expand All @@ -192,4 +192,40 @@
@test neighbors[2] == [1, 2]
end
end

@testset verbose=true "Periodicity 3D" begin
coords = [-0.08 0.0 0.18 0.1 -0.08
-0.12 -0.05 -0.09 0.15 0.39
0.14 0.34 0.12 0.06 0.13]

# 3 x 6 x 3 cells
nhs = GridNeighborhoodSearch{3}(0.1, size(coords, 2),
min_corner=[-0.1, -0.2, 0.05],
max_corner=[0.2, 0.4, 0.35])

TrixiParticles.initialize!(nhs, coords)

neighbors = [sort(collect(TrixiParticles.eachneighbor(coords[:, i], nhs)))
for i in 1:5]

# Note that (1, 2) and (2, 3) are not neighbors, but they are in neighboring cells
@test neighbors[1] == [1, 2, 3, 5]
@test neighbors[2] == [1, 2, 3]
@test neighbors[3] == [1, 2, 3]
@test neighbors[4] == [4]
@test neighbors[5] == [1, 5]

neighbors_loop = [Int[] for _ in axes(coords, 2)]

TrixiParticles.for_particle_neighbor(coords, coords,
nhs) do particle, neighbor, pos_diff, distance
append!(neighbors_loop[particle], neighbor)
end

@test sort(neighbors_loop[1]) == [1, 3, 5]
@test sort(neighbors_loop[2]) == [2]
@test sort(neighbors_loop[3]) == [1, 3]
@test sort(neighbors_loop[4]) == [4]
@test sort(neighbors_loop[5]) == [1, 5]
end
end