Repository: projectmesa/mesa
Pull request #2378 (2024-10-16 17:58:26)

Bugfix for deepcopy / pickling discrete spaces
As highlighted in #2373, deepcopy (and by extension) pickle breaks on discrete spaces. As a fix, in `Cell.__getstate__`, I set "connections" to an empty dict. In `DiscreteSpace.__setstate__`, I rebuild the connections via `DiscreteSpace._connect_cells()`. This latter method already existed in `Grid` and `VoronoiGrid`, but not yet in `Network`. Thus, I moved the method stub to DiscreteSpace and implemented it for Network.

This also adds unit tests for everything. I first confirmed that the unit tests failed and reproduce the original error. 

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔵 +0.2% [-0.2%, +0.7%] | 🔵 -0.3% [-0.5%, -0.1%] |
| BoltzmannWealth | large  | 🔵 -0.0% [-0.8%, +0.7%] | 🔵 -0.2% [-1.0%, +0.6%] |
| Schelling       | small  | 🔵 -0.2% [-0.5%, +0.2%] | 🔵 +0.2% [-0.1%, +0.5%] |
| Schelling       | large  | 🔵 +0.7% [+0.1%, +1.3%] | 🔵 -0.3% [-1.0%, +0.5%] |
| WolfSheep       | small  | 🔵 +0.1% [-0.3%, +0.5%] | 🔵 -0.3% [-0.6%, -0.1%] |
| WolfSheep       | large  | 🔵 +0.9% [-0.7%, +2.2%] | 🔵 +0.5% [-1.5%, +2.3%] |
| BoidFlockers    | small  | 🔵 +0.1% [-0.5%, +0.7%] | 🔵 +1.7% [+1.1%, +2.4%] |
| BoidFlockers    | large  | 🔵 +0.4% [-0.9%, +1.4%] | 🔵 +2.1% [+1.3%, +2.9%] |

quaquel: tests fail on 3.10 because `super.__getstate__()` was only added in 3.11. 
EwoutH: > tests fail on 3.10 because `super.__getstate__()` was only added in 3.11.

Finally a case in which having all these different CI runs pay off!

Easiest way would be to add a warning that the cell space only works with 3.11 and above.

SPEC 0 says we can drop 3.10 now. I say we do that directly after 3.0 is branched off, and require Python 3.11 for Mesa 3.1 and above.
quaquel: > Easiest way would be to add a warning that the cell space only works with 3.11 and above.

Using super here was me being lazy. I was surprised it worked (But not for `__setstate__` strangely enough). Its only a few lines of code more to have it work also for 3.10.
quaquel: now it errors not on the tests but the code coverage?
Corvince: Just to clarify something here, since I havent worked with getstate or setstate yet. This code changes that when I pickle a cell, its connections become an empty dictionary. And when I unpickle a CellSpace the connections the connections are restored. Is this correct? Wouldn't this mean if I pickle a cell by itself and unpickle it I lose the connections? Not sure how I feel about this. Its probably still reasonable for the time being, but it could also lead to very unexpected behaviors. 

Patch:
diff --git a/.github/release.yml b/.github/release.yml
index 0277adf3..e29dd2ed 100644
--- a/.github/release.yml
+++ b/.github/release.yml
@@ -19,6 +19,9 @@ changelog:
     - title: 🐛 Bugs fixed
       labels:
         - bug
+    - title: 🔍 Examples updated
+      labels:
+        - example
     - title: 📜 Documentation improvements
       labels:
         - docs
diff --git a/docs/tutorials/visualization_tutorial.ipynb b/docs/tutorials/visualization_tutorial.ipynb
index b536752f..03c594bc 100644
--- a/docs/tutorials/visualization_tutorial.ipynb
+++ b/docs/tutorials/visualization_tutorial.ipynb
@@ -194,7 +194,9 @@
     "\n",
     "If the visualization elements provided by Mesa aren't enough for you, you can build your own and plug them into the model server.\n",
     "\n",
