Repository: projectmesa/mesa
Pull request #1523 (2022-11-11 10:58:37)

Enhance schedulers to support intra-step removal of agents
As discussed in #1515 the only scheduler which checks for the presence of the agents was the BaseScheduler (and the ones which inherit the method `agent_buffer`), I suppose that this should be done also for the other schedulers except to the SimultaneousActivation one

codecov[bot]: # [Codecov](https://codecov.io/gh/projectmesa/mesa/pull/1523?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) Report
Base: **91.29**% // Head: **91.15**% // Decreases project coverage by **`-0.13%`** :warning:
> Coverage data is based on head [(`407ed24`)](https://codecov.io/gh/projectmesa/mesa/pull/1523?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) compared to base [(`a409675`)](https://codecov.io/gh/projectmesa/mesa/commit/a4096758ed96a41ddd84f4957781388660d44d46?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).
> Patch coverage: 57.14% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1523      +/-   ##
==========================================
- Coverage   91.29%   91.15%   -0.14%     
==========================================
  Files          15       15              
  Lines        1310     1312       +2     
  Branches      227      229       +2     
==========================================
  Hits         1196     1196              
  Misses         80       80              
- Partials       34       36       +2     
```


| [Impacted Files](https://codecov.io/gh/projectmesa/mesa/pull/1523?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) | Coverage Δ | |
|---|---|---|
| [mesa/time.py](https://codecov.io/gh/projectmesa/mesa/pull/1523/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS90aW1lLnB5) | `93.00% <57.14%> (-1.90%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/projectmesa/mesa/pull/1523?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).

rht: L105-107 has existed in the code (was added via #584, which didn't add the tests). I could merge this PR and add the tests later, so that people can use the feature sooner, but at this point I think we should add the tests now.
Tortar: Will put some tests
Tortar: I added some tests :-)
Tortar: I Is it better according to you to have a subclass of MockAgent with a different step instead of how it is done now?
rht: > @rht Is it better according to you to have a subclass of MockAgent with a different step instead of how it is done now?

I think the current implementation is fine. The step logic is still clear even with the addition of the `kill_other_agent` conditional.
rht: LGTM

Patch:
diff --git a/mesa/time.py b/mesa/time.py
index 54f6e25a..4c332e3e 100644
--- a/mesa/time.py
+++ b/mesa/time.py
@@ -61,7 +61,6 @@ class BaseScheduler:
             agent: An Agent to be added to the schedule. NOTE: The agent must
             have a step() method.
         """
-
         if agent.unique_id in self._agents:
             raise Exception(
                 f"Agent with unique id {repr(agent.unique_id)} already added to scheduler"
@@ -102,9 +101,9 @@ class BaseScheduler:
         if shuffled:
             self.model.random.shuffle(agent_keys)
 
-        for key in agent_keys:
-            if key in self._agents:
-                yield self._agents[key]
+        for agent_key in agent_keys:
+            if agent_key in self._agents:
+                yield self._agents[agent_key]
 
 
 class RandomActivation(BaseScheduler):
@@ -197,7 +196,8 @@ class StagedActivation(BaseScheduler):
             self.model.random.shuffle(agent_keys)
         for stage in self.stage_list:
             for agent_key in agent_keys:
-                getattr(self._agents[agent_key], stage)()  # Run stage
+                if agent_key in self._agents:
+                    getattr(self._agents[agent_key], stage)()  # Run stage
             # We recompute the keys because some agents might have been removed
             # in the previous loop.
             agent_keys = list(self._agents.keys())
@@ -239,7 +239,6 @@ class RandomActivationByType(BaseScheduler):
         Args:
             agent: An Agent to be added to the schedule.
         """
-
         super().add(agent)
         agent_class: type[Agent] = type(agent)
         self.agents_by_type[agent_class][agent.unique_id] = agent
@@ -248,7 +247,6 @@ class RandomActivationByType(BaseScheduler):
         """
         Remove all instances of a given agent from the schedule.
         """
-
         del self._agents[agent.unique_id]
 
         agent_class: type[Agent] = type(agent)
@@ -286,7 +284,8 @@ class RandomActivationByType(BaseScheduler):
         if shuffle_agents:
             self.model.random.shuffle(agent_keys)
         for agent_key in agent_keys:
-            self.agents_by_type[type_class][agent_key].step()
+            if agent_key in self.agents_by_type[type_class]:
+                self.agents_by_type[type_class][agent_key].step()
 
     def get_type_count(self, type_class: type[Agent]) -> int:
         """


Test patch:
diff --git a/tests/test_time.py b/tests/test_time.py
index d92ae04c..7bf198e9 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -29,7 +29,15 @@ class MockAgent(Agent):
         self.steps = 0
         self.advances = 0
 
+    def kill_other_agent(self):
+        for agent in self.model.schedule.agents:
+            if agent is not self:
+                self.model.schedule.remove(agent)
+                break
+
     def stage_one(self):
+        if self.model.enable_kill_other_agent:
+            self.kill_other_agent()
         self.model.log.append(self.unique_id + "_1")
 
     def stage_two(self):
@@ -39,11 +47,14 @@ class MockAgent(Agent):
         self.advances += 1
 
     def step(self):
+        if self.model.enable_kill_other_agent:
+            self.kill_other_agent()
         self.steps += 1
+        self.model.log.append(self.unique_id)
 
 
 class MockModel(Model):
-    def __init__(self, shuffle=False, activation=STAGED):
+    def __init__(self, shuffle=False, activation=STAGED, enable_kill_other_agent=False):
         """
         Creates a Model instance with a schedule
 
@@ -59,6 +70,7 @@ class MockModel(Model):
                               The default scheduler is a BaseScheduler.
         """
         self.log = []
+        self.enable_kill_other_agent = enable_kill_other_agent
 
         # Make scheduler
         if activation == STAGED:
@@ -91,7 +103,7 @@ class TestStagedActivation(TestCase):
 
     def test_no_shuffle(self):
         """
-        Testing staged activation without shuffling.
+        Testing the staged activation without shuffling.
         """
         model = MockModel(shuffle=False)
         model.step()
@@ -100,7 +112,7 @@ class TestStagedActivation(TestCase):
 
     def test_shuffle(self):
         """
-        Test staged activation with shuffling
+        Test the staged activation with shuffling
         """
         model = MockModel(shuffle=True)
         model.step()
@@ -118,7 +130,7 @@ class TestStagedActivation(TestCase):
 
     def test_remove(self):
         """
-        Test staged activation can remove an agent
+        Test the staged activation can remove an agent
         """
         model = MockModel(shuffle=True)
         agent_keys = list(model.schedule._agents.keys())
@@ -126,6 +138,15 @@ class TestStagedActivation(TestCase):
         model.schedule.remove(agent)
         assert agent not in model.schedule.agents
 
+    def test_intrastep_remove(self):
+        """
+        Test the staged activation can remove an agent in a
+        step of another agent so that the one removed doesn't step.
+        """
+        model = MockModel(shuffle=True, enable_kill_other_agent=True)
+        model.step()
+        assert len(model.log) == 2
+
     def test_add_existing_agent(self):
         model = MockModel()
         agent = model.schedule.agents[0]
@@ -162,13 +183,21 @@ class TestRandomActivation(TestCase):
         """
         Test the random activation step causes each agent to step
         """
-
         model = MockModel(activation=RANDOM)
         model.step()
         agent_steps = [i.steps for i in model.schedule.agents]
         # one step for each of 2 agents
         assert all(map(lambda x: x == 1, agent_steps))
 
+    def test_intrastep_remove(self):
+        """
+        Test the random activation can remove an agent in a
+        step of another agent so that the one removed doesn't step.
+        """
+        model = MockModel(activation=RANDOM, enable_kill_other_agent=True)
+        model.step()
+        assert len(model.log) == 1
+
 
 class TestSimultaneousActivation(TestCase):
     """
@@ -233,8 +262,8 @@ class TestRandomActivationByType(TestCase):
         RandomActivationByType.
         """
         model = MockModel(activation=RANDOM_BY_TYPE)
-        a = MockAgent(0, None)
-        b = MockAgent(0, None)
+        a = MockAgent(0, model)
+        b = MockAgent(0, model)
         model.schedule.add(a)
         with self.assertRaises(Exception):
             model.schedule.add(b)