Implement neighborhood search based on CellListMap.jl

Project.toml

@@ -13,9 +13,16 @@ Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[weakdeps]
+CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864"
+
+[extensions]
+PointNeighborsCellListMapExt = "CellListMap"
+
 [compat]
 Adapt = "3, 4"
 Atomix = "0.1, 1"
+CellListMap = "0.9"
 GPUArraysCore = "0.1, 0.2"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"

ext/PointNeighborsCellListMapExt.jl

@@ -0,0 +1,187 @@
+module PointNeighborsCellListMapExt
+
+using PointNeighbors
+using CellListMap: CellListMap, CellList, CellListPair
+
+"""
+    CellListMapNeighborhoodSearch(NDIMS; search_radius = 1.0, points_equal_neighbors = false)
+
+Neighborhood search based on the package [CellListMap.jl](https://github.com/m3g/CellListMap.jl).
+This package provides a similar implementation to the [`GridNeighborhoodSearch`](@ref)
+with [`FullGridCellList`](@ref), but with better support for periodic boundaries.
+This is just a wrapper to use CellListMap.jl with the PointNeighbors.jl API.
+Note that periodic boundaries are not yet supported.
+
+# Arguments
+- `NDIMS`: Number of dimensions.
+
+# Keywords
+- `search_radius = 1.0`:    The fixed search radius. The default of `1.0` is useful together
+                            with [`copy_neighborhood_search`](@ref).
+- `points_equal_neighbors = false`: If `true`, a `CellListMap.CellList` is used instead of
+                                    a `CellListMap.CellListPair`. This requires that `x === y`
+                                    in [`initialize!`](@ref) and [`update!`](@ref).
+                                    This option exists only for benchmarking purposes.
+                                    It makes the main loop awkward because CellListMap.jl
+                                    only computes pairs with `i < j` and PointNeighbors.jl
+                                    computes all pairs, so we have to manually use symmetry
+                                    to add the missing pairs.
+
+!!! warning "Experimental implementation"
+    This is an experimental feature and may change in future releases.
+"""
+mutable struct CellListMapNeighborhoodSearch{CL, B}
+    cell_list :: CL
+    box       :: B
+
+    # Add dispatch on `NDIMS` to avoid method overwriting of the function in PointNeighbors.jl
+    function PointNeighbors.CellListMapNeighborhoodSearch(NDIMS::Integer;
+                                                          search_radius = 1.0,
+                                                          points_equal_neighbors = false)
+        # Create a cell list with only one point and resize it later
+        x = zeros(NDIMS, 1)
+        box = CellListMap.Box(CellListMap.limits(x, x), search_radius)
+
+        if points_equal_neighbors
+            cell_list = CellListMap.CellList(x, box)
+        else
+            cell_list = CellListMap.CellList(x, x, box)
+        end
+
+        return new{typeof(cell_list), typeof(box)}(cell_list, box)
+    end
+end
+
+function PointNeighbors.search_radius(neighborhood_search::CellListMapNeighborhoodSearch)
+    return neighborhood_search.box.cutoff
+end
+
+function Base.ndims(neighborhood_search::CellListMapNeighborhoodSearch)
+    return length(neighborhood_search.box.cell_size)
+end
+
+function PointNeighbors.initialize!(neighborhood_search::CellListMapNeighborhoodSearch,
+                                    x::AbstractMatrix, y::AbstractMatrix)
+    PointNeighbors.update!(neighborhood_search, x, y)
+end
+
+# When `x !== y`, a `CellListPair` must be used
+function PointNeighbors.update!(neighborhood_search::CellListMapNeighborhoodSearch{<:CellListPair},
+                                x::AbstractMatrix, y::AbstractMatrix;
+                                points_moving = (true, true))
+    (; cell_list) = neighborhood_search
+
+    # Resize box
+    box = CellListMap.Box(CellListMap.limits(x, y), neighborhood_search.box.cutoff)
+    neighborhood_search.box = box
+
+    # Resize and update cell list
+    CellListMap.UpdateCellList!(x, y, box, cell_list)
+
+    # Recalculate number of batches for multithreading
+    CellListMap.set_number_of_batches!(cell_list)
+
+    return neighborhood_search
+end
+
+# When `points_equal_neighbors == true`, a `CellList` is used and `x` must be equal to `y`
+function PointNeighbors.update!(neighborhood_search::CellListMapNeighborhoodSearch{<:CellList},
+                                x::AbstractMatrix, y::AbstractMatrix;
+                                points_moving = (true, true))
+    (; cell_list) = neighborhood_search
+
+    @assert x===y "When `points_equal_neighbors == true`, `x` must be equal to `y`"
+
+    # Resize box
+    box = CellListMap.Box(CellListMap.limits(x), neighborhood_search.box.cutoff)
+    neighborhood_search.box = box
+
+    # Resize and update cell list
+    CellListMap.UpdateCellList!(x, box, cell_list)
+
+    # Recalculate number of batches for multithreading
+    CellListMap.set_number_of_batches!(cell_list)
+
+    # Due to https://github.com/m3g/CellListMap.jl/issues/106, we have to update again
+    CellListMap.UpdateCellList!(x, box, cell_list)
+
+    return neighborhood_search
+end
+
+# The type annotation is to make Julia specialize on the type of the function.
+# Otherwise, unspecialized code will cause a lot of allocations
+# and heavily impact performance.
+# See https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing
+function PointNeighbors.foreach_point_neighbor(f::T, system_coords, neighbor_coords,
+                                               neighborhood_search::CellListMapNeighborhoodSearch{<:CellListPair};
+                                               points = axes(system_coords, 2),
+                                               parallel = true) where {T}
+    (; cell_list, box) = neighborhood_search
+
+    # `0` is the returned output, which we don't use.
+    # Note that `parallel !== false` is `true` when `parallel` is a PointNeighbors backend.
+    CellListMap.map_pairwise!(0, box, cell_list,
+                              parallel = parallel !== false) do x, y, i, j, d2, output
+        # Skip all indices not in `points`
+        i in points || return output
+
+        pos_diff = x - y
+        distance = sqrt(d2)
+
+        @inline f(i, j, pos_diff, distance)
+
+        return output
+    end
+
+    return nothing
+end
+
+function PointNeighbors.foreach_point_neighbor(f::T, system_coords, neighbor_coords,
+                                               neighborhood_search::CellListMapNeighborhoodSearch{<:CellList};
+                                               points = axes(system_coords, 2),
+                                               parallel = true) where {T}
+    (; cell_list, box) = neighborhood_search
+
+    # `0` is the returned output, which we don't use.
+    # Note that `parallel !== false` is `true` when `parallel` is a PointNeighbors backend.
+    CellListMap.map_pairwise!(0, box, cell_list,
+                              parallel = parallel !== false) do x, y, i, j, d2, output
+        # Skip all indices not in `points`
+        i in points || return output
+
+        pos_diff = x - y
+        distance = sqrt(d2)
+
+        # When `points_equal_neighbors == true`, a `CellList` is used.
+        # With a `CellList`, we only see each pair once and have to use symmetry manually.
+        @inline f(i, j, pos_diff, distance)
+        @inline f(j, i, -pos_diff, distance)
+
+        return output
+    end
+
+    # With a `CellList`, interaction of a particle with itself is not included
+    PointNeighbors.@threaded system_coords for point in points
+        zero_pos_diff = zero(PointNeighbors.SVector{ndims(neighborhood_search),
+                                                    eltype(system_coords)})
+        @inline f(point, point, zero_pos_diff, zero(eltype(system_coords)))
+    end
+
+    return nothing
+end
+
+function PointNeighbors.copy_neighborhood_search(nhs::CellListMapNeighborhoodSearch{<:CellListPair},
+                                                 search_radius, n_points;
+                                                 eachpoint = 1:n_points)
+    return PointNeighbors.CellListMapNeighborhoodSearch(ndims(nhs); search_radius,
+                                                        points_equal_neighbors = false)
+end
+
+function PointNeighbors.copy_neighborhood_search(nhs::CellListMapNeighborhoodSearch{<:CellList},
+                                                 search_radius, n_points;
+                                                 eachpoint = 1:n_points)
+    return PointNeighbors.CellListMapNeighborhoodSearch(ndims(nhs); search_radius,
+                                                        points_equal_neighbors = true)
+end
+
+end

