add some agents

Author: Corvince

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 CellHolder
 
 
 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: CellHolder) -> 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: CellHolder) -> None:
         """Removes an agent from the cell.
 
         Args:

mesa/experimental/cell_space/cell_agent.py

@@ -1,37 +1,59 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
-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, Grid
 
 
-class CellAgent(Agent):
-    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces
+@runtime_checkable
+class CellHolder(Protocol):
+    cell: Cell | None
 
 
-    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
-    """
+class CellAgent:
+    cell: Cell | None
+    space: DiscreteSpace[Cell]
+    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces"""
 
-    def __init__(self, model: Model) -> None:
-        """
-        Create a new agent.
-
-        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
-
-    def move_to(self, cell) -> None:
+    def move_to(self: CellHolder, 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(CellAgent):
+    grid: Grid[Cell]
+
+    def move(self, 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

@@ -43,7 +43,8 @@ def __init__(
     def cutoff_empties(self):
         return 7.953 * len(self._cells) ** 0.384
 
-    def _connect_single_cell(self, cell: T): ...
+    def _connect_single_cell(self, cell: T):
+        ...
 
     @cached_property
     def all_cells(self):
@@ -56,7 +57,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:
@@ -51,9 +51,11 @@ def _connect_cells(self) -> None:
         else:
             self._connect_cells_nd()
 
-    def _connect_cells_2d(self) -> None: ...
+    def _connect_cells_2d(self) -> None:
+        ...
 
-    def _connect_cells_nd(self) -> None: ...
+    def _connect_cells_nd(self) -> None:
+        ...
 
     def _validate_parameters(self):
         if not all(isinstance(dim, int) and dim > 0 for dim in self.dimensions):

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