diff --git a/mesa/agent.py b/mesa/agent.py index 9ddd0ebf45f..5ed2c054ab2 100644 --- a/mesa/agent.py +++ b/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.""" diff --git a/mesa/model.py b/mesa/model.py index d1d4e60a6df..bd75e7141fb 100644 --- a/mesa/model.py +++ b/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