Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce AgentSet class #1916

Merged
merged 25 commits into from
Dec 23, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
3bb8dde
Introduce AgentSet class
EwoutH Dec 19, 2023
51d293b
move to WeakKeyDictionary, add inplace boolean, do_each can return re…
quaquel Dec 19, 2023
e397f9b
adds __getitem__, get_each, and code cleanup
quaquel Dec 19, 2023
8d31c3b
removal of set
quaquel Dec 20, 2023
ef76309
AgentSet subclasses abc.MutableSet and supports indexing and slicing
quaquel Dec 20, 2023
7f7ca58
unittests for AgentSet
quaquel Dec 20, 2023
ec36b9f
Update mesa/agent.py
quaquel Dec 20, 2023
1531ec5
Update mesa/agent.py
quaquel Dec 20, 2023
d34eef9
make AgentSet pickle-able
quaquel Dec 20, 2023
1e3b4d9
minor change to __setstate__ for pickleability
quaquel Dec 20, 2023
5c279b4
fix for pickle test
quaquel Dec 20, 2023
c1e64f6
additional keyword arguments for sort, select, renaming of some other…
quaquel Dec 21, 2023
d26c483
mimic how Agents handles random
quaquel Dec 21, 2023
2436e5b
fix for typo in docstring
quaquel Dec 21, 2023
503e05a
Black formatting
EwoutH Dec 21, 2023
3c42e50
Ruff fixes
EwoutH Dec 21, 2023
a513d74
Fix last ruff errors
EwoutH Dec 21, 2023
f298b7d
Model: Update docstring
EwoutH Dec 21, 2023
5c770a6
AgentSet: Add docstring
EwoutH Dec 21, 2023
175e2c4
tests: Add more tests for AgentSet
EwoutH Dec 21, 2023
daf8a23
AgentSet: Add agent_type argument to select() method
EwoutH Dec 21, 2023
a26e7c6
AgentSet: Rename reverse to ascending in sort
EwoutH Dec 21, 2023
8d2a9d2
Add tests for selecting by type
EwoutH Dec 21, 2023
7888007
Add experimental warning for AgentSet
EwoutH Dec 21, 2023
22b99cd
requested fixes
quaquel Dec 22, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 45 additions & 4 deletions mesa/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from random import Random

# mypy
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any, Callable, Iterator

if TYPE_CHECKING:
# We ensure that these are not imported during runtime to prevent cyclic
Expand Down Expand Up @@ -41,12 +41,15 @@
self.model = model
self.pos: Position | None = None

# Register the agent with the model using defaultdict
self.model.agents[type(self)][self] = None
# Directly register the agent with the model
if type(self) not in self.model._agents:
self.model._agents[type(self)] = set()
EwoutH marked this conversation as resolved.
Show resolved Hide resolved
self.model._agents[type(self)].add(self)

def remove(self) -> None:
"""Remove and delete the agent from the model."""
self.model.agents[type(self)].pop(self)
if type(self) in self.model._agents:
self.model._agents[type(self)].discard(self)

def step(self) -> None:
"""A single step of the agent."""
Expand All @@ -57,3 +60,41 @@
@property
def random(self) -> Random:
return self.model.random


class AgentSet:
def __init__(self, agents: set[Agent], model: Model):
self._agents = agents
EwoutH marked this conversation as resolved.
Show resolved Hide resolved
self.model = model
EwoutH marked this conversation as resolved.
Show resolved Hide resolved

def __len__(self):
return len(self._agents)

Check warning on line 71 in mesa/agent.py

View check run for this annotation

Codecov / codecov/patch

mesa/agent.py#L71

Added line #L71 was not covered by tests

def __iter__(self) -> Iterator[Agent]:
return iter(self._agents)

Check warning on line 74 in mesa/agent.py

View check run for this annotation

Codecov / codecov/patch

mesa/agent.py#L74

Added line #L74 was not covered by tests

def __contains__(self, agent: Agent) -> bool:
"""Check if an agent is in the AgentSet."""
return agent in self._agents

def select(self, filter_func: Callable[[Agent], bool] | None = None) -> AgentSet:
if filter_func is None:
return AgentSet(set(self._agents), self.model)

Check warning on line 82 in mesa/agent.py

View check run for this annotation

Codecov / codecov/patch

mesa/agent.py#L82

Added line #L82 was not covered by tests
return AgentSet(
{agent for agent in self._agents if filter_func(agent)}, self.model
)

def shuffle(self) -> AgentSet:
shuffled_agents = list(self._agents)
self.model.random.shuffle(shuffled_agents)
return AgentSet(set(shuffled_agents), self.model)

Check warning on line 90 in mesa/agent.py

View check run for this annotation

Codecov / codecov/patch

mesa/agent.py#L88-L90

Added lines #L88 - L90 were not covered by tests
EwoutH marked this conversation as resolved.
Show resolved Hide resolved

def sort(self, key: Callable[[Agent], Any], reverse: bool = False) -> AgentSet:
sorted_agents = sorted(self._agents, key=key, reverse=reverse)
return AgentSet(set(sorted_agents), self.model)

Check warning on line 94 in mesa/agent.py

View check run for this annotation

Codecov / codecov/patch

mesa/agent.py#L93-L94

Added lines #L93 - L94 were not covered by tests
EwoutH marked this conversation as resolved.
Show resolved Hide resolved

def do_each(self, method_name: str):
for agent in self._agents:
getattr(agent, method_name)()

Check warning on line 98 in mesa/agent.py

View check run for this annotation

Codecov / codecov/patch

mesa/agent.py#L98

Added line #L98 was not covered by tests
EwoutH marked this conversation as resolved.
Show resolved Hide resolved

# Additional methods like union, intersection, difference, etc., can be added as needed.
15 changes: 13 additions & 2 deletions mesa/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
# mypy
from typing import Any

from mesa.agent import AgentSet
from mesa.datacollection import DataCollector


Expand Down Expand Up @@ -51,12 +52,22 @@
self.running = True
self.schedule = None
self.current_id = 0
self.agents: defaultdict[type, dict] = defaultdict(dict)
self._agents: defaultdict[type, dict] = defaultdict(dict)

@property
def agents(self) -> AgentSet:
all_agents = set()
EwoutH marked this conversation as resolved.
Show resolved Hide resolved
for agent_type in self._agents:
all_agents.update(self._agents[agent_type])
return AgentSet(all_agents, self)

@property
def agent_types(self) -> list:
"""Return a list of different agent types."""
return list(self.agents.keys())
return list(self._agents.keys())

def select_agents(self, *args, **kwargs) -> AgentSet:
return self.agents.select(*args, **kwargs)

Check warning on line 70 in mesa/model.py

View check run for this annotation

Codecov / codecov/patch

mesa/model.py#L70

Added line #L70 was not covered by tests

def run_model(self) -> None:
"""Run the model until the end condition is reached. Overload as
Expand Down
4 changes: 2 additions & 2 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@ class TestAgent(Agent):
model = Model()
agent = TestAgent(model.next_id(), model)
# Check if the agent is added
assert agent in model.agents[type(agent)]
assert agent in model.agents

agent.remove()
# Check if the agent is removed
assert agent not in model.agents[type(agent)]
assert agent not in model.agents
2 changes: 1 addition & 1 deletion tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,5 +49,5 @@ class TestAgent(Agent):

model = Model()
test_agent = TestAgent(model.next_id(), model)
assert test_agent in model.agents[type(test_agent)]
assert test_agent in model.agents
assert type(test_agent) in model.agent_types