Track agents in model with defaultdict

- The `Model` class now uses a `defaultdict` to store agents, ensuring a set is automatically created for each new agent type.
- The `Agent` class has been updated to leverage this feature, simplifying the registration process when an agent is created.
- The `remove` method in the `Agent` class now uses `discard`, providing a safer way to remove agents from the model.

mesa/agent.py

@@ -25,7 +25,6 @@ class Agent:
 
     Attributes:
         unique_id (int): A unique identifier for this agent.
-        agent_type (type): The class of the agent, used for type categorization.
         model (Model): A reference to the model instance.
         self.pos: Position | None = None
     """
@@ -39,10 +38,17 @@ def __init__(self, unique_id: int, model: Model) -> None:
             model (Model): The model instance in which the agent exists.
         """
         self.unique_id = unique_id
-        self.agent_type = self.__class__
         self.model = model
         self.pos: Position | None = None
 
+        # Register the agent with the model using defaultdict
+        self.model.agents[type(self)].add(self)
+
+    def remove(self) -> None:
+        """Remove and delete the agent from the model."""
+        self.model.agents[type(self)].discard(self)
+        del self
+
     def step(self) -> None:
         """A single step of the agent."""
 

mesa/model.py

@@ -8,6 +8,8 @@
 from __future__ import annotations
 
 import random
+from collections import defaultdict
+from typing import Any, DefaultDict, Type
 
 # mypy
 from typing import Any
@@ -41,29 +43,12 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.running = True
         self.schedule = None
         self.current_id = 0
-        self.agents_by_type = {}
+        self.agents: DefaultDict[Type, set] = defaultdict(set)
 
-    def create_agent(self, agent_type, *args, **kwargs):
-        """
-        Creates an agent of the given type with a unique ID and adds it to the agents_by_type dictionary.
-
-        Parameters:
-        agent_type (class): The class of the agent to be created.
-        *args, **kwargs: Additional arguments passed to the agent's constructor.
-
-        Returns:
-        An instance of the specified agent_type.
-        """
-        # Create a new agent instance with a unique ID
-        agent = agent_type(self.next_id(), self, *args, **kwargs)
-
-        # Add the agent to the agents_by_type dictionary
-        agent_class = agent.__class__
-        if agent_class not in self.agents_by_type:
-            self.agents_by_type[agent_class] = []
-        self.agents_by_type[agent_class].append(agent)
-
-        return agent
+    @property
+    def agent_types(self) -> list:
+        """Return a list of different agent types."""
+        return list(self.agents.keys())
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as