src/PointNeighbors.jl

@@ -21,7 +21,8 @@ include("nhs_precomputed.jl")
 include("gpu.jl")
 
 export foreach_point_neighbor, foreach_neighbor
-export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch
+export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch,
+       CellListMapNeighborhoodSearch
 export DictionaryCellList, FullGridCellList
 export ParallelUpdate, SemiParallelUpdate, SerialUpdate
 export initialize!, update!, initialize_grid!, update_grid!

src/neighborhood_search.jl

@@ -260,3 +260,9 @@ end
 @inline function periodic_coords(coords, periodic_box::Nothing)
     return coords
 end
+
+function CellListMapNeighborhoodSearch(NDIMS; search_radius = 1.0,
+                                       points_equal_neighbors = false)
+    error("CellListMap.jl has to be imported to use `CellListMapNeighborhoodSearch`. " *
+          "If you did import CellListMap.jl, please try `]instantiate`.")
+end

test/Project.toml

@@ -1,12 +1,14 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TrixiParticles = "66699cd8-9c01-4e9d-a059-b96c86d16b3a"
 
 [compat]
 BenchmarkTools = "1"
+CellListMap = "0.9"
 Plots = "1"
 Test = "1"
 TrixiParticles = "0.2"

test/neighborhood_search.jl

