Bugfix for deepcopy / pickling discrete spaces (#2378)

* bugfix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix for 3.10

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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

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()

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)
 

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