Extract list of particles, neighbors, relative positions and distances

[JoKircher]

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?

[efaulhaber]

The PrecomputedNeighborhoodSearch stores the neighbors of each particle (but not the relative positions and the distances) and then behaves exactly like a GridNeighborhoodSearch. It is designed to be a different underlying implementation that has the same API but provides the fastest possible query (foreach_point_neighbor) at the cost of a very expensive update.

The first question is: why do you need this to be faster? Is this a bottleneck right now? I guess this is only done once in the initialization of your code.

If you really need more performance, you can only get it through parallelization, but this is not trivial in this case.
If you want to use parallelization, the count += 1 version does not work. So you will need to first collect the neighbors and corresponding data per point and then copy the data into your data format indexed per pair.
Something like this:

max_neighbors = 100
n_neighbors = zeros(Int, n_particles)
neighbors = zeros(Int, max_neighbors, n_particles)

foreach_point_neighbor(...) do i, j, _, _
    # This loop is parallelized over `i`, so for each `i`, the loop over the neighbors will be serial.
    # Therefore, the following does not cause any race conditions.
    n_neighbors[i] += 1
    neighbors[n_neighbors[i], i] = j
end

Then, you can copy this data into your format. In order to parallelize this, you can use a parallel loop over the particles and find the correct index in the large list with cumsum(n_neighbors).

[JoKircher]

I'm working with Graph Neural Networks where I have to recompute the graph and therefore the neighborhood for each solver step. In order to pass the graph to the Neural Network I need it to be completely constructed, so a full list of edges (points and their neighbors) with relative positions and distances is needed. The neighborhood search is happening quite frequently so I'm looking at other implementations to realize that.

I try your proposed solution and get back to you.

[efaulhaber]

Then I see three options:

1. Try what you did first and see if this is a bottleneck or not.
2. Implement the parallel version with cumsum as mentioned above.
3. If the construction of the neighborhood is still the bottleneck, I don't think there is another way to speed this up significantly. The full list is very large, so it will always be slow to fully assemble it. The only way to speed this up would be to modify the neural network so that it doesn't need a fully assembled list, similar to a matrix-free solver. Note that I have no idea of neural networks, so I don't know if this is possible.

Since you're working with neural networks, does that mean you work with GPUs? In this case, you can probably get a very significant speedup by implementing the parallel construction in step 2 because you can throw this on the GPU. I have a PR somewhere (I think #10) where I use a GPU-compatible data structure to do a fully parallel and allocation-free neighbor list construction. You could extend this approach to get the distances and relative positions as well.

[JoKircher]

I work with GPUs, so I will check that PR and the options you mentioned. Thanks!

[efaulhaber]

The PR is a bit messy because it is based on another PR.
What is happening is that instead of using a vector of vectors for the neighbors, I have a DynamicVectorOfVectors, which is a data structure that behaves like a vector of vectors, but the underlying data structure is a large matrix where each column corresponds to the neighbors of one particle.
This is one contiguous block of memory and allocation-free, so it is GPU-compatible.