Skip to content

Commit 9db99c9

Browse files
Ben-geoadamamer20
andauthored
Auto step counter (#117)
This pull request introduces a new step counter feature to the `ModelDF` class and adds corresponding tests. The main changes include adding a step counter, modifying the `step` method to increment the counter, and adding tests to verify the new functionality. Enhancements to `ModelDF` class: * [`mesa_frames/concrete/model.py`](diffhunk://#diff-c5db932a22d9e3b30ff7eead6d38e655cf3be4fee195a6fe18409bbe3c575ddcR88-R101): Introduced a new private attribute `_steps` to keep track of the step count and added a property `steps` to retrieve the current step count. Modified the `step` method to increment the step counter automatically by wrapping the user-defined `step` method. [[1]](diffhunk://#diff-c5db932a22d9e3b30ff7eead6d38e655cf3be4fee195a6fe18409bbe3c575ddcR88-R101) [[2]](diffhunk://#diff-c5db932a22d9e3b30ff7eead6d38e655cf3be4fee195a6fe18409bbe3c575ddcR150-R160) Testing the new step counter: * [`tests/test_modeldf.py`](diffhunk://#diff-2be7ad75672aa528662114edd128ae28ae41de2919b9e1740388a914a74122ffR1-R37): Added a new test class `Test_ModelDF` to verify the functionality of the step counter. This includes tests for the default step counter behavior and custom user-defined step behavior. --------- Co-authored-by: Adam Amer <136176500+adamamer20@users.noreply.github.com>
1 parent 3e4d349 commit 9db99c9

File tree

2 files changed

+62
-0
lines changed

2 files changed

+62
-0
lines changed

mesa_frames/concrete/model.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,20 @@ def __init__(self, seed: int | Sequence[int] | None = None) -> None:
8585
self.current_id = 0
8686
self._agents = AgentsDF(self)
8787
self._space = None
88+
self._steps = 0
89+
90+
self._user_step = self.step
91+
self.step = self._wrapped_step
92+
93+
def _wrapped_step(self) -> None:
94+
"""Automatically increments step counter and calls user-defined step()."""
95+
self._steps += 1
96+
self._user_step()
97+
98+
@property
99+
def steps(self) -> int:
100+
"""Get the current step count."""
101+
return self._steps
88102

89103
def get_agents_of_type(self, agent_type: type) -> "AgentSetDF":
90104
"""Retrieve the AgentSetDF of a specified type.
@@ -133,6 +147,17 @@ def step(self) -> None:
133147
"""
134148
self.agents.step()
135149

150+
@property
151+
def steps(self) -> int:
152+
"""Get the current step count.
153+
154+
Returns
155+
-------
156+
int
157+
The current step count of the model.
158+
"""
159+
return self._steps
160+
136161
@property
137162
def agents(self) -> AgentsDF:
138163
"""Get the AgentsDF object containing all agents in the model.

tests/test_modeldf.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
from mesa_frames import ModelDF
2+
3+
4+
class CustomModel(ModelDF):
5+
def __init__(self):
6+
super().__init__()
7+
self.custom_step_count = 0
8+
9+
def step(self):
10+
self.custom_step_count += 2
11+
12+
13+
class Test_ModelDF:
14+
def test_steps(self):
15+
model = ModelDF()
16+
17+
assert model.steps == 0
18+
19+
model.step()
20+
assert model.steps == 1
21+
22+
model.step()
23+
assert model.steps == 2
24+
25+
def test_user_defined_step(self):
26+
model = CustomModel()
27+
28+
assert model.steps == 0
29+
assert model.custom_step_count == 0
30+
31+
model.step()
32+
assert model.steps == 1
33+
assert model.custom_step_count == 2
34+
35+
model.step()
36+
assert model.steps == 2
37+
assert model.custom_step_count == 4

0 commit comments

Comments
 (0)