diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py index 6c92afbe162..4bc6e874a64 100644 --- a/mesa/experimental/cell_space/cell.py +++ b/mesa/experimental/cell_space/cell.py @@ -211,3 +211,12 @@ def modify_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__ + state = (self.__dict__, {k: getattr(self, k) for k in self.__slots__}) + state[1][ + "connections" + ] = {} # replace this with empty connections to avoid infinite recursion error in pickle/deepcopy + return state diff --git a/mesa/experimental/cell_space/discrete_space.py b/mesa/experimental/cell_space/discrete_space.py index 1cfaedeb167..7ed06a84aae 100644 --- a/mesa/experimental/cell_space/discrete_space.py +++ b/mesa/experimental/cell_space/discrete_space.py @@ -22,7 +22,7 @@ class DiscreteSpace(Generic[T]): all_cells (CellCollection): The cells composing the discrete space random (Random): The random number generator cell_klass (Type) : the type of cell class - empties (CellCollection) : collecction of all cells that are empty + empties (CellCollection) : collection of all cells that are empty property_layers (dict[str, PropertyLayer]): the property layers of the discrete space """ @@ -55,6 +55,7 @@ def __init__( def cutoff_empties(self): # noqa return 7.953 * len(self._cells) ** 0.384 + def _connect_cells(self): ... def _connect_single_cell(self, cell: T): ... @cached_property @@ -134,3 +135,8 @@ def modify_properties( 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 + self._connect_cells() diff --git a/mesa/experimental/cell_space/network.py b/mesa/experimental/cell_space/network.py index 551305267af..1d49f2ee1e5 100644 --- a/mesa/experimental/cell_space/network.py +++ b/mesa/experimental/cell_space/network.py @@ -34,6 +34,9 @@ def __init__( node_id, capacity, random=self.random ) + self._connect_cells() + + def _connect_cells(self) -> None: for cell in self.all_cells: self._connect_single_cell(cell) diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py index 4d52e159045..fd6f946902f 100644 --- a/tests/test_cell_space.py +++ b/tests/test_cell_space.py @@ -691,3 +691,42 @@ def test_patch(): # noqa: D103 agent.remove() assert agent not in model._agents + + +def test_copying_discrete_spaces(): # noqa: D103 + # inspired by #2373 + # we use deepcopy but this also applies to pickle + import copy + + import networkx as nx + + grid = OrthogonalMooreGrid((100, 100)) + grid_copy = copy.deepcopy(grid) + + c1 = grid[(5, 5)].connections + c2 = grid_copy[(5, 5)].connections + + for c1, c2 in zip(grid.all_cells, grid_copy.all_cells): + for k, v in c1.connections.items(): + assert v.coordinate == c2.connections[k].coordinate + + n = 10 + m = 20 + seed = 42 + G = nx.gnm_random_graph(n, m, seed=seed) # noqa: N806 + grid = Network(G) + grid_copy = copy.deepcopy(grid) + + for c1, c2 in zip(grid.all_cells, grid_copy.all_cells): + for k, v in c1.connections.items(): + assert v.coordinate == c2.connections[k].coordinate + + grid = HexGrid((100, 100)) + grid_copy = copy.deepcopy(grid) + + c1 = grid[(5, 5)].connections + c2 = grid_copy[(5, 5)].connections + + for c1, c2 in zip(grid.all_cells, grid_copy.all_cells): + for k, v in c1.connections.items(): + assert v.coordinate == c2.connections[k].coordinate