@@ -91,6 +91,7 @@
             # Run this for every neighborhood search
             @testset "$(names[j])" for j in eachindex(names)
                 nhs = neighborhood_searches[j]
+                @test PointNeighbors.search_radius(nhs) == search_radius
 
                 initialize!(nhs, coords, coords)
 
@@ -175,7 +176,10 @@
                                                                            max_corner,
                                                                            search_radius,
                                                                            backend = Vector{Vector{Int}})),
-                PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points)
+                PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points),
+                CellListMapNeighborhoodSearch(NDIMS; search_radius),
+                CellListMapNeighborhoodSearch(NDIMS; search_radius,
+                                              points_equal_neighbors = true)
             ]
 
             names = [
@@ -184,7 +188,9 @@
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors` and `ParallelUpdate`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors` and `SemiParallelUpdate`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
-                "`PrecomputedNeighborhoodSearch`"
+                "`PrecomputedNeighborhoodSearch`",
+                "`CellListMapNeighborhoodSearch`",
+                "`CellListMapNeighborhoodSearch` with `points_equal_neighbors = true`"
             ]
 
             # Also test copied templates
@@ -199,7 +205,9 @@
                 GridNeighborhoodSearch{NDIMS}(cell_list = FullGridCellList(; min_corner,
                                                                            max_corner,
                                                                            backend = Vector{Vector{Int32}})),
-                PrecomputedNeighborhoodSearch{NDIMS}()
+                PrecomputedNeighborhoodSearch{NDIMS}(),
+                CellListMapNeighborhoodSearch(NDIMS),
+                CellListMapNeighborhoodSearch(NDIMS, points_equal_neighbors = true)
             ]
             copied_nhs = copy_neighborhood_search.(template_nhs, search_radius, n_points)
             append!(neighborhood_searches, copied_nhs)
@@ -209,6 +217,7 @@
 
             @testset "$(names[i])" for i in eachindex(names)
                 nhs = neighborhood_searches[i]
+                @test PointNeighbors.search_radius(nhs) == search_radius
 
                 # Initialize with `seed = 1`
                 initialize!(nhs, coords_initialize, coords_initialize)

test/test_util.jl

@@ -3,6 +3,9 @@
 using Test: @test, @testset, @test_throws
 using PointNeighbors
 
+# Load `PointNeighborsCellListMapExt`
+import CellListMap
+
 """
     @trixi_testset "name of the testset" #= code to test #=