Remove remaining mesa.space usage (#333)

Author: falloficaruss

examples/conways_game_of_life_fast/model.py

@@ -1,7 +1,7 @@
 import numpy as np
 from mesa import Model
 from mesa.datacollection import DataCollector
-from mesa.space import PropertyLayer
+from mesa.discrete_space import PropertyLayer
 from scipy.signal import convolve2d
 
 
@@ -10,7 +10,9 @@ class GameOfLifeModel(Model):
     def __init__(self, width=10, height=10, alive_fraction=0.2):
         super().__init__()
         # Initialize the property layer for cell states
-        self.cell_layer = PropertyLayer("cells", width, height, False, dtype=bool)
+        self.cell_layer = PropertyLayer(
+            "cells", (width, height), default_value=False, dtype=bool
+        )
         # Randomly set cells to alive
         self.cell_layer.data = np.random.choice([True, False], size=(width, height), p=[alive_fraction, 1 - alive_fraction])
 

examples/dining_philosophers/app.py

@@ -44,8 +44,8 @@ def DiningTable(model):
         )
     )
 
-    for agent in model.grid.get_all_cell_contents():
-        node_id = agent.pos
+    for agent in model.grid.agents:
+        node_id = agent.cell.coordinate
         x, y = pos[node_id]
 
         # Map (-1, 1) to (50, 350)
@@ -151,10 +151,9 @@ def SpaghettiBarChart(model):
     # We want them sorted by ID (position)
 
     philosophers = []
-    # Access internal agent storage for efficiency or just filter grid
-    # model.philosophers is an AgentSet which is iterable
-
-    philosophers = sorted(model.philosophers, key=lambda p: p.position)
+    philosophers = sorted(
+        model.agents_by_type[PhilosopherAgent], key=lambda p: p.position
+    )
 
     labels = [f"P{p.position // 2}" for p in philosophers]
     values = [p.total_eaten for p in philosophers]

examples/dining_philosophers/dining_philosophers/agent.py

@@ -1,6 +1,6 @@
 from enum import Enum, auto
 
-import mesa
+from mesa.discrete_space import FixedAgent
 
 
 class State(Enum):
@@ -9,28 +9,34 @@ class State(Enum):
     EATING = auto()
 
 
-class ForkAgent(mesa.Agent):
-    def __init__(self, model, position):
+class ForkAgent(FixedAgent):
+    def __init__(self, model):
         super().__init__(model)
-        self.position = position
         self.is_used = False
         self.owner = None
 
+    @property
+    def position(self):
+        return self.cell.coordinate
+
     def __repr__(self):
         owner = None if self.owner is None else getattr(self.owner, "position", None)
         return f"Fork-{self.position}(used={self.is_used}, owner={owner})"
 
 
-class PhilosopherAgent(mesa.Agent):
-    def __init__(self, model, position):
+class PhilosopherAgent(FixedAgent):
+    def __init__(self, model):
         super().__init__(model)
-        self.position = position
         self.state = State.THINKING
         self.total_eaten = 0
         self.ticks_since_state_change = 0
         self.total_wait_time = 0
         self.eating_count = 0
 
+    @property
+    def position(self):
+        return self.cell.coordinate
+
     def _start_eating(self):
         self.total_wait_time += self.ticks_since_state_change
         self.eating_count += 1
@@ -59,8 +65,7 @@ def step(self):
             self.ticks_since_state_change = 0
 
     def put_down_forks(self):
-        neighbors = self.model.grid.get_neighbors(self.pos, include_center=False)
-        my_forks = [a for a in neighbors if isinstance(a, ForkAgent)]
+        my_forks = self._get_neighbor_forks()
 
         for fork in my_forks:
             if fork.owner == self:
@@ -78,8 +83,7 @@ def try_to_eat(self):
             self.eat_strategy_cooperative()
 
     def eat_strategy_naive(self):
-        neighbors = self.model.grid.get_neighbors(self.pos, include_center=False)
-        my_forks = [a for a in neighbors if isinstance(a, ForkAgent)]
+        my_forks = self._get_neighbor_forks()
 
         left_pos = (self.position - 1) % self.model.num_nodes
         right_pos = (self.position + 1) % self.model.num_nodes
@@ -100,8 +104,7 @@ def eat_strategy_naive(self):
                 pass
 
     def eat_strategy_atomic(self):
-        neighbors = self.model.grid.get_neighbors(self.pos, include_center=False)
-        my_forks = [a for a in neighbors if isinstance(a, ForkAgent)]
+        my_forks = self._get_neighbor_forks()
 
         if all(not fork.is_used for fork in my_forks):
             for fork in my_forks:
@@ -110,8 +113,7 @@ def eat_strategy_atomic(self):
             self._start_eating()
 
     def eat_strategy_cooperative(self):
-        neighbors = self.model.grid.get_neighbors(self.pos, include_center=False)
-        my_forks = [a for a in neighbors if isinstance(a, ForkAgent)]
+        my_forks = self._get_neighbor_forks()
 
         # If any fork is used, we can't eat anyway
         if any(fork.is_used for fork in my_forks):
