Generalize CellAgent

mesa/experimental/cell_space/cell.py

@@ -7,7 +7,7 @@
 from mesa.experimental.cell_space.cell_collection import CellCollection
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell_agent import CellAgent
+    from mesa.experimental.cell_space.cell_agent import DiscreteSpaceAgent
 
 
 class Cell:
@@ -80,7 +80,7 @@ def disconnect(self, other: Cell) -> None:
         """
         self._connections.remove(other)
 
-    def add_agent(self, agent: CellAgent) -> None:
+    def add_agent(self, agent: DiscreteSpaceAgent[Cell]) -> None:
         """Adds an agent to the cell.
 
         Args:
@@ -96,7 +96,7 @@ def add_agent(self, agent: CellAgent) -> None:
 
         self.agents.append(agent)
 
-    def remove_agent(self, agent: CellAgent) -> None:
+    def remove_agent(self, agent: DiscreteSpaceAgent[Cell]) -> None:
         """Removes an agent from the cell.
 
         Args:

mesa/experimental/cell_space/cell_agent.py

@@ -1,37 +1,83 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
 
-from mesa import Agent, Model
+from mesa.experimental.cell_space.discrete_space import DiscreteSpace
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell import Cell
+    from mesa.experimental.cell_space import Cell
 
+T = TypeVar("T", bound="Cell")
 
-class CellAgent(Agent):
-    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces
 
+class DiscreteSpaceAgent(Protocol[T]):
+    cell: T | None
+    space: DiscreteSpace[T]
+
+    def move_to(self, cell: T) -> None: ...
+
+    def move_relative(self, directions: tuple[int, ...], distance: int = 1): ...
+
+
+class CellAgent:
+    """Cell Agent is an Agent class that adds behavior for moving in discrete spaces
 
     Attributes:
-        unique_id (int): A unique identifier for this agent.
-        model (Model): The model instance to which the agent belongs
-        pos: (Position | None): The position of the agent in the space
-        cell: (Cell | None): the cell which the agent occupies
+        space (DiscreteSpace): the discrete space the agent is in
+        cell (Cell): the cell the agent is in
     """
 
-    def __init__(self, model: Model) -> None:
-        """
-        Create a new agent.
+    def __init__(
+        self,
+        space: DiscreteSpace[Cell],
+        cell: Cell | None = None,
+        *args: tuple[Any],
+        **kwargs: dict[str, Any],
+    ):
+        super().__init__(*args, **kwargs)
+        self.space = space
+        self.cell = cell
+        if cell is not None:
+            cell.add_agent(self)
 
-        Args:
-            unique_id (int): A unique identifier for this agent.
-            model (Model): The model instance in which the agent exists.
-        """
-        super().__init__(model)
-        self.cell: Cell | None = None
+    @property
+    def coordinate(self) -> tuple[int, ...]:
+        return self.cell.coordinate if self.cell else ()
 
-    def move_to(self, cell) -> None:
+    def move_to(self, cell: Cell) -> None:
         if self.cell is not None:
             self.cell.remove_agent(self)
         self.cell = cell
         cell.add_agent(self)
+
+    def move_relative(self, directions: tuple[int, ...], distance: int = 1):
+        new_position = tuple(
+            self.cell.coordinate[i] + directions[i] * distance
+            for i in range(len(directions))
+            if self.cell
+        )
+        new_cell = self.space[new_position]
+        self.move_to(new_cell)
+
+
+class Grid2DMovingAgent:
+    def move(self: DiscreteSpaceAgent[Cell], direction: str, distance: int = 1):
+        match direction:
+            case "N" | "North" | "Up":
+                self.move_relative((-1, 0), distance)
+            case "S" | "South" | "Down":
+                self.move_relative((1, 0), distance)
+            case "E" | "East" | "Right":
+                self.move_relative((0, 1), distance)
+            case "W" | "West" | "Left":
+                self.move_relative((0, -1), distance)
+            case "NE" | "NorthEast" | "UpRight":
+                self.move_relative((-1, 1), distance)
+            case "NW" | "NorthWest" | "UpLeft":
+                self.move_relative((-1, -1), distance)
+            case "SE" | "SouthEast" | "DownRight":
+                self.move_relative((1, 1), distance)
+            case "SW" | "SouthWest" | "DownLeft":
+                self.move_relative((1, -1), distance)
+            case _:
+                raise ValueError(f"Invalid direction: {direction}")

mesa/experimental/cell_space/discrete_space.py

@@ -56,7 +56,7 @@ def __getitem__(self, key):
         return self._cells[key]
 
     @property
-    def empties(self) -> CellCollection:
+    def empties(self) -> CellCollection[T]:
         return self.all_cells.select(lambda cell: cell.is_empty)
 
     def select_random_empty_cell(self) -> T:

mesa/experimental/cell_space/grid.py

@@ -10,7 +10,7 @@
 T = TypeVar("T", bound=Cell)
 
 
-class Grid(DiscreteSpace, Generic[T]):
+class Grid(DiscreteSpace[T], Generic[T]):
     """Base class for all grid classes
 
     Attributes:

mesa/visualization/__init__.py

@@ -1,12 +1,11 @@
 from .components.altair import make_space_altair
 from .components.matplotlib import make_plot_measure, make_space_matplotlib
-from .solara_viz import JupyterViz, SolaraViz, make_text
+from .solara_viz import JupyterViz, SolaraViz
 from .UserParam import Slider
 
 __all__ = [
     "JupyterViz",
     "SolaraViz",
-    "make_text",
     "Slider",
     "make_space_altair",
     "make_space_matplotlib",

tests/test_cell_space.py

@@ -3,6 +3,7 @@
 import pytest
 
 from mesa import Model
+from mesa.agent import Agent
 from mesa.experimental.cell_space import (
     Cell,
     CellAgent,
@@ -15,6 +16,10 @@
 )
 
 
+class TestAgent(Agent, CellAgent):
+    pass
+
+
 def test_orthogonal_grid_neumann():
     width = 10
     height = 10
@@ -408,7 +413,7 @@ def test_empties_space():
 
     model = Model()
     for i in range(8):
-        grid._cells[i].add_agent(CellAgent(model))
+        grid._cells[i].add_agent(TestAgent(model))
 
     cell = grid.select_random_empty_cell()
     assert cell.coordinate in {8, 9}
@@ -432,7 +437,7 @@ def test_cell():
 
     # add_agent
     model = Model()
-    agent = CellAgent(model)
+    agent = TestAgent(model)
 
     cell1.add_agent(agent)
     assert agent in cell1.agents
@@ -445,11 +450,11 @@ def test_cell():
         cell1.remove_agent(agent)
 
     cell1 = Cell((1,), capacity=1, random=random.Random())
-    cell1.add_agent(CellAgent(model))
+    cell1.add_agent(TestAgent(model))
     assert cell1.is_full
 
     with pytest.raises(Exception):
-        cell1.add_agent(CellAgent(model))
+        cell1.add_agent(TestAgent(model))
 
 
 def test_cell_collection():
@@ -475,9 +480,9 @@ def test_cell_collection():
 
     cells = collection.cells
     model = Model()
-    cells[0].add_agent(CellAgent(model))
-    cells[3].add_agent(CellAgent(model))
-    cells[7].add_agent(CellAgent(model))
+    cells[0].add_agent(TestAgent(model))
+    cells[3].add_agent(TestAgent(model))
+    cells[7].add_agent(TestAgent(model))
     agents = collection.agents
     assert len(list(agents)) == 3