diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py index b7c002bb07a..8c7fb9bb632 100644 --- a/mesa/experimental/cell_space/cell.py +++ b/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 + ) diff --git a/mesa/experimental/cell_space/discrete_space.py b/mesa/experimental/cell_space/discrete_space.py index 0d3af169e98..1cfaedeb167 100644 --- a/mesa/experimental/cell_space/discrete_space.py +++ b/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) diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py index 0dd3855bc6b..a8e4abad336 100644 --- a/tests/test_cell_space.py +++ b/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())