Repository: projectmesa/mesa
Pull request #2251 (2024-08-27 07:26:24)

Performance enhancements for Model.agents
This PR is a performance enhancement for Model.agents. It emerged out of [a discussion](#2224) on the weird scaling performance of the Boltzman wealth model. 

# Key changes
* `model.agents` now returns the agentset as maintained by the model, rather than a new copy based on the hard references
* agent registration and deregistration have been moved from the Agent into the model. The agent now calls `model.register` and `model.deregister`. This encapsulates everything cleanly inside the model class and makes Agent less dependent on the inner details of how Model manages the hard references to agents
* the setup of the relevant datastructures is moved into its own helper method, again this cleans up code.



github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]            | Run time [95% CI]          |
|--------------|--------|-------------------------------|----------------------------|
| Schelling    | small  | 🔴 +77.0% [+76.3%, +77.6%]    | 🔵 -0.4% [-0.7%, -0.1%]    |
| Schelling    | large  | 🔴 +94.2% [+92.9%, +95.3%]    | 🔵 -0.8% [-2.9%, +0.7%]    |
| WolfSheep    | small  | 🔴 +130.1% [+128.1%, +131.9%] | 🟢 -97.5% [-97.5%, -97.4%] |
| WolfSheep    | large  | 🔴 +130.8% [+129.1%, +132.4%] | 🟢 -99.9% [-99.9%, -99.9%] |
| BoidFlockers | small  | 🔴 +83.6% [+82.2%, +84.7%]    | 🔵 -0.3% [-1.0%, +0.6%]    |
| BoidFlockers | large  | 🔴 +83.1% [+82.0%, +84.1%]    | 🔵 -0.6% [-1.2%, +0.0%]    |

EwoutH: @quaquel can you reproduce the speedup of WolfSheep locally?
quaquel: > Performance benchmarks:
> 
> Model	Size	Init time [95% CI]	Run time [95% CI]
> Schelling	small	🔴 +77.0% [+76.3%, +77.6%]	🔵 -0.4% [-0.7%, -0.1%]
> Schelling	large	🔴 +94.2% [+92.9%, +95.3%]	🔵 -0.8% [-2.9%, +0.7%]
> WolfSheep	small	🔴 +130.1% [+128.1%, +131.9%]	🟢 -97.5% [-97.5%, -97.4%]
> WolfSheep	large	🔴 +130.8% [+129.1%, +132.4%]	🟢 -99.9% [-99.9%, -99.9%]
> BoidFlockers	small	🔴 +83.6% [+82.2%, +84.7%]	🔵 -0.3% [-1.0%, +0.6%]
> BoidFlockers	large	🔴 +83.1% [+82.0%, +84.1%]	🔵 -0.6% [-1.2%, +0.0%]

I'll try to figure out what is happening with the init times here. Its probably overhead from setting up the additional datastructures, which we get back when running the model (at least for wolf sheep?).
Corvince: Thanks for this PR, I like the way it encapsulates the logic much better. 

I think it is a good default to return the actual agentset and not just a copy. I think this is actually more intuitive.

Regarding the copy function, doesn't `agentset.select()` already do that? I'm not arguing against adding an additional method but I just want to mention that it's already possible and it should rightfully be discussed separately from this PR 
quaquel: > Regarding the copy function, doesn't agentset.select() already do that? I'm not arguing against adding an additional method but I just want to mention that it's already possible and it should rightfully be discussed separately from this PR

Yes, select does this as a corner case. But fair enought to seperate `copy` into a seperate PR
quaquel: This is ready for a final review.

quaquel: \trigger-benchmark
github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]         | Run time [95% CI]          |
|--------------|--------|----------------------------|----------------------------|
| Schelling    | small  | 🔴 +11.9% [+11.5%, +12.2%] | 🔵 +2.3% [+2.1%, +2.6%]    |
| Schelling    | large  | 🔴 +15.6% [+14.9%, +16.5%] | 🔴 +12.9% [+10.2%, +15.3%] |
| WolfSheep    | small  | 🔴 +23.1% [+21.8%, +24.3%] | 🔵 +2.8% [+0.7%, +4.9%]    |
| WolfSheep    | large  | 🔴 +28.6% [+27.5%, +29.7%] | 🔴 +27.6% [+25.2%, +29.9%] |
| BoidFlockers | small  | 🔴 +13.8% [+12.8%, +14.7%] | 🔵 +0.3% [-0.3%, +1.0%]    |
| BoidFlockers | large  | 🔴 +12.2% [+11.3%, +13.1%] | 🔵 +0.0% [-0.7%, +0.8%]    |

