Add full support for property layers to cell spaces (#2512)

Author: quaquel

mesa/experimental/cell_space/__init__.py

@@ -20,6 +20,7 @@
     OrthogonalVonNeumannGrid,
 )
 from mesa.experimental.cell_space.network import Network
+from mesa.experimental.cell_space.property_layer import PropertyLayer
 from mesa.experimental.cell_space.voronoi import VoronoiGrid
 
 __all__ = [
@@ -34,5 +35,6 @@
     "Network",
     "OrthogonalMooreGrid",
     "OrthogonalVonNeumannGrid",
+    "PropertyLayer",
     "VoronoiGrid",
 ]

mesa/experimental/cell_space/cell.py

@@ -2,14 +2,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 from functools import cache, cached_property
 from random import Random
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from mesa.experimental.cell_space.cell_agent import CellAgent
 from mesa.experimental.cell_space.cell_collection import CellCollection
-from mesa.space import PropertyLayer
 
 if TYPE_CHECKING:
     from mesa.agent import Agent
@@ -24,14 +22,12 @@ class Cell:
         coordinate (Tuple[int, int]) : the position of the cell in the discrete space
         agents (List[Agent]): the agents occupying the cell
         capacity (int): the maximum number of agents that can simultaneously occupy the cell
-        properties (dict[str, Any]): the properties of the cell
         random (Random): the random number generator
 
     """
 
     __slots__ = [
         "__dict__",
-        "_mesa_property_layers",
         "agents",
         "capacity",
         "connections",
@@ -40,15 +36,6 @@ class Cell:
         "random",
     ]
 
-    # def __new__(cls,
-    #     coordinate: tuple[int, ...],
-    #     capacity: float | None = None,
-    #     random: Random | None = None,):
-    #     if capacity != 1:
-    #         return object.__new__(cls)
-    #     else:
-    #         return object.__new__(SingleAgentCell)
-
     def __init__(
         self,
         coordinate: Coordinate,
@@ -70,9 +57,10 @@ def __init__(
             Agent
         ] = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
         self.capacity: int | None = capacity
-        self.properties: dict[Coordinate, object] = {}
+        self.properties: dict[
+            Coordinate, object
+        ] = {}  # fixme still used by voronoi mesh
         self.random = random
-        self._mesa_property_layers: dict[str, PropertyLayer] = {}
 
     def connect(self, other: Cell, key: Coordinate | None = None) -> None:
         """Connects this cell to another cell.
@@ -105,6 +93,7 @@ def add_agent(self, agent: CellAgent) -> None:
 
         """
         n = len(self.agents)
+        self.empty = False
 
         if self.capacity and n >= self.capacity:
             raise Exception(
@@ -121,6 +110,7 @@ def remove_agent(self, agent: CellAgent) -> None:
 
         """
         self.agents.remove(agent)
+        self.empty = self.is_empty
 
     @property
     def is_empty(self) -> bool:
@@ -195,23 +185,6 @@ def _neighborhood(
                 neighborhood.pop(self, None)
             return neighborhood
 
-    # PropertyLayer methods
-    def get_property(self, property_name: str) -> Any:
-        """Get the value of a property."""
-        return self._mesa_property_layers[property_name].data[self.coordinate]
-
-    def set_property(self, property_name: str, value: Any):
-        """Set the value of a property."""
-        self._mesa_property_layers[property_name].set_cell(self.coordinate, value)
-
-    def modify_property(
-        self, property_name: str, operation: Callable, value: Any = None
-    ):
-        """Modify the value of a property."""
-        self._mesa_property_layers[property_name].modify_cell(
-            self.coordinate, operation, value
-        )
-
     def __getstate__(self):
         """Return state of the Cell with connections set to empty."""
         # fixme, once we shift to 3.11, replace this with super. __getstate__

mesa/experimental/cell_space/discrete_space.py

@@ -3,15 +3,13 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Callable
 from functools import cached_property
 from random import Random
-from typing import Any, Generic, TypeVar
+from typing import Generic, TypeVar
 
 from mesa.agent import AgentSet
 from mesa.experimental.cell_space.cell import Cell
 from mesa.experimental.cell_space.cell_collection import CellCollection
-from mesa.space import PropertyLayer
 
 T = TypeVar("T", bound=Cell)
 
@@ -61,8 +59,6 @@ def __init__(
         self.cell_klass = cell_klass
 
         self._empties: dict[tuple[int, ...], None] = {}
-        self._empties_initialized = False
-        self.property_layers: dict[str, PropertyLayer] = {}
 
     @property
     def cutoff_empties(self):  # noqa
@@ -98,64 +94,6 @@ def select_random_empty_cell(self) -> T:
         """Select random empty cell."""
         return self.random.choice(list(self.empties))
 
-    # PropertyLayer methods
-    def add_property_layer(
-        self, property_layer: PropertyLayer, add_to_cells: bool = True
-    ):
-        """Add a property layer to the grid.
-
-        Args:
-            property_layer: the property layer to add
-            add_to_cells: whether to add the property layer to all cells (default: True)
-        """
-        if property_layer.name in self.property_layers:
-            raise ValueError(f"Property layer {property_layer.name} already exists.")
-        self.property_layers[property_layer.name] = property_layer
-        if add_to_cells:
-            for cell in self._cells.values():
-                cell._mesa_property_layers[property_layer.name] = property_layer
-
-    def remove_property_layer(self, property_name: str, remove_from_cells: bool = True):
-        """Remove a property layer from the grid.
-
-        Args:
-            property_name: the name of the property layer to remove
-            remove_from_cells: whether to remove the property layer from all cells (default: True)
-        """
-        del self.property_layers[property_name]
-        if remove_from_cells:
-            for cell in self._cells.values():
-                del cell._mesa_property_layers[property_name]
-
-    def set_property(
-        self, property_name: str, value, condition: Callable[[T], bool] | None = None
-    ):
-        """Set the value of a property for all cells in the grid.
-
-        Args:
-            property_name: the name of the property to set
-            value: the value to set
-            condition: a function that takes a cell and returns a boolean
-        """
-        self.property_layers[property_name].set_cells(value, condition)
-
-    def modify_properties(
-        self,
-        property_name: str,
-        operation: Callable,
-        value: Any = None,
-        condition: Callable[[T], bool] | None = None,
-    ):
-        """Modify the values of a specific property for all cells in the grid.
-
-        Args:
-            property_name: the name of the property to modify
-            operation: the operation to perform
-            value: the value to use in the operation
-            condition: a function that takes a cell and returns a boolean (used to filter cells)
-        """
-        self.property_layers[property_name].modify_cells(operation, value, condition)
-
     def __setstate__(self, state):
         """Set the state of the discrete space and rebuild the connections."""
         self.__dict__ = state

mesa/experimental/cell_space/grid.py

@@ -2,17 +2,50 @@
 
 from __future__ import annotations
 
+import copyreg
 from collections.abc import Sequence
 from itertools import product
 from random import Random
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 from mesa.experimental.cell_space import Cell, DiscreteSpace
+from mesa.experimental.cell_space.property_layer import (
+    HasPropertyLayers,
+    PropertyDescriptor,
+)
 
 T = TypeVar("T", bound=Cell)
 
 
-class Grid(DiscreteSpace[T], Generic[T]):
+def pickle_gridcell(obj):
+    """Helper function for pickling GridCell instances."""
+    # we have the base class and the state via __getstate__
+    args = obj.__class__.__bases__[0], obj.__getstate__()
+    return unpickle_gridcell, args
+
+
+def unpickle_gridcell(parent, fields):
+    """Helper function for unpickling GridCell instances."""
+    # since the class is dynamically created, we recreate it here
+    cell_klass = type(
+        "GridCell",
+        (parent,),
+        {"_mesa_properties": set()},
+    )
+    instance = cell_klass(
+        (0, 0)
+    )  # we use a default coordinate and overwrite it with the correct value next
+
+    # __gestate__ returns a tuple with dict and slots, but slots contains the dict so we can just use the
+    # second item only
+    for k, v in fields[1].items():
+        if k != "__dict__":
+            setattr(instance, k, v)
+
+    return instance
+
+
+class Grid(DiscreteSpace[T], Generic[T], HasPropertyLayers):
     """Base class for all grid classes.
 
     Attributes:
@@ -60,14 +93,23 @@ def __init__(
         self._try_random = True
         self._ndims = len(dimensions)
         self._validate_parameters()
+        self.cell_klass = type(
+            "GridCell",
+            (self.cell_klass,),
+            {"_mesa_properties": set()},
+        )
+
+        # we register the pickle_gridcell helper function
+        copyreg.pickle(self.cell_klass, pickle_gridcell)
 
         coordinates = product(*(range(dim) for dim in self.dimensions))
 
         self._cells = {
-            coord: cell_klass(coord, capacity, random=self.random)
+            coord: self.cell_klass(coord, capacity, random=self.random)
             for coord in coordinates
         }
         self._connect_cells()
+        self.create_property_layer("empty", default_value=True, dtype=bool)
 
     def _connect_cells(self) -> None:
         if self._ndims == 2:
@@ -126,6 +168,23 @@ def _connect_single_cell_2d(self, cell: T, offsets: list[tuple[int, int]]) -> No
             if 0 <= ni < height and 0 <= nj < width:
                 cell.connect(self._cells[ni, nj], (di, dj))
 
+    def __getstate__(self) -> dict[str, Any]:
+        """Custom __getstate__ for handling dynamic GridCell class and PropertyDescriptors."""
+        state = super().__getstate__()
+        state = {k: v for k, v in state.items() if k != "cell_klass"}
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Custom __setstate__ for handling dynamic GridCell class and PropertyDescriptors."""
+        self.__dict__ = state
+        self._connect_cells()  # using super fails for this for some reason, so we repeat ourselves
+
+        self.cell_klass = type(
+            self._cells[(0, 0)]
+        )  # the __reduce__ function handles this for us nicely
+        for layer in self._mesa_property_layers.values():
+            setattr(self.cell_klass, layer.name, PropertyDescriptor(layer))
+
 
 class OrthogonalMooreGrid(Grid[T]):
     """Grid where cells are connected to their 8 neighbors.