@@ -123,7 +125,7 @@ def eat_strategy_cooperative(self):
 
         neighbors_p = [
             p
-            for p in self.model.philosophers
+            for p in self.model.agents_by_type[PhilosopherAgent]
             if p.position in (left_p_pos, right_p_pos)
         ]
 
@@ -151,3 +153,10 @@ def eat_strategy_cooperative(self):
 
     def __repr__(self):
         return f"Phil-{self.position}({self.state.name})"
+
+    def _get_neighbor_forks(self):
+        return [
+            agent
+            for agent in self.cell.neighborhood.agents
+            if isinstance(agent, ForkAgent)
+        ]

examples/dining_philosophers/dining_philosophers/model.py

@@ -1,5 +1,6 @@
 import mesa
 import networkx as nx
+from mesa.discrete_space import Network
 
 from .agent import ForkAgent, PhilosopherAgent, State
 
@@ -20,41 +21,52 @@ def __init__(
         self.num_nodes = num_philosophers * 2
 
         self.G = nx.circulant_graph(self.num_nodes, [1])
-        self.grid = mesa.space.NetworkGrid(self.G)
-
-        philosophers_list = []
+        self.grid = Network(self.G, random=self.random)
 
         for node_id in range(self.num_nodes):
+            node_cell = self.grid[node_id]
             if node_id % 2 == 0:
-                p = PhilosopherAgent(self, node_id)
-                self.grid.place_agent(p, node_id)
-                philosophers_list.append(p)
+                p = PhilosopherAgent(self)
+                p.cell = node_cell
             else:
-                f = ForkAgent(self, node_id)
-                self.grid.place_agent(f, node_id)
-
-        self.philosophers = mesa.agent.AgentSet(philosophers_list, random=self.random)
+                f = ForkAgent(self)
+                f.cell = node_cell
 
         model_reporters = {
             "Eating": lambda m: len(
-                [p for p in m.philosophers if p.state == State.EATING]
+                [
+                    p
+                    for p in m.agents_by_type[PhilosopherAgent]
+                    if p.state == State.EATING
+                ]
             ),
             "Hungry": lambda m: len(
-                [p for p in m.philosophers if p.state == State.HUNGRY]
+                [
+                    p
+                    for p in m.agents_by_type[PhilosopherAgent]
+                    if p.state == State.HUNGRY
+                ]
             ),
             "Thinking": lambda m: len(
-                [p for p in m.philosophers if p.state == State.THINKING]
+                [
+                    p
+                    for p in m.agents_by_type[PhilosopherAgent]
+                    if p.state == State.THINKING
+                ]
             ),
             "Avg Wait Time": lambda m: (
                 (
-                    sum(p.total_wait_time for p in m.philosophers)
-                    / sum(p.eating_count for p in m.philosophers)
+                    sum(p.total_wait_time for p in m.agents_by_type[PhilosopherAgent])
+                    / sum(p.eating_count for p in m.agents_by_type[PhilosopherAgent])
                 )
-                if sum(p.eating_count for p in m.philosophers) > 0
+                if sum(p.eating_count for p in m.agents_by_type[PhilosopherAgent]) > 0
                 else 0
             ),
             "Throughput": lambda m: (
-                (sum(p.total_eaten for p in m.philosophers) / m.time)
+                (
+                    sum(p.total_eaten for p in m.agents_by_type[PhilosopherAgent])
+                    / m.time
+                )
                 if m.time > 0
                 else 0
             ),
@@ -68,11 +80,18 @@ def __init__(
 
             # Use default argument in lambda to capture the current value of p_id
             model_reporters[f"P{i}"] = lambda m, pid=p_id: (
-                next((p.total_eaten for p in m.philosophers if p.position == pid), 0)
+                next(
+                    (
+                        p.total_eaten
+                        for p in m.agents_by_type[PhilosopherAgent]
+                        if p.position == pid
+                    ),
+                    0,
+                )
             )
 
         self.datacollector = mesa.DataCollector(model_reporters=model_reporters)
 
     def step(self):
-        self.philosophers.shuffle_do("step")
+        self.agents_by_type[PhilosopherAgent].shuffle_do("step")
         self.datacollector.collect(self)

examples/hex_ant/model.py

@@ -1,6 +1,8 @@
 import mesa
 from mesa.discrete_space import HexGrid
 
+from .agent import Ant
+
 
 class AntForaging(mesa.Model):
     """
@@ -80,8 +82,6 @@ def _init_agents(self, num_ants):
         center = (self.grid.width // 2, self.grid.height // 2)
         center_cell = self.grid[center]
 
-        from .agent import Ant
-
         for _ in range(num_ants):
             ant = Ant(self)
             # Add agent to the cell (spatial placement)

examples/warehouse/warehouse/agents.py

@@ -132,14 +132,17 @@ def continue_task(self):
 
         if status == "movement complete" and self.meta_agent.status == "inventory":
             # Pick up item and bring to loading dock
+            source_coordinate = self.meta_agent.cell.coordinate
+            target_level = self.item.cell.coordinate[2]
             self.meta_agent.cell = self.model.warehouse[
-                *self.meta_agent.cell.coordinate[:2], self.item.cell.coordinate[2]
+                (source_coordinate[0], source_coordinate[1], target_level)
             ]
             self.meta_agent.status = "loading"
             self.carrying = self.item.item
             self.item.quantity -= 1
+            loading_coordinate = self.meta_agent.cell.coordinate
             self.meta_agent.cell = self.model.warehouse[
-                *self.meta_agent.cell.coordinate[:2], 0
+                (loading_coordinate[0], loading_coordinate[1], 0)
             ]
             self.path = self.find_path(self.cell, self.loading_dock)
 

gis/agents_and_networks/src/agent/building.py

@@ -8,16 +8,18 @@
 import pyproj
 from shapely.geometry import Polygon
 
+FloatCoordinate = tuple[float, float]
+
 
 class Building(mg.GeoAgent):
     unique_id: int  # an ID that represents the building
     model: mesa.Model
     geometry: Polygon
     crs: pyproj.CRS
-    centroid: mesa.space.FloatCoordinate
+    centroid: FloatCoordinate
     name: str
     function: float  # 1.0 for work, 2.0 for home, 0.0 for neither
-    entrance_pos: mesa.space.FloatCoordinate  # nearest vertex on road
+    entrance_pos: FloatCoordinate  # nearest vertex on road
 
     def __init__(self, model, geometry, crs) -> None:
         super().__init__(model=model, geometry=geometry, crs=crs)

gis/agents_and_networks/src/agent/commuter.py

@@ -11,6 +11,8 @@
 from ..space.utils import UnitTransformer, redistribute_vertices
 from .building import Building
 
+FloatCoordinate = tuple[float, float]
+
 
 class Commuter(mg.GeoAgent):
     unique_id: int  # commuter_id, used to link commuters and nodes
@@ -20,7 +22,7 @@ class Commuter(mg.GeoAgent):
     origin: Building  # where he begins his trip
     destination: Building  # the destination he wants to arrive at
     my_path: list[
-        mesa.space.FloatCoordinate
+        FloatCoordinate
     ]  # a set containing nodes to visit in the shortest path
     step_in_path: int  # the number of step taking in the walk
     my_home: Building

gis/agents_and_networks/src/logger.py

@@ -1,9 +1,8 @@
 import logging
+from functools import wraps
 
 
 def logger(func):
-    from functools import wraps
-
     @wraps(func)
     def wrapper(*args, **kwargs):
         logger = logging.getLogger(func.__name__)

gis/agents_and_networks/src/space/campus.py

@@ -2,21 +2,22 @@
 from collections import defaultdict
 from typing import DefaultDict
 
-import mesa
 import mesa_geo as mg
 from shapely.geometry import Point
 
 from ..agent.building import Building
 from ..agent.commuter import Commuter
 
+FloatCoordinate = tuple[float, float]
+
 
 class Campus(mg.GeoSpace):
     homes: tuple[Building]
     works: tuple[Building]
     other_buildings: tuple[Building]
-    home_counter: DefaultDict[mesa.space.FloatCoordinate, int]
+    home_counter: DefaultDict[FloatCoordinate, int]
     _buildings: dict[int, Building]
-    _commuters_pos_map: DefaultDict[mesa.space.FloatCoordinate, set[Commuter]]
+    _commuters_pos_map: DefaultDict[FloatCoordinate, set[Commuter]]
     _commuter_id_map: dict[int, Commuter]
 
     def __init__(self, crs: str) -> None:
@@ -54,9 +55,7 @@ def add_buildings(self, agents) -> None:
         self.works = self.works + tuple(works)
         self.homes = self.homes + tuple(homes)
 
-    def get_commuters_by_pos(
-        self, float_pos: mesa.space.FloatCoordinate
-    ) -> set[Commuter]:
+    def get_commuters_by_pos(self, float_pos: FloatCoordinate) -> set[Commuter]:
         return self._commuters_pos_map[float_pos]
 
     def get_commuter_by_id(self, commuter_id: int) -> Commuter:
@@ -69,16 +68,14 @@ def add_commuter(self, agent: Commuter) -> None:
 
     def update_home_counter(
         self,
-        old_home_pos: mesa.space.FloatCoordinate | None,
-        new_home_pos: mesa.space.FloatCoordinate,
+        old_home_pos: FloatCoordinate | None,
+        new_home_pos: FloatCoordinate,
     ) -> None:
         if old_home_pos is not None:
             self.home_counter[old_home_pos] -= 1
         self.home_counter[new_home_pos] += 1
 
-    def move_commuter(
-        self, commuter: Commuter, pos: mesa.space.FloatCoordinate
-    ) -> None:
+    def move_commuter(self, commuter: Commuter, pos: FloatCoordinate) -> None:
         self.__remove_commuter(commuter)
         commuter.geometry = Point(pos)
         self.add_commuter(commuter)