Performance enhancements for Model.agents (#2251)

This PR is a performance enhancement for Model.agents. It emerged from a discussion on [the weird scaling performance of the Boltzman wealth model](#2224).

# 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.

Author: quaquel

mesa/agent.py

@@ -49,19 +49,12 @@ def __init__(self, unique_id: int, model: Model) -> None:
         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."""

mesa/model.py

@@ -8,9 +8,8 @@
 # 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 @@ def __new__(cls, *args: Any, **kwargs: Any) -> Any:
             # 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,39 +77,98 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         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
 
     @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