Make cell connections public and named (#2296)

make connections public and use a dict with relative coordinates as key
---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Corvince

mesa/experimental/cell_space/cell.py

@@ -9,8 +9,11 @@
 from mesa.experimental.cell_space.cell_collection import CellCollection
 
 if TYPE_CHECKING:
+    from mesa.agent import Agent
     from mesa.experimental.cell_space.cell_agent import CellAgent
 
+Coordinate = tuple[int, ...]
+
 
 class Cell:
     """The cell represents a position in a discrete space.
@@ -26,7 +29,7 @@ class Cell:
 
     __slots__ = [
         "coordinate",
-        "_connections",
+        "connections",
         "agents",
         "capacity",
         "properties",
@@ -44,7 +47,7 @@ class Cell:
 
     def __init__(
         self,
-        coordinate: tuple[int, ...],
+        coordinate: Coordinate,
         capacity: float | None = None,
         random: Random | None = None,
     ) -> None:
@@ -58,20 +61,25 @@ def __init__(
         """
         super().__init__()
         self.coordinate = coordinate
-        self._connections: list[Cell] = []  # TODO: change to CellCollection?
-        self.agents = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
+        self.connections: dict[Coordinate, Cell] = {}
+        self.agents: list[
+            Agent
+        ] = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
         self.capacity = capacity
-        self.properties: dict[str, object] = {}
+        self.properties: dict[Coordinate, object] = {}
         self.random = random
 
-    def connect(self, other: Cell) -> None:
+    def connect(self, other: Cell, key: Coordinate | None = None) -> None:
         """Connects this cell to another cell.
 
         Args:
             other (Cell): other cell to connect to
+            key (Tuple[int, ...]): key for the connection. Should resemble a relative coordinate
 
         """
-        self._connections.append(other)
+        if key is None:
+            key = other.coordinate
+        self.connections[key] = other
 
     def disconnect(self, other: Cell) -> None:
         """Disconnects this cell from another cell.
@@ -80,7 +88,9 @@ def disconnect(self, other: Cell) -> None:
             other (Cell): other cell to remove from connections
 
         """
-        self._connections.remove(other)
+        keys_to_remove = [k for k, v in self.connections.items() if v == other]
+        for key in keys_to_remove:
+            del self.connections[key]
 
     def add_agent(self, agent: CellAgent) -> None:
         """Adds an agent to the cell.
@@ -123,30 +133,34 @@ def __repr__(self):  # noqa
 
     # FIXME: Revisit caching strategy on methods
     @cache  # noqa: B019
-    def neighborhood(self, radius=1, include_center=False):
+    def neighborhood(self, radius: int = 1, include_center: bool = False):
         """Returns a list of all neighboring cells."""
-        return CellCollection(
+        return CellCollection[Cell](
             self._neighborhood(radius=radius, include_center=include_center),
             random=self.random,
         )
 
     # FIXME: Revisit caching strategy on methods
     @cache  # noqa: B019
-    def _neighborhood(self, radius=1, include_center=False):
+    def _neighborhood(
+        self, radius: int = 1, include_center: bool = False
+    ) -> dict[Cell, list[Agent]]:
         # if radius == 0:
         #     return {self: self.agents}
         if radius < 1:
             raise ValueError("radius must be larger than one")
         if radius == 1:
-            neighborhood = {neighbor: neighbor.agents for neighbor in self._connections}
+            neighborhood = {
+                neighbor: neighbor.agents for neighbor in self.connections.values()
+            }
             if not include_center:
                 return neighborhood
             else:
                 neighborhood[self] = self.agents
                 return neighborhood
         else:
-            neighborhood = {}
-            for neighbor in self._connections:
+            neighborhood: dict[Cell, list[Agent]] = {}
+            for neighbor in self.connections.values():
                 neighborhood.update(
                     neighbor._neighborhood(radius - 1, include_center=True)
                 )

mesa/experimental/cell_space/discrete_space.py

@@ -62,11 +62,11 @@ def all_cells(self):
     def __iter__(self):  # noqa
         return iter(self._cells.values())
 
-    def __getitem__(self, key):  # noqa
+    def __getitem__(self, key: tuple[int, ...]) -> T:  # noqa: D105
         return self._cells[key]
 
     @property
-    def empties(self) -> CellCollection:
+    def empties(self) -> CellCollection[T]:
         """Return all empty in spaces."""
         return self.all_cells.select(lambda cell: cell.is_empty)
 

mesa/experimental/cell_space/grid.py

@@ -12,7 +12,7 @@
 T = TypeVar("T", bound=Cell)
 
 
-class Grid(DiscreteSpace, Generic[T]):
+class Grid(DiscreteSpace[T], Generic[T]):
     """Base class for all grid classes.
 
     Attributes:
@@ -100,7 +100,7 @@ def _connect_single_cell_nd(self, cell: T, offsets: list[tuple[int, ...]]) -> No
             if self.torus:
                 n_coord = tuple(nc % d for nc, d in zip(n_coord, self.dimensions))
             if all(0 <= nc < d for nc, d in zip(n_coord, self.dimensions)):
-                cell.connect(self._cells[n_coord])
+                cell.connect(self._cells[n_coord], d_coord)
 
     def _connect_single_cell_2d(self, cell: T, offsets: list[tuple[int, int]]) -> None:
         i, j = cell.coordinate
@@ -111,7 +111,7 @@ def _connect_single_cell_2d(self, cell: T, offsets: list[tuple[int, int]]) -> No
             if self.torus:
                 ni, nj = ni % height, nj % width
             if 0 <= ni < height and 0 <= nj < width:
-                cell.connect(self._cells[ni, nj])
+                cell.connect(self._cells[ni, nj], (di, dj))
 
 
 class OrthogonalMooreGrid(Grid[T]):
@@ -133,7 +133,6 @@ def _connect_cells_2d(self) -> None:
             ( 1, -1), ( 1, 0), ( 1, 1),
         ]
         # fmt: on