github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]         | Run time [95% CI]          |
|--------------|--------|----------------------------|----------------------------|
| Schelling    | small  | 🔴 +9.6% [+9.2%, +10.0%]   | 🔵 +1.2% [+0.6%, +1.8%]    |
| Schelling    | large  | 🔴 +14.3% [+13.5%, +15.2%] | 🔵 +5.0% [+2.6%, +7.8%]    |
| WolfSheep    | small  | 🔴 +22.2% [+20.7%, +23.7%] | 🔵 +4.0% [+1.8%, +6.0%]    |
| WolfSheep    | large  | 🔴 +22.3% [+20.8%, +24.0%] | 🔴 +12.0% [+10.0%, +13.6%] |
| BoidFlockers | small  | 🔴 +14.9% [+14.0%, +15.7%] | 🔵 +1.7% [+0.8%, +2.6%]    |
| BoidFlockers | large  | 🔴 +12.8% [+12.2%, +13.5%] | 🔵 +1.6% [+1.0%, +2.3%]    |

EwoutH: Thanks for this. I think the regressions in init time are perfectably acceptable. What I'm curious about is why wolf_sheep has become slower. Because a lot of agents get added and removed?

I would like to get #2252 in first to get the full picture.
github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]         | Run time [95% CI]          |
|-----------------|--------|----------------------------|----------------------------|
| BoltzmannWealth | small  | 🔴 +16.5% [+14.0%, +19.0%] | 🟢 -12.1% [-12.2%, -11.9%] |
| BoltzmannWealth | large  | 🔴 +47.5% [+21.2%, +84.9%] | 🟢 -17.8% [-20.0%, -15.5%] |
| Schelling       | small  | 🔴 +10.5% [+10.1%, +10.9%] | 🔵 -0.7% [-1.0%, -0.4%]    |
| Schelling       | large  | 🔴 +13.7% [+12.8%, +14.7%] | 🔵 +0.6% [-1.5%, +2.6%]    |
| WolfSheep       | small  | 🔴 +21.8% [+20.3%, +23.3%] | 🔵 -0.2% [-3.6%, +3.3%]    |
| WolfSheep       | large  | 🔴 +20.0% [+19.1%, +20.9%] | 🔵 +4.3% [+2.1%, +6.7%]    |
| BoidFlockers    | small  | 🔴 +12.9% [+12.0%, +13.8%] | 🔵 +0.6% [-0.2%, +1.2%]    |
| BoidFlockers    | large  | 🔴 +12.4% [+11.7%, +13.2%] | 🔵 -0.1% [-0.3%, +0.1%]    |

EwoutH: Let me review it in depth tonight (I'm trying to stay in the model-building flow)

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index f137ba09..33a7e619 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -49,19 +49,12 @@ class Agent:
         self.model = model
         self.pos: Position | None = None
 
-        # register agent
-        try:
-            self.model.agents_[type(self)][self] = None
-        except AttributeError as err:
-            # model super has not been called
-            raise RuntimeError(
-                "The Mesa Model class was not initialized. You must explicitly initialize the Model by calling super().__init__() on initialization."
-            ) from err
+        self.model.register_agent(self)
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
         with contextlib.suppress(KeyError):
-            self.model.agents_[type(self)].pop(self)
+            self.model.deregister_agent(self)
 
     def step(self) -> None:
         """A single step of the agent."""
@@ -316,6 +309,21 @@ class AgentSet(MutableSet, Sequence):
                 for agent in self._agents
             ]
 
+    def set(self, attr_name: str, value: Any) -> AgentSet:
+        """
+        Set a specified attribute to a given value for all agents in the AgentSet.
+
+        Args:
+            attr_name (str): The name of the attribute to set.
+            value (Any): The value to set the attribute to.
+
+        Returns:
+            AgentSet: The AgentSet instance itself, after setting the attribute.
+        """
+        for agent in self:
+            setattr(agent, attr_name, value)
+        return self
+
     def __getitem__(self, item: int | slice) -> Agent:
         """
         Retrieve an agent or a slice of agents from the AgentSet.
diff --git a/mesa/model.py b/mesa/model.py
index bcd5902f..6261c79b 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -8,9 +8,8 @@ Core Objects: Model
 # Remove this __future__ import once the oldest supported Python is 3.10
 from __future__ import annotations
 
-import itertools
 import random
-from collections import defaultdict
+import warnings
 
 # mypy
 from typing import Any
@@ -32,11 +31,9 @@ class Model:
         running: A boolean indicating if the model should continue running.
         schedule: An object to manage the order and execution of agent steps.
         current_id: A counter for assigning unique IDs to agents.
-        agents_: A defaultdict mapping each agent type to a dict of its instances.
-                 This private attribute is used internally to manage agents.
 
     Properties:
-        agents: An AgentSet containing all agents in the model, generated from the _agents attribute.
+        agents: An AgentSet containing all agents in the model
         agent_types: A list of different agent types present in the model.
 
     Methods:
@@ -46,6 +43,14 @@ class Model:
         next_id: Generates and returns the next unique identifier for an agent.
         reset_randomizer: Resets the model's random number generator with a new or existing seed.
         initialize_data_collector: Sets up the data collector for the model, requiring an initialized scheduler and agents.
+        register_agent : register an agent with the model
+        deregister_agent : remove an agent from the model
+
+    Notes:
+        Model.agents returns the AgentSet containing all agents registered with the model. Changing
+        the content of the AgentSet directly can result in strange behavior. If you want change the
+        composition of this AgentSet, ensure you operate on a copy.
+
     """
 
     def __new__(cls, *args: Any, **kwargs: Any) -> Any:
