Skip to content

Commit

Permalink
Add full support for property layers to cell spaces (#2512)
Browse files Browse the repository at this point in the history
  • Loading branch information
quaquel authored Dec 4, 2024
1 parent adad7a2 commit 567339f
Show file tree
Hide file tree
Showing 8 changed files with 740 additions and 161 deletions.
2 changes: 2 additions & 0 deletions mesa/experimental/cell_space/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__ = [
Expand All @@ -34,5 +35,6 @@
"Network",
"OrthogonalMooreGrid",
"OrthogonalVonNeumannGrid",
"PropertyLayer",
"VoronoiGrid",
]
39 changes: 6 additions & 33 deletions mesa/experimental/cell_space/cell.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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",
Expand All @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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(
Expand All @@ -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:
Expand Down Expand Up @@ -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__
Expand Down
64 changes: 1 addition & 63 deletions mesa/experimental/cell_space/discrete_space.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
65 changes: 62 additions & 3 deletions mesa/experimental/cell_space/grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down
Loading

0 comments on commit 567339f

Please sign in to comment.