-        height, width = self.dimensions
 
         for cell in self.all_cells:
             self._connect_single_cell_2d(cell, offsets)
@@ -165,13 +164,12 @@ def _connect_cells_2d(self) -> None:
                     ( 1, 0),
         ]
         # fmt: on
-        height, width = self.dimensions
 
         for cell in self.all_cells:
             self._connect_single_cell_2d(cell, offsets)
 
     def _connect_cells_nd(self) -> None:
-        offsets = []
+        offsets: list[tuple[int, ...]] = []
         dimensions = len(self.dimensions)
         for dim in range(dimensions):
             for delta in [

mesa/experimental/cell_space/network.py

@@ -7,7 +7,7 @@
 from mesa.experimental.cell_space.discrete_space import DiscreteSpace
 
 
-class Network(DiscreteSpace):
+class Network(DiscreteSpace[Cell]):
     """A networked discrete space."""
 
     def __init__(
@@ -37,6 +37,6 @@ def __init__(
         for cell in self.all_cells:
             self._connect_single_cell(cell)
 
-    def _connect_single_cell(self, cell):
+    def _connect_single_cell(self, cell: Cell):
         for node_id in self.G.neighbors(cell.coordinate):
-            cell.connect(self._cells[node_id])
+            cell.connect(self._cells[node_id], node_id)

mesa/experimental/cell_space/voronoi.py

@@ -216,8 +216,8 @@ def _connect_cells(self) -> None:
 
         for point in self.triangulation.export_triangles():
             for i, j in combinations(point, 2):
-                self._cells[i].connect(self._cells[j])
-                self._cells[j].connect(self._cells[i])
+                self._cells[i].connect(self._cells[j], (i, j))
+                self._cells[j].connect(self._cells[i], (j, i))
 
     def _validate_parameters(self) -> None:
         if self.capacity is not None and not isinstance(self.capacity, float | int):
@@ -233,9 +233,10 @@ def _validate_parameters(self) -> None:
 
     def _get_voronoi_regions(self) -> tuple:
         if self.voronoi_coordinates is None or self.regions is None:
-            self.voronoi_coordinates, self.regions = (
-                self.triangulation.export_voronoi_regions()
-            )
+            (
+                self.voronoi_coordinates,
+                self.regions,
+            ) = self.triangulation.export_voronoi_regions()
         return self.voronoi_coordinates, self.regions
 
     @staticmethod