NetworkGrid: modify get_neighbors and create get_neighborhood

Author: Tortar

examples/boltzmann_wealth_model_network/boltzmann_wealth_model_network/model.py

@@ -57,7 +57,7 @@ def __init__(self, unique_id, model):
     def move(self):
         possible_steps = [
             node
-            for node in self.model.grid.get_neighbors(self.pos, include_center=False)
+            for node in self.model.grid.get_neighborhood(self.pos, include_center=False)
             if self.model.grid.is_cell_empty(node)
         ]
         if len(possible_steps) > 0:
@@ -66,8 +66,7 @@ def move(self):
 
     def give_money(self):
 
-        neighbors_nodes = self.model.grid.get_neighbors(self.pos, include_center=False)
-        neighbors = self.model.grid.get_cell_list_contents(neighbors_nodes)
+        neighbors = self.model.grid.get_neighbors(self.pos, include_center=False)
         if len(neighbors) > 0:
             other = self.random.choice(neighbors)
             other.wealth += 1

examples/virus_on_network/virus_on_network/model.py

@@ -124,7 +124,9 @@ def __init__(
         self.gain_resistance_chance = gain_resistance_chance
 
     def try_to_infect_neighbors(self):
-        neighbors_nodes = self.model.grid.get_neighbors(self.pos, include_center=False)
+        neighbors_nodes = self.model.grid.get_neighborhood(
+            self.pos, include_center=False
+        )
         susceptible_neighbors = [
             agent
             for agent in self.model.grid.get_cell_list_contents(neighbors_nodes)

mesa/space.py

@@ -1015,12 +1015,17 @@ def place_agent(self, agent: Agent, node_id: int) -> None:
         self.G.nodes[node_id]["agent"].append(agent)
         agent.pos = node_id
 
-    def get_neighbors(self, node_id: int, include_center: bool = False) -> list[int]:
+    def get_neighborhood(self, node_id: int, include_center: bool = False) -> list[int]:
         """Get all adjacent nodes"""
-        neighbors = list(self.G.neighbors(node_id))
+        neighborhood = list(self.G.neighbors(node_id))
         if include_center:
-            neighbors.append(node_id)
-        return neighbors
+            neighborhood.append(node_id)
+        return neighborhood
+
+    def get_neighbors(self, node_id: int, include_center: bool = False) -> list[Agent]:
+        """Get all agents in adjacent nodes."""
+        neighborhood = self.get_neighborhood(node_id, include_center)
+        return self.get_cell_list_contents(neighborhood)
 
     def move_agent(self, agent: Agent, node_id: int) -> None:
         """Move an agent from its current node to a new node."""
@@ -1046,7 +1051,7 @@ def get_cell_list_contents(self, cell_list: list[int]) -> list[Agent]:
             for node_id in cell_list
             if not self.is_cell_empty(node_id)
         ]
-        return [item for sublist in list_of_lists for item in sublist]
+        return [agent for sublist in list_of_lists for agent in sublist]
 
     def get_all_cell_contents(self) -> list[Agent]:
         """Returns a list of all the agents in the network."""

tests/test_space.py

@@ -358,11 +358,11 @@ def test_agent_positions(self):
 
     def test_get_neighbors(self):
         assert (
-            len(self.space.get_neighbors(0, include_center=True))
+            len(self.space.get_neighborhood(0, include_center=True))
             == TestSingleNetworkGrid.GRAPH_SIZE
         )
         assert (
-            len(self.space.get_neighbors(0, include_center=False))
+            len(self.space.get_neighborhood(0, include_center=False))
             == TestSingleNetworkGrid.GRAPH_SIZE - 1
         )
 
@@ -433,11 +433,11 @@ def test_agent_positions(self):
 
     def test_get_neighbors(self):
         assert (
-            len(self.space.get_neighbors(0, include_center=True))
+            len(self.space.get_neighborhood(0, include_center=True))
             == TestMultipleNetworkGrid.GRAPH_SIZE
         )
         assert (
-            len(self.space.get_neighbors(0, include_center=False))
+            len(self.space.get_neighborhood(0, include_center=False))
             == TestMultipleNetworkGrid.GRAPH_SIZE - 1
         )