From 82350ddcc007e2b157c399e1f0454229d986e4d9 Mon Sep 17 00:00:00 2001
From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com>
Date: Wed, 8 May 2024 15:41:25 +0200
Subject: [PATCH] Implement neighborhood search based on CellListMap.jl

---
 Project.toml           |  2 ++
 src/PointNeighbors.jl  |  1 +
 src/celllistmap_nhs.jl | 63 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 src/celllistmap_nhs.jl

diff --git a/Project.toml b/Project.toml
index fc36226..b1cee12 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,12 +4,14 @@ authors = ["Erik Faulhaber <erik.faulhaber@uni-koeln.de>"]
 version = "0.2.2-dev"
 
 [deps]
+CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
+CellListMap = "0.9"
 LinearAlgebra = "1"
 Polyester = "0.7.5"
 Reexport = "1"
diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl
index 8c18400..0b7d896 100644
--- a/src/PointNeighbors.jl
+++ b/src/PointNeighbors.jl
@@ -10,6 +10,7 @@ include("util.jl")
 include("neighborhood_search.jl")
 include("trivial_nhs.jl")
 include("grid_nhs.jl")
+include("celllistmap_nhs.jl")
 
 export for_particle_neighbor
 export TrivialNeighborhoodSearch, GridNeighborhoodSearch
diff --git a/src/celllistmap_nhs.jl b/src/celllistmap_nhs.jl
new file mode 100644
index 0000000..3e92726
--- /dev/null
+++ b/src/celllistmap_nhs.jl
@@ -0,0 +1,63 @@
+mutable struct CellListMapNeighborhoodSearch{CL, B}
+    cell_list :: CL
+    box       :: B
+
+    function CellListMapNeighborhoodSearch{NDIMS}(search_radius) where {NDIMS}
+        # Create a cell list with only one particle and resize it later
+        x = zeros(NDIMS, 1)
+        box = CellListMap.Box(CellListMap.limits(x, x), search_radius)
+        cell_list = CellListMap.CellList(x, x, box)
+
+        return new{typeof(cell_list), typeof(box)}(cell_list, box)
+    end
+end
+
+function initialize!(neighborhood_search::CellListMapNeighborhoodSearch,
+                     x::AbstractMatrix, y::AbstractMatrix)
+    update!(neighborhood_search, x, y)
+end
+
+function update!(neighborhood_search::CellListMapNeighborhoodSearch,
+                 x::AbstractMatrix, y::AbstractMatrix;
+                 particles_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
+
+# 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 for_particle_neighbor(f::T, system_coords, neighbor_coords,
+                               neighborhood_search::CellListMapNeighborhoodSearch;
+                               particles = axes(system_coords, 2),
+                               parallel = true) where {T}
+    (; cell_list, box) = neighborhood_search
+
+    # `0` is the returned output, which we don't use
+    CellListMap.map_pairwise!(0, box, cell_list,
+                              parallel = parallel) do x, y, i, j, d2, output
+        # Skip all indices not in `particles`
+        i in particles || return output
+
+        pos_diff = x - y
+        distance = sqrt(d2)
+
+        @inline f(i, j, pos_diff, distance)
+
+        return output
+    end
+
+    return nothing
+end