-    "For this example, let's build a simple histogram visualization, which can count the number of agents with each value of wealth."
+    "For this example, let's build a simple histogram visualization, which can count the number of agents with each value of wealth.\n",
+    "\n",
+    "**Note:** Due to the way solara works we need to trigger an update whenever the underlying model changes. For this you need to register an update counter with every component."
    ]
   },
   {
@@ -205,9 +207,11 @@
    "source": [
     "import solara\n",
     "from matplotlib.figure import Figure\n",
+    "from mesa.visualization.utils import update_counter\n",
     "\n",
     "@solara.component\n",
     "def Histogram(model):\n",
+    "    update_counter.get() # This is required to update the counter\n",
     "    # Note: you must initialize a figure using this method instead of\n",
     "    # plt.figure(), for thread safety purpose\n",
     "    fig = Figure()\n",
diff --git a/mesa/agent.py b/mesa/agent.py
index 2b6d2f99..18454164 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -62,9 +62,9 @@ class Agent:
         super().__init__(*args, **kwargs)
 
         self.model: Model = model
-        self.model.register_agent(self)
         self.unique_id: int = next(self._ids[model])
         self.pos: Position | None = None
+        self.model.register_agent(self)
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py
index 6c92afbe..4bc6e874 100644
--- a/mesa/experimental/cell_space/cell.py
+++ b/mesa/experimental/cell_space/cell.py
@@ -211,3 +211,12 @@ class Cell:
         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 1cfaedeb..7ed06a84 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 @@ class DiscreteSpace(Generic[T]):
     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 @@ class DiscreteSpace(Generic[T]):
             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 55130526..1d49f2ee 100644
--- a/mesa/experimental/cell_space/network.py
+++ b/mesa/experimental/cell_space/network.py
@@ -34,6 +34,9 @@ class Network(DiscreteSpace[Cell]):
                 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)
 


Test patch:
diff --git a/tests/test_cell_space.py b/tests/test_cell_space.py
index 4d52e159..fd6f9469 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
diff --git a/tests/test_time.py b/tests/test_time.py
new file mode 100644
index 00000000..f180ac25
--- /dev/null
+++ b/tests/test_time.py
@@ -0,0 +1,317 @@
+"""Test the advanced schedulers."""
+
+import unittest
+from unittest import TestCase, mock
+
+from mesa.agent import Agent
+from mesa.model import Model
+from mesa.time import (
+    BaseScheduler,
+    RandomActivation,
+    RandomActivationByType,
+    SimultaneousActivation,
+    StagedActivation,
+)
+
+RANDOM = "random"
+STAGED = "staged"
+SIMULTANEOUS = "simultaneous"
+RANDOM_BY_TYPE = "random_by_type"
+
+
+class MockAgent(Agent):
+    """Minimalistic agent for testing purposes."""
+
+    def __init__(self, model):  # noqa: D107
+        super().__init__(model)
+        self.steps = 0
+        self.advances = 0
+
+    def kill_other_agent(self):  # noqa: D102
+        for agent in self.model.schedule.agents:
+            if agent is not self:
+                agent.remove()
+
+    def stage_one(self):  # noqa: D102
+        if self.model.enable_kill_other_agent:
+            self.kill_other_agent()
+        self.model.log.append(f"{self.unique_id}_1")
+
+    def stage_two(self):  # noqa: D102
+        self.model.log.append(f"{self.unique_id}_2")
+
+    def advance(self):  # noqa: D102
+        self.advances += 1
+
+    def step(self):  # noqa: D102
+        if self.model.enable_kill_other_agent:
+            self.kill_other_agent()
+        self.steps += 1
+        self.model.log.append(self.unique_id)
+
+
+class MockModel(Model):  # Noqa: D101
+    def __init__(
+        self, shuffle=False, activation=STAGED, enable_kill_other_agent=False, seed=None
+    ):
+        """Creates a Model instance with a schedule.
+
+        Args:
+            shuffle (Bool): whether or not to instantiate a scheduler
+                            with shuffling.
+                            This option is only used for
+                            StagedActivation schedulers.
+
+            activation (str): which kind of scheduler to use.
+                              'random' creates a RandomActivation scheduler.
+                              'staged' creates a StagedActivation scheduler.
+                              The default scheduler is a BaseScheduler.
+            enable_kill_other_agent: whether or not to enable kill_other_agent
+            seed : rng
+        """
+        super().__init__(seed=seed)
+        self.log = []
+        self.enable_kill_other_agent = enable_kill_other_agent
+
+        # Make scheduler
+        if activation == STAGED:
+            model_stages = ["stage_one", "model.model_stage", "stage_two"]
+            self.schedule = StagedActivation(
+                self, stage_list=model_stages, shuffle=shuffle
+            )
+        elif activation == RANDOM:
+            self.schedule = RandomActivation(self)
+        elif activation == SIMULTANEOUS:
+            self.schedule = SimultaneousActivation(self)
+        elif activation == RANDOM_BY_TYPE:
+            self.schedule = RandomActivationByType(self)
+        else:
+            self.schedule = BaseScheduler(self)
+
+        # Make agents
+        for _ in range(2):
+            agent = MockAgent(self)
+            self.schedule.add(agent)
+
+    def step(self):  # noqa: D102
+        self.schedule.step()
+
+    def model_stage(self):  # noqa: D102
+        self.log.append("model_stage")
+
+
+class TestStagedActivation(TestCase):
+    """Test the staged activation."""
+
+    expected_output = ["1_1", "1_1", "model_stage", "1_2", "1_2"]
+
+    def test_no_shuffle(self):
+        """Testing the staged activation without shuffling."""
+        model = MockModel(shuffle=False)
+        model.step()
+        model.step()
+        assert all(i == j for i, j in zip(model.log[:5], model.log[5:]))
+
+    def test_shuffle(self):
+        """Test the staged activation with shuffling."""
+        model = MockModel(shuffle=True)
+        model.step()
+        for output in self.expected_output[:2]:
+            assert output in model.log[:2]
+        for output in self.expected_output[3:]:
+            assert output in model.log[3:]
+        assert self.expected_output[2] == model.log[2]
+
+    def test_shuffle_shuffles_agents(self):  # noqa: D102
+        model = MockModel(shuffle=True)
+        a = mock.Mock()
+        model.schedule._agents.random = a
+        assert a.shuffle.call_count == 0
+        model.step()
+        assert a.shuffle.call_count == 1
+
+    def test_remove(self):
+        """Test the staged activation can remove an agent."""
+        model = MockModel(shuffle=True)
+        agents = list(model.schedule._agents)
+        agent = agents[0]
+        model.schedule.remove(agents[0])
+        assert agent not in model.schedule.agents
+
+    def test_intrastep_remove(self):
+        """Test the staged activation can remove an agent in a step of another agent.
+
+        so that the one removed doesn't step.
+        """
+        model = MockModel(shuffle=True, enable_kill_other_agent=True)
+        model.step()
+        assert len(model.log) == 3
+
+    def test_add_existing_agent(self):  # noqa: D102
+        model = MockModel()
+        agent = model.schedule.agents[0]
+        with self.assertRaises(Exception):
+            model.schedule.add(agent)
+
+
+class TestRandomActivation(TestCase):
+    """Test the random activation."""
+
+    def test_init(self):  # noqa: D102
+        model = Model()
+        agents = [MockAgent(model) for _ in range(10)]
+
+        scheduler = RandomActivation(model, agents)
+        assert all(agent in scheduler.agents for agent in agents)
+
+    def test_random_activation_step_shuffles(self):
+        """Test the random activation step."""
+        model = MockModel(activation=RANDOM)
+        a = mock.Mock()
+        model.schedule._agents.random = a
+        model.schedule.step()
+        assert a.shuffle.call_count == 1
+
+    def test_random_activation_step_increments_step_and_time_counts(self):
+        """Test the random activation step increments step and time counts."""
+        model = MockModel(activation=RANDOM)
+        assert model.schedule.steps == 0
+        assert model.schedule.time == 0
+        model.schedule.step()
+        assert model.schedule.steps == 1
+        assert model.schedule.time == 1
+
+    def test_random_activation_step_steps_each_agent(self):
+        """Test the random activation step causes each agent to step."""
+        model = MockModel(activation=RANDOM)
+        model.step()
+        agent_steps = [i.steps for i in model.schedule.agents]
+        # one step for each of 2 agents
+        assert all(x == 1 for x in agent_steps)
+
+    def test_intrastep_remove(self):
+        """Test the random activation can remove an agent in a step of another agent.
+
+        so that the one removed doesn't step.
+        """
+        model = MockModel(activation=RANDOM, enable_kill_other_agent=True)
+        model.step()
+        assert len(model.log) == 1
+
+    def test_get_agent_keys(self):  # noqa: D102
+        model = MockModel(activation=RANDOM)
+
+        keys = model.schedule.get_agent_keys()
+        agent_ids = [agent.unique_id for agent in model.agents]
+        assert all(entry_i == entry_j for entry_i, entry_j in zip(keys, agent_ids))
+
+        keys = model.schedule.get_agent_keys(shuffle=True)
+        agent_ids = {agent.unique_id for agent in model.agents}
+        assert all(entry in agent_ids for entry in keys)
+
+    def test_not_sequential(self):  # noqa: D102
+        model = MockModel(activation=RANDOM)
+        # Create 10 agents
+        for _ in range(10):
+            model.schedule.add(MockAgent(model))
+        # Run 3 steps
+        for _ in range(3):
+            model.step()
+        # Filter out non-integer elements from the log
+        filtered_log = [item for item in model.log if isinstance(item, int)]
+
+        # Check that there are no 18 consecutive agents id's in the filtered log
+        total_agents = 10
+        assert not any(
+            all(
+                (filtered_log[(i + j) % total_agents] - filtered_log[i]) % total_agents
+                == j % total_agents
+                for j in range(18)
+            )
+            for i in range(len(filtered_log))
+        ), f"Agents are activated sequentially:\n{filtered_log}"
+
+
+class TestSimultaneousActivation(TestCase):
+    """Test the simultaneous activation."""
+
+    def test_simultaneous_activation_step_steps_and_advances_each_agent(self):
+        """Test the simultaneous activation step causes each agent to step."""
+        model = MockModel(activation=SIMULTANEOUS)
+        model.step()
+        # one step for each of 2 agents
+        agent_steps = [i.steps for i in model.schedule.agents]
+        agent_advances = [i.advances for i in model.schedule.agents]
+        assert all(x == 1 for x in agent_steps)
+        assert all(x == 1 for x in agent_advances)
+
+
+class TestRandomActivationByType(TestCase):
+    """Test the random activation by type.
+
+    TODO implement at least 2 types of agents, and test that step_type only
+    does step for one type of agents, not the entire agents.
+    """
+
+    def test_init(self):  # noqa: D102
+        model = Model()
+        agents = [MockAgent(model) for _ in range(10)]
+        agents += [Agent(model) for _ in range(10)]
+
+        scheduler = RandomActivationByType(model, agents)
+        assert all(agent in scheduler.agents for agent in agents)
+
+    def test_random_activation_step_shuffles(self):
+        """Test the random activation by type step."""
+        model = MockModel(activation=RANDOM_BY_TYPE)
+        a = mock.Mock()
+        model.random = a
+        for agentset in model.schedule._agents_by_type.values():
+            agentset.random = a
+        model.schedule.step()
+        assert a.shuffle.call_count == 2
+
+    def test_random_activation_step_increments_step_and_time_counts(self):
+        """Test the random activation by type step increments step and time counts."""
+        model = MockModel(activation=RANDOM_BY_TYPE)
+        assert model.schedule.steps == 0
+        assert model.schedule.time == 0
+        model.schedule.step()
+        assert model.schedule.steps == 1
+        assert model.schedule.time == 1
+
+    def test_random_activation_step_steps_each_agent(self):
+        """Test the random activation by type step causes each agent to step."""
+        model = MockModel(activation=RANDOM_BY_TYPE)
+        model.step()
+        agent_steps = [i.steps for i in model.schedule.agents]
+        # one step for each of 2 agents
+        assert all(x == 1 for x in agent_steps)
+
+    def test_random_activation_counts(self):
+        """Test the random activation by type step causes each agent to step."""
+        model = MockModel(activation=RANDOM_BY_TYPE)
+
+        agent_types = model.agent_types
+        for agent_type in agent_types:
+            assert model.schedule.get_type_count(agent_type) == len(
+                model.agents_by_type[agent_type]
+            )
+
+    # def test_add_non_unique_ids(self):
+    #     """
+    #     Test that adding agent with duplicate ids result in an error.
+    #     TODO: we need to run this test on all schedulers, not just
+    #     TODO:: identical IDs is something for the agent, not the scheduler and should be tested there
+    #     RandomActivationByType.
+    #     """
+    #     model = MockModel(activation=RANDOM_BY_TYPE)
+    #     a = MockAgent(0, model)
+    #     b = MockAgent(0, model)
+    #     model.schedule.add(a)
+    #     with self.assertRaises(Exception):
+    #         model.schedule.add(b)
+
+
+if __name__ == "__main__":
+    unittest.main()