diff --git a/mesa/experimental/cell_space/__init__.py b/mesa/experimental/cell_space/__init__.py index 5c790075909..22dddcaf30b 100644 --- a/mesa/experimental/cell_space/__init__.py +++ b/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", ] diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py index 3020dad885d..e101b017071 100644 --- a/mesa/experimental/cell_space/cell.py +++ b/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__ diff --git a/mesa/experimental/cell_space/discrete_space.py b/mesa/experimental/cell_space/discrete_space.py index a017d06bda8..88a6dc5476e 100644 --- a/mesa/experimental/cell_space/discrete_space.py +++ b/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 diff --git a/mesa/experimental/cell_space/grid.py b/mesa/experimental/cell_space/grid.py index 299572faf5e..7f1d5751f64 100644 --- a/mesa/experimental/cell_space/grid.py +++ b/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. diff --git a/mesa/experimental/cell_space/property_layer.py b/mesa/experimental/cell_space/property_layer.py new file mode 100644 index 00000000000..5e2dbf66c94 --- /dev/null +++ b/mesa/experimental/cell_space/property_layer.py @@ -0,0 +1,429 @@ +"""This module provides functionality for working with property layers in grids.""" + +import warnings +from collections.abc import Callable, Sequence +from typing import Any, TypeVar + +import numpy as np + +from mesa.experimental.cell_space import Cell + +Coordinate = Sequence[int] +T = TypeVar("T", bound=Cell) + + +class PropertyLayer: + """A class representing a layer of properties in a two-dimensional grid. + + Each cell in the grid can store a value of a specified data type. + + Attributes: + name: The name of the property layer. + dimensions: The width of the grid (number of columns). + data: A NumPy array representing the grid data. + + """ + + # Fixme + # can't we simplify this a lot + # in essence, this is just a numpy array with a name and fixed dimensions + # all other functionality seems redundant to me? + + @property + def data(self): # noqa: D102 + return self._mesa_data + + @data.setter + def data(self, value): + self.set_cells(value) + + propertylayer_experimental_warning_given = False + + def __init__( + self, name: str, dimensions: Sequence[int], default_value=0.0, dtype=float + ): + """Initializes a new PropertyLayer instance. + + Args: + name: The name of the property layer. + dimensions: the dimensions of the property layer. + default_value: The default value to initialize each cell in the grid. Should ideally + be of the same type as specified by the dtype parameter. + dtype (data-type, optional): The desired data-type for the grid's elements. Default is float. + + Notes: + A UserWarning is raised if the default_value is not of a type compatible with dtype. + The dtype parameter can accept both Python data types (like bool, int or float) and NumPy data types + (like np.int64 or np.float64). + """ + self.name = name + self.dimensions = dimensions + + # Check if the dtype is suitable for the data + if not isinstance(default_value, dtype): + warnings.warn( + f"Default value {default_value} ({type(default_value).__name__}) might not be best suitable with dtype={dtype.__name__}.", + UserWarning, + stacklevel=2, + ) + + # fixme why not initialize with empty? + self._mesa_data = np.full(self.dimensions, default_value, dtype=dtype) + + if not self.__class__.propertylayer_experimental_warning_given: + warnings.warn( + "The property layer functionality and associated classes are experimental. It may be changed or removed in any and all future releases, including patch releases.\n" + "We would love to hear what you think about this new feature. If you have any thoughts, share them with us here: https://github.com/projectmesa/mesa/discussions/1932", + FutureWarning, + stacklevel=2, + ) + self.__class__.propertylayer_experimental_warning_given = True + + @classmethod + def from_data(cls, name: str, data: np.ndarray): + """Create a property layer from a NumPy array. + + Args: + name: The name of the property layer. + data: A NumPy array representing the grid data. + + """ + layer = cls( + name, + data.shape, + default_value=data[*[0 for _ in range(len(data.shape))]], + dtype=data.dtype.type, + ) + layer.set_cells(data) + return layer + + def set_cells(self, value, condition: Callable | None = None): + """Perform a batch update either on the entire grid or conditionally, in-place. + + Args: + value: The value to be used for the update. + condition: (Optional) A callable that returns a boolean array when applied to the data. + """ + if condition is None: + np.copyto(self._mesa_data, value) # In-place update + else: + vectorized_condition = np.vectorize(condition) + condition_result = vectorized_condition(self._mesa_data) + np.copyto(self._mesa_data, value, where=condition_result) + + def modify_cells( + self, + operation: Callable, + value=None, + condition: Callable | None = None, + ): + """Modify cells using an operation, which can be a lambda function or a NumPy ufunc. + + If a NumPy ufunc is used, an additional value should be provided. + + Args: + operation: A function to apply. Can be a lambda function or a NumPy ufunc. + value: The value to be used if the operation is a NumPy ufunc. Ignored for lambda functions. + condition: (Optional) A callable that returns a boolean array when applied to the data. + """ + condition_array = np.ones_like( + self._mesa_data, dtype=bool + ) # Default condition (all cells) + if condition is not None: + vectorized_condition = np.vectorize(condition) + condition_array = vectorized_condition(self._mesa_data) + + # Check if the operation is a lambda function or a NumPy ufunc + if isinstance(operation, np.ufunc): + if ufunc_requires_additional_input(operation): + if value is None: + raise ValueError("This ufunc requires an additional input value.") + modified_data = operation(self._mesa_data, value) + else: + modified_data = operation(self._mesa_data) + else: + # Vectorize non-ufunc operations + vectorized_operation = np.vectorize(operation) + modified_data = vectorized_operation(self._mesa_data) + + self._mesa_data = np.where(condition_array, modified_data, self._mesa_data) + + def select_cells(self, condition: Callable, return_list=True): + """Find cells that meet a specified condition using NumPy's boolean indexing, in-place. + + Args: + condition: A callable that returns a boolean array when applied to the data. + return_list: (Optional) If True, return a list of (x, y) tuples. Otherwise, return a boolean array. + + Returns: + A list of (x, y) tuples or a boolean array. + """ + # fixme: consider splitting into two separate functions + # select_cells_boolean + # select_cells_index + + condition_array = condition(self._mesa_data) + if return_list: + return list(zip(*np.where(condition_array))) + else: + return condition_array + + def aggregate(self, operation: Callable): + """Perform an aggregate operation (e.g., sum, mean) on a property across all cells. + + Args: + operation: A function to apply. Can be a lambda function or a NumPy ufunc. + """ + return operation(self._mesa_data) + + +class HasPropertyLayers: + """Mixin-like class to add property layer functionality to Grids. + + Property layers can be added to a grid using create_property_layer or add_property_layer. Once created, property + layers can be accessed as attributes if the name used for the layer is a valid python identifier. + + """ + + # fixme is there a way to indicate that a mixin only works with specific classes? + def __init__(self, *args, **kwargs): + """Initialize a HasPropertyLayers instance.""" + super().__init__(*args, **kwargs) + self._mesa_property_layers = {} + + def create_property_layer( + self, + name: str, + default_value=0.0, + dtype=float, + ): + """Add a property layer to the grid. + + Args: + name: The name of the property layer. + default_value: The default value of the property layer. + dtype: The data type of the property layer. + + Returns: + Property layer instance. + + """ + layer = PropertyLayer( + name, self.dimensions, default_value=default_value, dtype=dtype + ) + self.add_property_layer(layer) + return layer + + def add_property_layer(self, layer: PropertyLayer): + """Add a predefined property layer to the grid. + + Args: + layer: The property layer to add. + + Raises: + ValueError: If the dimensions of the layer and the grid are not the same. + + """ + if layer.dimensions != self.dimensions: + raise ValueError( + "Dimensions of property layer do not match the dimensions of the grid" + ) + if layer.name in self._mesa_property_layers: + raise ValueError(f"Property layer {layer.name} already exists.") + if ( + layer.name in self.cell_klass.__slots__ + or layer.name in self.cell_klass.__dict__ + ): + raise ValueError( + f"Property layer {layer.name} clashes with existing attribute in {self.cell_klass.__name__}" + ) + + self._mesa_property_layers[layer.name] = layer + setattr(self.cell_klass, layer.name, PropertyDescriptor(layer)) + self.cell_klass._mesa_properties.add(layer.name) + + def remove_property_layer(self, property_name: str): + """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._mesa_property_layers[property_name] + delattr(self.cell_klass, property_name) + self.cell_klass._mesa_properties.remove(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._mesa_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._mesa_property_layers[property_name].modify_cells( + operation, value, condition + ) + + def get_neighborhood_mask( + self, coordinate: Coordinate, include_center: bool = True, radius: int = 1 + ) -> np.ndarray: + """Generate a boolean mask representing the neighborhood. + + Args: + coordinate: Center of the neighborhood. + include_center: Include the central cell in the neighborhood. + radius: The radius of the neighborhood. + + Returns: + np.ndarray: A boolean mask representing the neighborhood. + """ + cell = self._cells[coordinate] + neighborhood = cell.get_neighborhood( + include_center=include_center, radius=radius + ) + mask = np.zeros(self.dimensions, dtype=bool) + + # Convert the neighborhood list to a NumPy array and use advanced indexing + coords = np.array([c.coordinate for c in neighborhood]) + indices = [coords[:, i] for i in range(coords.shape[1])] + mask[*indices] = True + return mask + + def select_cells( + self, + conditions: dict | None = None, + extreme_values: dict | None = None, + masks: np.ndarray | list[np.ndarray] = None, + only_empty: bool = False, + return_list: bool = True, + ) -> list[Coordinate] | np.ndarray: + """Select cells based on property conditions, extreme values, and/or masks, with an option to only select empty cells. + + Args: + conditions (dict): A dictionary where keys are property names and values are callables that return a boolean when applied. + extreme_values (dict): A dictionary where keys are property names and values are either 'highest' or 'lowest'. + masks (np.ndarray | list[np.ndarray], optional): A mask or list of masks to restrict the selection. + only_empty (bool, optional): If True, only select cells that are empty. Default is False. + return_list (bool, optional): If True, return a list of coordinates, otherwise return a mask. + + Returns: + Union[list[Coordinate], np.ndarray]: Coordinates where conditions are satisfied or the combined mask. + """ + # fixme: consider splitting into two separate functions + # select_cells_boolean + # select_cells_index + # also we might want to change the naming to avoid classes with PropertyLayer + + # Initialize the combined mask + combined_mask = np.ones(self.dimensions, dtype=bool) + + # Apply the masks + if masks is not None: + if isinstance(masks, list): + for mask in masks: + combined_mask = np.logical_and(combined_mask, mask) + else: + combined_mask = np.logical_and(combined_mask, masks) + + # Apply the empty mask if only_empty is True + if only_empty: + combined_mask = np.logical_and( + combined_mask, self._mesa_property_layers["empty"] + ) + + # Apply conditions + if conditions: + for prop_name, condition in conditions.items(): + prop_layer = self._mesa_property_layers[prop_name].data + prop_mask = condition(prop_layer) + combined_mask = np.logical_and(combined_mask, prop_mask) + + # Apply extreme values + if extreme_values: + for property_name, mode in extreme_values.items(): + prop_values = self._mesa_property_layers[property_name].data + + # Create a masked array using the combined_mask + masked_values = np.ma.masked_array(prop_values, mask=~combined_mask) + + if mode == "highest": + target_value = masked_values.max() + elif mode == "lowest": + target_value = masked_values.min() + else: + raise ValueError( + f"Invalid mode {mode}. Choose from 'highest' or 'lowest'." + ) + + extreme_value_mask = prop_values == target_value + combined_mask = np.logical_and(combined_mask, extreme_value_mask) + + # Generate output + if return_list: + selected_cells = list(zip(*np.where(combined_mask))) + return selected_cells + else: + return combined_mask + + def __getattr__(self, name: str) -> Any: # noqa: D105 + try: + return self._mesa_property_layers[name] + except KeyError as e: + raise AttributeError( + f"'{type(self).__name__}' object has no property layer called '{name}'" + ) from e + + def __setattr__(self, key, value): # noqa: D105 + # fixme + # this might be done more elegantly, the main problem is that _mesa_property_layers must already be defined to avoid infinite recursion errors from happening + # also, this protection only works if the attribute is added after the layer, not the other way around + try: + layers = self.__dict__["_mesa_property_layers"] + except KeyError: + super().__setattr__(key, value) + else: + if key in layers: + raise AttributeError( + f"'{type(self).__name__}' object already has a property layer with name '{key}'" + ) + else: + super().__setattr__(key, value) + + +class PropertyDescriptor: + """Descriptor for giving cells attribute like access to values defined in property layers.""" + + def __init__(self, property_layer: PropertyLayer): # noqa: D107 + self.layer: PropertyLayer = property_layer + + def __get__(self, instance: Cell, owner): # noqa: D105 + return self.layer.data[instance.coordinate] + + def __set__(self, instance: Cell, value): # noqa: D105 + self.layer.data[instance.coordinate] = value + + +def ufunc_requires_additional_input(ufunc): # noqa: D103 + # NumPy ufuncs have a 'nargs' attribute indicating the number of input arguments + # For binary ufuncs (like np.add), nargs is 2 + return ufunc.nargs > 1 diff --git a/mesa/visualization/mpl_space_drawing.py b/mesa/visualization/mpl_space_drawing.py index ea7687c4210..5fc7578a34e 100644 --- a/mesa/visualization/mpl_space_drawing.py +++ b/mesa/visualization/mpl_space_drawing.py @@ -178,7 +178,7 @@ def draw_property_layers( property_layers = space.properties except AttributeError: # new style spaces - property_layers = space.property_layers + property_layers = space._mesa_property_layers for layer_name, portrayal in propertylayer_portrayal.items(): layer = property_layers.get(layer_name, None) @@ -186,7 +186,7 @@ def draw_property_layers( continue data = layer.data.astype(float) if layer.data.dtype == bool else layer.data - width, height = data.shape if space is None else (space.width, space.height) + width, height = data.shape # if space is None else (space.width, space.height) if space and data.shape != (width, height): warnings.warn( diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py index b11701745cd..53fb737d5a7 100644 --- a/tests/test_cell_space.py +++ b/tests/test_cell_space.py @@ -16,9 +16,9 @@ Network, OrthogonalMooreGrid, OrthogonalVonNeumannGrid, + PropertyLayer, VoronoiGrid, ) -from mesa.space import PropertyLayer def test_orthogonal_grid_neumann(): @@ -163,7 +163,7 @@ def test_orthogonal_grid_moore(): for connection in grid._cells[(5, 5)].connections.values(): # fmt: off assert connection.coordinate in {(4, 4), (4, 5), (4, 6), - (5, 4), (5, 6), + (5, 4), (5, 6), (6, 4), (6, 5), (6, 6)} # fmt: on @@ -175,7 +175,7 @@ def test_orthogonal_grid_moore(): for connection in grid._cells[(0, 0)].connections.values(): # fmt: off assert connection.coordinate in {(9, 9), (9, 0), (9, 1), - (0, 9), (0, 1), + (0, 9), (0, 1), (1, 9), (1, 0), (1, 1)} # fmt: on @@ -206,9 +206,12 @@ def test_orthogonal_grid_moore_3d(): assert len(grid._cells[(5, 5, 5)].connections.values()) == 26 for connection in grid._cells[(5, 5, 5)].connections.values(): # fmt: off - assert connection.coordinate in {(4, 4, 4), (4, 4, 5), (4, 4, 6), (4, 5, 4), (4, 5, 5), (4, 5, 6), (4, 6, 4), (4, 6, 5), (4, 6, 6), - (5, 4, 4), (5, 4, 5), (5, 4, 6), (5, 5, 4), (5, 5, 6), (5, 6, 4), (5, 6, 5), (5, 6, 6), - (6, 4, 4), (6, 4, 5), (6, 4, 6), (6, 5, 4), (6, 5, 5), (6, 5, 6), (6, 6, 4), (6, 6, 5), (6, 6, 6)} + assert connection.coordinate in {(4, 4, 4), (4, 4, 5), (4, 4, 6), (4, 5, 4), (4, 5, 5), (4, 5, 6), (4, 6, 4), + (4, 6, 5), (4, 6, 6), + (5, 4, 4), (5, 4, 5), (5, 4, 6), (5, 5, 4), (5, 5, 6), (5, 6, 4), (5, 6, 5), + (5, 6, 6), + (6, 4, 4), (6, 4, 5), (6, 4, 6), (6, 5, 4), (6, 5, 5), (6, 5, 6), (6, 6, 4), + (6, 6, 5), (6, 6, 6)} # fmt: on # Moore neighborhood, torus True, top corner @@ -218,9 +221,12 @@ def test_orthogonal_grid_moore_3d(): assert len(grid._cells[(0, 0, 0)].connections.values()) == 26 for connection in grid._cells[(0, 0, 0)].connections.values(): # fmt: off - assert connection.coordinate in {(9, 9, 9), (9, 9, 0), (9, 9, 1), (9, 0, 9), (9, 0, 0), (9, 0, 1), (9, 1, 9), (9, 1, 0), (9, 1, 1), - (0, 9, 9), (0, 9, 0), (0, 9, 1), (0, 0, 9), (0, 0, 1), (0, 1, 9), (0, 1, 0), (0, 1, 1), - (1, 9, 9), (1, 9, 0), (1, 9, 1), (1, 0, 9), (1, 0, 0), (1, 0, 1), (1, 1, 9), (1, 1, 0), (1, 1, 1)} + assert connection.coordinate in {(9, 9, 9), (9, 9, 0), (9, 9, 1), (9, 0, 9), (9, 0, 0), (9, 0, 1), (9, 1, 9), + (9, 1, 0), (9, 1, 1), + (0, 9, 9), (0, 9, 0), (0, 9, 1), (0, 0, 9), (0, 0, 1), (0, 1, 9), (0, 1, 0), + (0, 1, 1), + (1, 9, 9), (1, 9, 0), (1, 9, 1), (1, 0, 9), (1, 0, 0), (1, 0, 1), (1, 1, 9), + (1, 1, 0), (1, 1, 1)} # fmt: on @@ -262,15 +268,24 @@ def test_orthogonal_grid_moore_4d(): assert len(grid._cells[(5, 5, 5, 5)].connections.values()) == 80 for connection in grid._cells[(5, 5, 5, 5)].connections.values(): # fmt: off - assert connection.coordinate in {(4, 4, 4, 4), (4, 4, 4, 5), (4, 4, 4, 6), (4, 4, 5, 4), (4, 4, 5, 5), (4, 4, 5, 6), (4, 4, 6, 4), (4, 4, 6, 5), (4, 4, 6, 6), - (4, 5, 4, 4), (4, 5, 4, 5), (4, 5, 4, 6), (4, 5, 5, 4), (4, 5, 5, 5), (4, 5, 5, 6), (4, 5, 6, 4), (4, 5, 6, 5), (4, 5, 6, 6), - (4, 6, 4, 4), (4, 6, 4, 5), (4, 6, 4, 6), (4, 6, 5, 4), (4, 6, 5, 5), (4, 6, 5, 6), (4, 6, 6, 4), (4, 6, 6, 5), (4, 6, 6, 6), - (5, 4, 4, 4), (5, 4, 4, 5), (5, 4, 4, 6), (5, 4, 5, 4), (5, 4, 5, 5), (5, 4, 5, 6), (5, 4, 6, 4), (5, 4, 6, 5), (5, 4, 6, 6), - (5, 5, 4, 4), (5, 5, 4, 5), (5, 5, 4, 6), (5, 5, 5, 4), (5, 5, 5, 6), (5, 5, 6, 4), (5, 5, 6, 5), (5, 5, 6, 6), - (5, 6, 4, 4), (5, 6, 4, 5), (5, 6, 4, 6), (5, 6, 5, 4), (5, 6, 5, 5), (5, 6, 5, 6), (5, 6, 6, 4), (5, 6, 6, 5), (5, 6, 6, 6), - (6, 4, 4, 4), (6, 4, 4, 5), (6, 4, 4, 6), (6, 4, 5, 4), (6, 4, 5, 5), (6, 4, 5, 6), (6, 4, 6, 4), (6, 4, 6, 5), (6, 4, 6, 6), - (6, 5, 4, 4), (6, 5, 4, 5), (6, 5, 4, 6), (6, 5, 5, 4), (6, 5, 5, 5), (6, 5, 5, 6), (6, 5, 6, 4), (6, 5, 6, 5), (6, 5, 6, 6), - (6, 6, 4, 4), (6, 6, 4, 5), (6, 6, 4, 6), (6, 6, 5, 4), (6, 6, 5, 5), (6, 6, 5, 6), (6, 6, 6, 4), (6, 6, 6, 5), (6, 6, 6, 6)} + assert connection.coordinate in {(4, 4, 4, 4), (4, 4, 4, 5), (4, 4, 4, 6), (4, 4, 5, 4), (4, 4, 5, 5), + (4, 4, 5, 6), (4, 4, 6, 4), (4, 4, 6, 5), (4, 4, 6, 6), + (4, 5, 4, 4), (4, 5, 4, 5), (4, 5, 4, 6), (4, 5, 5, 4), (4, 5, 5, 5), + (4, 5, 5, 6), (4, 5, 6, 4), (4, 5, 6, 5), (4, 5, 6, 6), + (4, 6, 4, 4), (4, 6, 4, 5), (4, 6, 4, 6), (4, 6, 5, 4), (4, 6, 5, 5), + (4, 6, 5, 6), (4, 6, 6, 4), (4, 6, 6, 5), (4, 6, 6, 6), + (5, 4, 4, 4), (5, 4, 4, 5), (5, 4, 4, 6), (5, 4, 5, 4), (5, 4, 5, 5), + (5, 4, 5, 6), (5, 4, 6, 4), (5, 4, 6, 5), (5, 4, 6, 6), + (5, 5, 4, 4), (5, 5, 4, 5), (5, 5, 4, 6), (5, 5, 5, 4), (5, 5, 5, 6), + (5, 5, 6, 4), (5, 5, 6, 5), (5, 5, 6, 6), + (5, 6, 4, 4), (5, 6, 4, 5), (5, 6, 4, 6), (5, 6, 5, 4), (5, 6, 5, 5), + (5, 6, 5, 6), (5, 6, 6, 4), (5, 6, 6, 5), (5, 6, 6, 6), + (6, 4, 4, 4), (6, 4, 4, 5), (6, 4, 4, 6), (6, 4, 5, 4), (6, 4, 5, 5), + (6, 4, 5, 6), (6, 4, 6, 4), (6, 4, 6, 5), (6, 4, 6, 6), + (6, 5, 4, 4), (6, 5, 4, 5), (6, 5, 4, 6), (6, 5, 5, 4), (6, 5, 5, 5), + (6, 5, 5, 6), (6, 5, 6, 4), (6, 5, 6, 5), (6, 5, 6, 6), + (6, 6, 4, 4), (6, 6, 4, 5), (6, 6, 4, 6), (6, 6, 5, 4), (6, 6, 5, 5), + (6, 6, 5, 6), (6, 6, 6, 4), (6, 6, 6, 5), (6, 6, 6, 6)} # fmt: on @@ -379,7 +394,7 @@ def test_hexgrid(): for connection in grid._cells[(1, 0)].connections.values(): # fmt: off assert connection.coordinate in {(0, 0), (0, 1), - (1, 1), + (1, 1), (2, 0), (2, 1)} # middle odd row @@ -387,7 +402,7 @@ def test_hexgrid(): for connection in grid._cells[(5, 5)].connections.values(): # fmt: off assert connection.coordinate in {(4, 5), (4, 6), - (5, 4), (5, 6), + (5, 4), (5, 6), (6, 5), (6, 6)} # fmt: on @@ -397,7 +412,7 @@ def test_hexgrid(): for connection in grid._cells[(4, 4)].connections.values(): # fmt: off assert connection.coordinate in {(3, 3), (3, 4), - (4, 3), (4, 5), + (4, 3), (4, 5), (5, 3), (5, 4)} # fmt: on @@ -410,7 +425,7 @@ def test_hexgrid(): for connection in grid._cells[(0, 0)].connections.values(): # fmt: off assert connection.coordinate in {(9, 9), (9, 0), - (0, 9), (0, 1), + (0, 9), (0, 1), (1, 9), (1, 0)} # fmt: on @@ -432,6 +447,10 @@ def test_networkgrid(): for connection in cell.connections.values(): assert connection.coordinate in G.neighbors(i) + import pickle + + pickle.loads(pickle.dumps(grid)) # noqa: S301 + def test_voronoigrid(): """Test VoronoiGrid.""" @@ -620,88 +639,247 @@ def test_empty_cell_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, random=random.Random(42)) + dimensions = (10, 10) + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) # Test adding a PropertyLayer to the grid - elevation = PropertyLayer("elevation", width, height, default_value=0) + elevation = PropertyLayer("elevation", dimensions, default_value=0.0) grid.add_property_layer(elevation) - assert "elevation" in grid.property_layers - assert len(grid.property_layers) == 1 + assert "elevation" in grid._mesa_property_layers + assert len(grid._mesa_property_layers) == 2 ## empty is always there + assert grid.elevation is elevation + + with pytest.raises(AttributeError): + grid.elevation = 0 # Test accessing PropertyLayer from a cell cell = grid._cells[(0, 0)] - assert "elevation" in cell._mesa_property_layers - assert cell.get_property("elevation") == 0 + assert hasattr(cell, "elevation") + assert cell.elevation == 0.0 # Test setting property value for a cell - cell.set_property("elevation", 100) - assert cell.get_property("elevation") == 100 + cell.elevation = 100 + assert cell.elevation == 100 + assert elevation.data[0, 0] == 100 # Test modifying property value for a cell - cell.modify_property("elevation", lambda x: x + 50) - assert cell.get_property("elevation") == 150 + cell.elevation += 50 + assert cell.elevation == 150 + assert elevation.data[0, 0] == 150 - cell.modify_property("elevation", np.add, 50) - assert cell.get_property("elevation") == 200 + cell.elevation = np.add(cell.elevation, 50) + assert cell.elevation == 200 + assert elevation.data[0, 0] == 200 # Test modifying PropertyLayer values grid.set_property("elevation", 100, condition=lambda value: value == 200) - assert cell.get_property("elevation") == 100 + assert cell.elevation == 100 # Test modifying PropertyLayer using numpy operations grid.modify_properties("elevation", np.add, 50) - assert cell.get_property("elevation") == 150 + assert cell.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 + assert "elevation" not in grid._mesa_property_layers + assert not hasattr(cell, "elevation") + + # what happens if we add a layer whose name clashes with an existing cell attribute? + dimensions = (10, 10) + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + + with pytest.raises(ValueError): + grid.create_property_layer("capacity", 1, dtype=int) + + +def test_copy_pickle_with_propertylayers(): + """Test deepcopy and pickle with dynamically created GridClass and ProperyLayer descriptors.""" + import copy + import pickle + + dimensions = (10, 10) + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + + grid2 = copy.deepcopy(grid) + assert grid2._cells[(0, 0)].empty + + data = grid2._mesa_property_layers["empty"].data + grid2._cells[(0, 0)].empty = False + assert grid2._cells[(0, 0)].empty == data[0, 0] + + dimensions = (10, 10) + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + dump = pickle.dumps(grid) + grid2 = pickle.loads(dump) # noqa: S301 + assert grid2._cells[(0, 0)].empty + data = grid2._mesa_property_layers["empty"].data + grid2._cells[(0, 0)].empty = False + assert grid2._cells[(0, 0)].empty == data[0, 0] 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) + dimensions = (5, 5) + elevation = PropertyLayer("elevation", dimensions, default_value=0.0) + temperature = PropertyLayer("temperature", dimensions, default_value=20.0) # Test initialization with a single PropertyLayer - grid1 = OrthogonalMooreGrid((width, height), torus=False, random=random.Random(42)) + grid1 = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) grid1.add_property_layer(elevation) - assert "elevation" in grid1.property_layers - assert len(grid1.property_layers) == 1 + assert "elevation" in grid1._mesa_property_layers + assert len(grid1._mesa_property_layers) == 2 ## empty is already there # Test initialization with multiple PropertyLayers - grid2 = OrthogonalMooreGrid((width, height), torus=False, random=random.Random(42)) - 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 + grid2 = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + grid2.add_property_layer(temperature) + grid2.add_property_layer(elevation) + # + assert "temperature" in grid2._mesa_property_layers + assert "elevation" in grid2._mesa_property_layers + assert len(grid2._mesa_property_layers) == 3 # 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") + assert cell.elevation == 10 + assert cell.temperature == 25 + + +def test_get_neighborhood_mask(): + """Test get_neighborhood_mask.""" + dimensions = (5, 5) + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + grid.create_property_layer("elevation", default_value=0.0) + + grid.get_neighborhood_mask((2, 2)) + + mask = grid.get_neighborhood_mask((2, 2)) + for cell in grid._cells[(2, 2)].connections.values(): + assert mask[cell.coordinate] + assert mask[grid._cells[(2, 2)].coordinate] + + mask = grid.get_neighborhood_mask((2, 2), include_center=False) + for cell in grid._cells[(2, 2)].connections.values(): + assert mask[cell.coordinate] + assert not mask[grid._cells[(2, 2)].coordinate] + + +def test_select_cells(): + """Test select_cells.""" + dimensions = (5, 5) + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + + data = np.random.default_rng(12456).random((5, 5)) + grid.add_property_layer(PropertyLayer.from_data("elevation", data)) + + # fixme, add an agent and update the np.all test accordingly + mask = grid.select_cells( + conditions={"elevation": lambda x: x > 0.5}, return_list=False, only_empty=True + ) + assert mask.shape == (5, 5) + assert np.all(mask == (data > 0.5)) + + mask = grid.select_cells( + conditions={"elevation": lambda x: x > 0.5}, return_list=False, only_empty=False + ) + assert mask.shape == (5, 5) + assert np.all(mask == (data > 0.5)) + + # fixme add extreme_values highest and lowest + mask = grid.select_cells( + extreme_values={"elevation": "highest"}, return_list=False, only_empty=False + ) + assert mask.shape == (5, 5) + assert np.all(mask == (data == data.max())) + + mask = grid.select_cells( + extreme_values={"elevation": "lowest"}, return_list=False, only_empty=False + ) + assert mask.shape == (5, 5) + assert np.all(mask == (data == data.min())) + + with pytest.raises(ValueError): + grid.select_cells( + extreme_values={"elevation": "weird"}, return_list=False, only_empty=False + ) + + # fixme add pre-specified mask to any other option + + +def test_property_layer(): + """Test various property layer methods.""" + elevation = PropertyLayer("elevation", (5, 5), default_value=0.0) + + # test set_cells + elevation.set_cells(10) + assert np.all(elevation.data == 10) + + elevation.set_cells(np.ones((5, 5))) + assert np.all(elevation.data == 1) + + with pytest.raises(ValueError): + elevation.set_cells(np.ones((6, 6))) + + data = np.random.default_rng(42).random((5, 5)) + layer = PropertyLayer.from_data("some_name", data) + + def condition(x): + return x > 0.5 + + layer.set_cells(1, condition=condition) + assert np.all((layer.data == 1) == (data > 0.5)) + + # modify_cells + data = np.zeros((10, 10)) + layer = PropertyLayer.from_data("some_name", data) + + layer.data = np.zeros((10, 10)) + layer.modify_cells(lambda x: x + 2) + assert np.all(layer.data == 2) + + layer.data = np.ones((10, 10)) + layer.modify_cells(np.multiply, 3) + assert np.all(layer.data[3, 3] == 3) + + data = np.random.default_rng(42).random((10, 10)) + layer.data = np.random.default_rng(42).random((10, 10)) + layer.modify_cells(np.add, value=3, condition=condition) + assert np.all((layer.data > 3.5) == (data > 0.5)) + + with pytest.raises(ValueError): + layer.modify_cells(np.add) # Missing value for ufunc + + # aggregate + layer.data = np.ones((10, 10)) + assert layer.aggregate(np.sum) == 100 def test_property_layer_errors(): """Test error handling for PropertyLayers.""" - width, height = 5, 5 - grid = OrthogonalMooreGrid((width, height), torus=False, random=random.Random(42)) - elevation = PropertyLayer("elevation", width, height, default_value=0) + dimensions = 5, 5 + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + elevation = PropertyLayer("elevation", dimensions, default_value=0.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) + # test adding a layer with different dimensions than space + dimensions = 5, 5 + grid = OrthogonalMooreGrid(dimensions, torus=False, random=random.Random(42)) + elevation = PropertyLayer("elevation", (10, 10), default_value=0.0) + with pytest.raises( + ValueError, + match="Dimensions of property layer do not match the dimensions of the grid", + ): + grid.add_property_layer(elevation) + + with pytest.warns(UserWarning): + PropertyLayer("elevation", (10, 10), default_value=0, dtype=float) + def test_cell_agent(): # noqa: D103 cell1 = Cell((1,), capacity=None, random=random.Random()) diff --git a/tests/test_components_matplotlib.py b/tests/test_components_matplotlib.py index f258b58d90b..8b773464dea 100644 --- a/tests/test_components_matplotlib.py +++ b/tests/test_components_matplotlib.py @@ -231,7 +231,7 @@ def test_draw_property_layers(): model = Model(seed=42) grid = OrthogonalMooreGrid((10, 10), torus=True, random=model.random, capacity=1) - grid.add_property_layer(PropertyLayer("test", grid.width, grid.height, 0)) + grid.create_property_layer("test", 0.0) fig, ax = plt.subplots() draw_property_layers(grid, {"test": {"colormap": "viridis", "colorbar": True}}, ax)