experimental: Integrate PropertyLayers into cell space

mesa/experimental/cell_space/cell.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from functools import cache, cached_property
 from random import Random
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from mesa.experimental.cell_space.cell_collection import CellCollection
+from mesa.space import PropertyLayer
 
 if TYPE_CHECKING:
     from mesa.agent import Agent
@@ -34,6 +36,7 @@ class Cell:
         "capacity",
         "properties",
         "random",
+        "_mesa_property_layers",
         "__dict__",
     ]
 
@@ -69,6 +72,7 @@ def __init__(
         self.capacity: int = capacity
         self.properties: dict[Coordinate, object] = {}
         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.
@@ -190,3 +194,20 @@ def _neighborhood(
             if not include_center:
                 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
+        )

mesa/experimental/cell_space/discrete_space.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from functools import cached_property
 from random import Random
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 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)
 
@@ -21,7 +23,7 @@ class DiscreteSpace(Generic[T]):
         random (Random): The random number generator
         cell_klass (Type) : the type of cell class
         empties (CellCollection) : collecction of all cells that are empty
-
+        property_layers (dict[str, PropertyLayer]): the property layers of the discrete space
     """
 
     def __init__(
@@ -47,6 +49,7 @@ def __init__(
 
         self._empties: dict[tuple[int, ...], None] = {}
         self._empties_initialized = False
+        self.property_layers: dict[str, PropertyLayer] = {}
 
     @property
     def cutoff_empties(self):  # noqa
@@ -73,3 +76,61 @@ def empties(self) -> CellCollection[T]:
     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)

tests/test_cell_space.py

@@ -2,6 +2,7 @@
 
 import random
 
+import numpy as np
 import pytest
 
 from mesa import Model
@@ -15,6 +16,7 @@
     OrthogonalVonNeumannGrid,
     VoronoiGrid,
 )
+from mesa.space import PropertyLayer
 
 
 def test_orthogonal_grid_neumann():
@@ -526,6 +528,92 @@ def test_cell_collection():
     assert len(cells) == len(collection)
 
 
+### PropertyLayer tests
+def test_property_layer_integration():
+    """Test integration of PropertyLayer with DiscrateSpace and Cell."""
+    width, height = 10, 10
+    grid = OrthogonalMooreGrid((width, height), torus=False)
+
+    # Test adding a PropertyLayer to the grid
+    elevation = PropertyLayer("elevation", width, height, default_value=0)
+    grid.add_property_layer(elevation)
+    assert "elevation" in grid.property_layers
+    assert len(grid.property_layers) == 1
+
+    # Test accessing PropertyLayer from a cell
+    cell = grid._cells[(0, 0)]
+    assert "elevation" in cell._mesa_property_layers
+    assert cell.get_property("elevation") == 0
+
+    # Test setting property value for a cell
+    cell.set_property("elevation", 100)
+    assert cell.get_property("elevation") == 100
+
+    # Test modifying property value for a cell
+    cell.modify_property("elevation", lambda x: x + 50)
+    assert cell.get_property("elevation") == 150
+
+    cell.modify_property("elevation", np.add, 50)
+    assert cell.get_property("elevation") == 200
+
+    # Test modifying PropertyLayer values
+    grid.set_property("elevation", 100, condition=lambda value: value == 200)
+    assert cell.get_property("elevation") == 100
+
+    # Test modifying PropertyLayer using numpy operations
+    grid.modify_properties("elevation", np.add, 50)
+    assert cell.get_property("elevation") == 150
+
+    # Test removing a PropertyLayer
+    grid.remove_property_layer("elevation")
+    assert "elevation" not in grid.property_layers
+    assert "elevation" not in cell._mesa_property_layers
+
+
+def test_multiple_property_layers():
+    """Test initialization of DiscrateSpace with PropertyLayers."""
+    width, height = 5, 5
+    elevation = PropertyLayer("elevation", width, height, default_value=0)
+    temperature = PropertyLayer("temperature", width, height, default_value=20)
+
+    # Test initialization with a single PropertyLayer
+    grid1 = OrthogonalMooreGrid((width, height), torus=False)
+    grid1.add_property_layer(elevation)
+    assert "elevation" in grid1.property_layers
+    assert len(grid1.property_layers) == 1
+
+    # Test initialization with multiple PropertyLayers
+    grid2 = OrthogonalMooreGrid((width, height), torus=False)
+    grid2.add_property_layer(temperature, add_to_cells=False)
+    grid2.add_property_layer(elevation, add_to_cells=True)
+
+    assert "temperature" in grid2.property_layers
+    assert "elevation" in grid2.property_layers
+    assert len(grid2.property_layers) == 2
+
+    # Modify properties
+    grid2.modify_properties("elevation", lambda x: x + 10)
+    grid2.modify_properties("temperature", lambda x: x + 5)
+
+    for cell in grid2.all_cells:
+        assert cell.get_property("elevation") == 10
+        # Assert error temperature, since it was not added to cells
+        with pytest.raises(KeyError):
+            cell.get_property("temperature")
+
+
+def test_property_layer_errors():
+    """Test error handling for PropertyLayers."""
+    width, height = 5, 5
+    grid = OrthogonalMooreGrid((width, height), torus=False)
+    elevation = PropertyLayer("elevation", width, height, default_value=0)
+
+    # Test adding a PropertyLayer with an existing name
+    grid.add_property_layer(elevation)
+    with pytest.raises(ValueError, match="Property layer elevation already exists."):
+        grid.add_property_layer(elevation)
+
+
 def test_cell_agent():  # noqa: D103
     cell1 = Cell((1,), capacity=None, random=random.Random())
     cell2 = Cell((2,), capacity=None, random=random.Random())