@@ -57,6 +62,7 @@ class Model:
             # advance.
             obj._seed = random.random()
         obj.random = random.Random(obj._seed)
+
         # TODO: Remove these 2 lines just before Mesa 3.0
         obj._steps = 0
         obj._time = 0
@@ -71,7 +77,8 @@ class Model:
         self.running = True
         self.schedule = None
         self.current_id = 0
-        self.agents_: defaultdict[type, dict] = defaultdict(dict)
+
+        self._setup_agent_registration()
 
         self._steps: int = 0
         self._time: TimeT = 0  # the model's clock
@@ -79,31 +86,89 @@ class Model:
     @property
     def agents(self) -> AgentSet:
         """Provides an AgentSet of all agents in the model, combining agents from all types."""
-
-        if hasattr(self, "_agents"):
-            return self._agents
-        else:
-            all_agents = itertools.chain.from_iterable(self.agents_.values())
-            return AgentSet(all_agents, self)
+        return self._all_agents
 
     @agents.setter
     def agents(self, agents: Any) -> None:
         raise AttributeError(
-            "You are trying to set model.agents. In Mesa 3.0 and higher, this attribute will be "
+            "You are trying to set model.agents. In Mesa 3.0 and higher, this attribute is "
             "used by Mesa itself, so you cannot use it directly anymore."
             "Please adjust your code to use a different attribute name for custom agent storage."
         )
 
-        self._agents = agents
-
     @property
     def agent_types(self) -> list[type]:
-        """Return a list of different agent types."""
-        return list(self.agents_.keys())
+        """Return a list of all unique agent types registered with the model."""
+        return list(self._agents_by_type.keys())
 
     def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
-        """Retrieves an AgentSet containing all agents of the specified type."""
-        return AgentSet(self.agents_[agenttype].keys(), self)
+        """Retrieves an AgentSet containing all agents of the specified type.
+
+        Args:
+            agenttype: The type of agent to retrieve.
+
+        Raises:
+            KeyError: If agenttype does not exist
+
+
+        """
+        return self._agents_by_type[agenttype]
+
+    def _setup_agent_registration(self):
+        """helper method to initialize the agent registration datastructures"""
+        self._agents = {}  # the hard references to all agents in the model
+        self._agents_by_type: dict[
+            type, AgentSet
+        ] = {}  # a dict with an agentset for each class of agents
+        self._all_agents = AgentSet([], self)  # an agenset with all agents
+
+    def register_agent(self, agent):
+        """Register the agent with the model
+
+        Args:
+            agent: The agent to register.
+
+        Notes:
+            This method is called automatically by ``Agent.__init__``, so there is no need to use this
+            if you are subclassing Agent and calling its super in the ``__init__`` method.
+
+        """
+        if not hasattr(self, "_agents"):
+            self._setup_agent_registration()
+
+            warnings.warn(
+                "The Mesa Model class was not initialized. In the future, you need to explicitly initialize "
+                "the Model by calling super().__init__() on initialization.",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        self._agents[agent] = None
+
+        # because AgentSet requires model, we cannot use defaultdict
+        # tricks with a function won't work because model then cannot be pickled
+        try:
+            self._agents_by_type[type(agent)].add(agent)
+        except KeyError:
+            self._agents_by_type[type(agent)] = AgentSet(
+                [
+                    agent,
+                ],
+                self,
+            )
+
+        self._all_agents.add(agent)
+
+    def deregister_agent(self, agent):
+        """Deregister the agent with the model
+
+        Notes::
+        This method is called automatically by ``Agent.remove``
+
+        """
+        del self._agents[agent]
+        self._agents_by_type[type(agent)].remove(agent)
+        self._all_agents.remove(agent)
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 516902f9..8c5bf351 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -270,6 +270,30 @@ def test_agentset_do_callable():
     assert len(agentset) == 0
 
 
+def test_agentset_set_method():
+    # Initialize the model and agents with and without existing attributes
+    class TestAgentWithAttribute(Agent):
+        def __init__(self, unique_id, model, age=None):
+            super().__init__(unique_id, model)
+            self.age = age
+
+    model = Model()
+    agents = [TestAgentWithAttribute(model.next_id(), model, age=i) for i in range(5)]
+    agentset = AgentSet(agents, model)
+
+    # Set a new attribute "health" and an existing attribute "age" for all agents
+    agentset.set("health", 100).set("age", 50).set("status", "active")
+
+    # Check if all agents have the "health", "age", and "status" attributes correctly set
+    for agent in agentset:
+        assert hasattr(agent, "health")
+        assert agent.health == 100
+        assert hasattr(agent, "age")
+        assert agent.age == 50
+        assert hasattr(agent, "status")
+        assert agent.status == "active"
+
+
 def test_agentset_map_str():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]