Extract list of particles, neighbors, relative positions and distances

Hey,

I want to extract the full list of particles, their neighbors, relative positions and distances. So far, I got a small MWE working on CPU with the SerialBackend()

using PointNeighbors

# Generate grid of particles
n_particles_per_dimension = (100, 100)
n_particles = prod(n_particles_per_dimension)
coordinates = Array{Float32}(undef, 2, n_particles)
cartesian_indices = CartesianIndices(n_particles_per_dimension)

for i in axes(coordinates, 2)
    coordinates[:, i] .= Tuple(cartesian_indices[i])
end

# `FullGridCellList` requires a bounding box
min_corner = minimum(coordinates, dims=2)
max_corner = maximum(coordinates, dims=2)
search_radius = 3.0f0
cell_list = FullGridCellList(; search_radius, min_corner, max_corner)
nhs = GridNeighborhoodSearch{2}(; search_radius, cell_list)

# Initialize the NHS to find neighbors in `coordinates` of particles in `coordinates`
initialize!(nhs, coordinates, coordinates)

# Simple example: just count the neighbors of each particle
n_neighbors = zeros(Int, n_particles)

# Use a function for performance reasons
function count_neighbors!(n_neighbors, coordinates, nhs)
    n_neighbors .= 0
    foreach_point_neighbor(coordinates, coordinates, nhs) do i, j, pos_diff, distance
        n_neighbors[i] += 1
    end
end

function insert_neighbor!(senders, receivers, rel_deplacement, rel_dist_norm, coordinates, nhs)
    count = 1
    foreach_point_neighbor(coordinates, coordinates, nhs; parallelization_backend= SerialBackend()) do i, j, pos_diff, distance
        senders[count] = i
        receivers[count] = j
        rel_deplacement[:,count] = pos_diff
        rel_dist_norm[count] = distance
        count += 1
    end
end

count_neighbors!(n_neighbors, coordinates, nhs)

n_edges= sum(n_neighbors)
senders = zeros(Int, 1, n_edges)
receivers = zeros(Int, 1, n_edges)
rel_deplacement = zeros(2, n_edges)
rel_dist_norm = zeros(1, n_edges)
insert_neighbor!(senders, receivers, rel_deplacement, rel_dist_norm, coordinates, nhs)

Is there a faster, more straight forward way to achieve this? In the docs, I read about PrecomputedNeighborhoodSearch which sounds like it provides the functionality I'm looking for, but I'm unsure of how to use it.

Could you provide some assistance to the matter?

Start with the provided MWE and the documentation for PrecomputedNeighborhoodSearch; compare its API with foreach_point_neighbor and SerialBackend(). Determine how to obtain senders, receivers, relative positions, and distances in one pass, then verify the result and performance against the shown extraction.