Network: Generates square lattice neighbours

Generates a neighbour list such that agents are on a square lattice, and
each agent has 4 neighbours. PBCs are implemented. TODO: perhaps move
from size_t to int?

Co-authored-by: Moritz Sallermann <[email protected]>

include/network.hpp

@@ -258,9 +258,9 @@ class Network
             std::iota( sorting_indices.begin(), sorting_indices.end(), 0 );
 
             // Then, we figure out how to sort the neighbour indices list
-            std::sort( sorting_indices.begin(), sorting_indices.end(), [&]( auto i1, auto i2 ) {
-                return neighbours[i1] < neighbours[i2];
-            } );
+            std::sort(
+                sorting_indices.begin(), sorting_indices.end(),
+                [&]( auto i1, auto i2 ) { return neighbours[i1] < neighbours[i2]; } );
 
             std::optional<size_t> last_neighbour_index = std::nullopt;
             for( size_t i = 0; i < n_neighbours; i++ )

include/network_generation.hpp

@@ -4,6 +4,7 @@
 #include <random>
 #include <util/math.hpp>
 #include <util/misc.hpp>
+#include <vector>
 
 namespace Seldon::NetworkGeneration
 {
@@ -72,6 +73,7 @@ generate_n_connections( size_t n_agents, size_t n_connections, bool self_interac
     return NetworkT( std::move( neighbour_list ), std::move( weight_list ), NetworkT::EdgeDirection::Incoming );
 }
 
+// @TODO generate_fully_connected does not need to be overloaded..perhaps a std::optional instead to reduce code duplication?
 template<typename AgentType>
 Network<AgentType> generate_fully_connected( size_t n_agents, typename Network<AgentType>::WeightT weight = 0.0 )
 {
@@ -234,4 +236,62 @@ Network<AgentType> generate_from_file( const std::string & file )
 
     return NetworkT( std::move( neighbour_list ), std::move( weight_list ), NetworkT::EdgeDirection::Incoming );
 }
+
+/* Constructs a new network on a square lattice of edge length n_edge (with PBCs)
+ */
+template<typename AgentType>
+Network<AgentType> generate_square_lattice( size_t n_edge, typename Network<AgentType>::WeightT weight = 0.0 )
+{
+    using NetworkT = Network<AgentType>;
+    using WeightT  = typename NetworkT::WeightT;
+    auto n_agents  = n_edge * n_edge;
+
+    // Create an empty Network
+    auto network = NetworkT( n_agents );
+
+    auto wrap_edge_index = [&]( int k )
+    {
+        if( k >= int( n_edge ) )
+        {
+            return k - int( n_edge );
+        }
+        else if( k < 0 )
+        {
+            return int( n_edge ) + k;
+        }
+        else
+        {
+            return k;
+        }
+    };
+
+    auto linear_index = [&]( int i, int j )
+    {
+        auto idx = wrap_edge_index( i ) + n_edge * wrap_edge_index( j );
+        return idx;
+    };
+
+    for( int i = 0; i < int( n_edge ); i++ )
+    {
+        // other edge
+        for( int j = 0; j < int( n_edge ); j++ )
+        {
+            // Central agent
+            auto central_index = linear_index( i, j );
+
+            // clang-format off
+            std::vector<size_t> neighbours = {
+                linear_index( i - 1, j ), 
+                linear_index( i + 1, j ), 
+                linear_index( i, j - 1 ),
+                linear_index( i, j + 1 )
+            };
+
+            // clang-format on
+            network.set_neighbours_and_weights( central_index, neighbours, weight );
+        }
+    }
+
+    return network;
+}
 } // namespace Seldon::NetworkGeneration

test/test_network.cpp

@@ -171,4 +171,29 @@ TEST_CASE( "Testing the network class" )
             REQUIRE_THAT( weights, Catch::Matchers::RangeEquals( weights_no_double_counting[i_agent] ) );
         }
     }
+
+    SECTION( "Test the generation of a square lattice neighbour list for three agents" )
+    {
+        // clang-format off
+        std::vector<std::vector<size_t>> desired_neighbour_list =  { 
+                { 2, 1, 3, 6 }, 
+                { 2, 0, 4, 7 }, 
+                { 0, 1, 5, 8 },
+                { 4, 5, 6, 0 },
+                { 5, 3, 1, 7 },
+                { 3, 4, 2, 8 },
+                { 7, 8, 3, 0 },
+                { 8, 6, 4, 1 },
+                { 7, 6, 5, 2 }, 
+            };
+        // clang-format on
+
+        auto network = Seldon::NetworkGeneration::generate_square_lattice<double>( 3 );
+
+        for( size_t i_agent = 0; i_agent < network.n_agents(); i_agent++ )
+        {
+            auto neighbours = network.get_neighbours( i_agent );
+            REQUIRE_THAT( neighbours, Catch::Matchers::UnorderedRangeEquals( desired_neighbour_list[i_agent] ) );
+        }
+    }
 }