refactor: neighbours are now exposed by the neighbour lists themselves (#60)

**Summary**
This PR moves all neighbour-list related logic to the neighour
subpackage, so that all other modules are agnostic of the underlying
implementation.

This is accomplished by making each neighbour list type callable (with
`system` and `i` are arguments). The call to the instances return
generators, that when iterated upon return the indices of all the
particles adjacent to `i`.
The methods to compute energies leverage these newly defined methods
instead of encoding the neighbour list logic themselves.

**Implementation discussion**
For linked lists, I wasn't able to return a simple generator in the
dedicated calling function. It looks like Julia doesn't have some of the
features that would be required to create custom generators (something
like the `yield` keyword in Python).
Instead, I created a new struct that is used to iterate over.
This adds some complexity and makes the iteration process less clear,
which is opposite to the goal of this PR.
I believe this complexity is inherent to the linked list algorithm, and
is now fully contained in the implementation of this list (rather than
spread to other parts), so I'd argue that this PR still improves overall
clarity.

**Test plan**
CI, which contains tests that rely on these lists.
Locally, tests on a polymer chain.
Currently, unit tests on molecules fail - need to be fixed.

Author: V-Francois

src/atoms.jl

@@ -76,64 +76,13 @@ end
 
 
 """
-Compute the energy of particle `i` by brute force (no neighbour list).
-
-`compute_energy_particle(system, i, ::EmptyList)` sums interactions of particle `i`
-with all other particles using `compute_energy_ij`.
+Compute the energy of particle `i` using the provided neighbour list.
 """
-function compute_energy_particle(system::Atoms, i, ::EmptyList)
+function compute_energy_particle(system::Atoms, i, neighbour_list::NeighbourList)
     energy_i = zero(typeof(system.density))
     position_i = get_position(system, i)
-    for (j, _) in enumerate(system)
+    for j in neighbour_list(system, i)
         energy_i += compute_energy_ij(system, i, j, position_i)
     end
     return energy_i
 end
-
-
-"""
-Compute the energy of particle `i` using a `CellList` neighbour list.
-
-This restricts pair evaluations to particles in neighbouring cells of `i`.
-"""
-function compute_energy_particle(system::Atoms, i, neighbour_list::CellList)
-    energy_i = zero(typeof(system.density))
-    position_i = get_position(system, i)
-    c = get_cell_index(position_i, neighbour_list)
-    neighbour_cells = neighbour_list.neighbour_cells[c]
-
-    # Scan the neighbourhood of cell mc (including itself)
-    @inbounds for c2 in neighbour_cells
-        # Scan atoms in cell c2
-        neighbours = neighbour_list.cells[c2]
-        @inbounds for j in neighbours
-            energy_i += compute_energy_ij(system, i, j, position_i)
-        end
-    end
-    return energy_i
-end
-
-"""
-Compute the energy of particle `i` using a `LinkedList` neighbour list.
-
-This variant iterates linked list heads for neighbouring cells and accumulates
-pair energies computed with `compute_energy_ij`.
-"""
-function compute_energy_particle(system::Atoms, i, neighbour_list::LinkedList)
-    energy_i = zero(typeof(system.density))
-    # Get cell of particle i
-    position_i = get_position(system, i)
-    c = get_cell_index(position_i, neighbour_list)
-    neighbour_cells = neighbour_list.neighbour_cells[c]
-    # Scan the neighbourhood of cell mc (including itself)
-    @inbounds for c2 in neighbour_cells
-        # Scan atoms in cell c2
-        j = neighbour_list.head[c2]
-        while (j != -1)
-            energy_ij = compute_energy_ij(system, i, j, position_i)
-            energy_i += energy_ij
-            j = neighbour_list.list[j]
-        end
-    end
-    return energy_i
-end

src/molecules.jl

@@ -112,15 +112,15 @@ Return arrays of first indices and counts for consecutive blocks in `vec`.
 function get_first_and_counts(vec::Vector{Int})
     firsts = Int[]
     counts = Int[]
-    
+
     # Handle empty vector case
     isempty(vec) && return firsts, counts
-    
+
     # Initialize with first element
     current = vec[1]
     push!(firsts, 1)
     count = 1
-    
+
     # Scan through vector
     @inbounds for i in 2:length(vec)
         if vec[i] != current
@@ -134,7 +134,7 @@ function get_first_and_counts(vec::Vector{Int})
     end
     # Add last count
     push!(counts, count)
-    
+
     return firsts, counts
 end
 
@@ -198,69 +198,18 @@ function compute_energy_ij(system::Molecules, position_i, position_j, model_ij::
 end
 
 """
-Compute particle energy by brute force (no neighbour list).
-
-`compute_energy_particle(system, i, ::EmptyList)` sums interactions of particle `i`
-with all particles (including bonded and non-bonded contributions via helper
-functions). Used when no neighbour list is available.
-"""
-function compute_energy_particle(system::Molecules, i, ::EmptyList)
-    energy = zero(typeof(system.density))
-    position_i = system.position[i]
-    bonds_i = system.bonds[i]
-    @inbounds for j in eachindex(system)
-        energy += check_compute_energy_ij(system, i, j, position_i, bonds_i)
-    end
-    return energy
-end
-
-# With linked list
-"""
-Compute particle energy using a `LinkedList` neighbour list.
+Compute particle energy using a the provided neighbour list.
 
-This variant restricts non-bonded pair evaluation to particles in neighbouring
-cells defined by the linked list; bonded contributions are added explicitly.
+Non-bonded pair energy evaluations are restricted to particles in the neighbour list;
+bonded contributions are added explicitly.
 """
-function compute_energy_particle(system::Molecules, i, neighbour_list::LinkedList)
-    energy_i = zero(typeof(system.density))
-    # Get cell of particle i
+function compute_energy_particle(system::Molecules, i, neighbour_list::NeighbourList)
     position_i = system.position[i]
-    c = get_cell_index(i, neighbour_list)
-    cells = neighbour_list.neighbour_cells[c]
-    # Scan the neighbourhood of cell mc (including itself)
     bonds_i = system.bonds[i]
-    energy_i += compute_energy_bonded_i(system, i, position_i, bonds_i)
-    @inbounds for c2 in cells
-        # Calculate the scalar cell index of the neighbour cell (with PBC)
-        j = neighbour_list.head[c2]
-        while (j != -1)
-            energy_i += check_nonbonded_compute_energy_ij(system, i, j, position_i, bonds_i)
-            j = neighbour_list.list[j]
-        end
-    end
-    return energy_i
-end
 
-"""
-Compute particle energy using a `CellList` neighbour list.
-
-This variant restricts non-bonded pair evaluation to particles in neighbouring
-cells defined by the cell list; bonded contributions are added explicitly.
-"""
-function compute_energy_particle(system::Molecules, i, neighbour_list::CellList)
-    energy_i = zero(typeof(system.density))
-    position_i = get_position(system, i)
-    c = get_cell_index(i, neighbour_list)
-    neighbour_cells = neighbour_list.neighbour_cells[c]
-    # Scan the neighbourhood of cell mc (including itself)
-    bonds_i = system.bonds[i]
-    energy_i += compute_energy_bonded_i(system, i, position_i, bonds_i)
-    @inbounds for c2 in neighbour_cells
-        # Scan atoms in cell c2
-        neighbours = neighbour_list.cells[c2]
-        @inbounds for j in neighbours
-            energy_i += check_nonbonded_compute_energy_ij(system, i, j, position_i, bonds_i)
-        end
+    energy_i = compute_energy_bonded_i(system, i, position_i, bonds_i)
+    for j in neighbour_list(system, i)
+        energy_i += check_nonbonded_compute_energy_ij(system, i, j, position_i, bonds_i)
     end
     return energy_i
 end
@@ -290,4 +239,4 @@ function compute_chain_correlation(system::Molecules)
         end
     end
     return sum(correlation_array.^2)
-end
+end

src/neighbours.jl

@@ -42,6 +42,15 @@ function old_new_cell(::Particles, i, ::EmptyList)
     return 1, 1
 end
 
+"""Calling an EmptyList objects return an object which can be iterated upon.
+
+This iteration will return the indices of the neighbours (which for this list is all the other particles in the system).
+"""
+function (empty_list::EmptyList)(system::Particles, ::Int)
+    return (j for j in 1:length(system))
+end
+
+
 """Return the scalar cell index of particle `i` stored in `neighbour_list`.
 """
 function get_cell_index(i::Int, neighbour_list::NeighbourList)
@@ -195,6 +204,20 @@ function old_new_cell(system::Particles, i, neighbour_list::CellList)
     return c, c2
 end
 
+"""Calling a CellList objects return an object which can be iterated upon.
+
+This iteration will return the indices of the neighbours of particle i.
+"""
+function (cell_list::CellList)(system::Particles, i::Int)
+    position_i = get_position(system, i)
+    c = get_cell_index(position_i, cell_list)
+    neighbour_cells = cell_list.neighbour_cells[c]
+    # Scan the neighbourhood of cell mc (including itself)
+    # and from there scan atoms in cell c2
+    return (j for c2 in neighbour_cells for j in @inbounds cell_list.cells[c2])
+end
+
+
 """Linked-list neighbour list implementation.
 
 Uses arrays `head` and `list` to store per-cell linked lists of particle indices.
@@ -281,3 +304,68 @@ function old_new_cell(system::Particles, i, neighbour_list::LinkedList)
     c2 = cell_index(neighbour_list, mc2)
     return c, c2
 end
+
+""" This struct is used to iterate over neighbours of a Linked list
+"""
+struct LinkedIterator
+    neighbour_cells::Vector{Int}
+    head::Vector{Int}
+    list::Vector{Int}
+end
+
+# To iterate over the neighbours of a linked list, one could write the following loops
+#@inbounds for c in neighbour_list.neighbour_cells
+#    j = neighbour_list.head[c]
+#    while (j != -1)
+#        do stuff
+#        j = neighbour_list.list[j]
+#    end
+#end
+# This is however impossible to rewrite as a simple generator
+# So we implement the following function, which uses a state to carry over the needed information
+function Base.iterate(neighbour_list::LinkedIterator, state=-1)
+    # First time in
+    if state == -1
+        j = -1
+        c_state = 1
+        # The while loop is necessary, in case the first head is -1
+        while j == -1
+            next = iterate(neighbour_list.neighbour_cells, c_state)
+            if next == nothing
+                return nothing
+            end
+            c, c_state = next
+            @inbounds j = neighbour_list.head[c]
+        end
+    else
+        c_state, j = state
+        @inbounds j = neighbour_list.list[j]
+        # The while loop is necessary, in case a head is -1
+        while j == -1
+            next = iterate(neighbour_list.neighbour_cells, c_state)
+            if next == nothing
+                return nothing
+            end
+            c, c_state = next
+            @inbounds j = neighbour_list.head[c]
+        end
+    end
+
+    if j == -1
+        return nothing
+    end
+    state = (c_state, j)
+    return j, state
+end
+
+"""Calling a LinkedList objects return an object which can be iterated upon.
+
+This iteration will return the indices of the neighbours of particle i.
+"""
+function (linked_list::LinkedList)(system::Particles, i::Int)
+    position_i = get_position(system, i)
+    c = get_cell_index(position_i, linked_list)
+    neighbour_cells = linked_list.neighbour_cells[c]
+
+    return LinkedIterator(neighbour_cells, linked_list.head, linked_list.list)
+end