refactor: neighbours are now exposed by the neighbour lists themselves

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