remove_all_agents method added to model

mesa/agent.py

@@ -67,7 +67,13 @@ def __init__(self, model: Model, *args, **kwargs) -> None:
         self.model.register_agent(self)
 
     def remove(self) -> None:
-        """Remove and delete the agent from the model."""
+        """Remove and delete the agent from the model.
+
+        Notes:
+            If you need to do additional cleanup when removing an agent by for example removing
+            it from a space, consider extending this method in your own agent class.
+
+        """
         with contextlib.suppress(KeyError):
             self.model.deregister_agent(self)
 

mesa/model.py

@@ -276,3 +276,16 @@ def initialize_data_collector(
         )
         # Collect data for the first time during initialization.
         self.datacollector.collect(self)
+
+    def remove_all_agents(self):
+        """Remove all agents from the model.
+
+        Notes:
+            This method calls agent.remove for all agents in the model. If you need to remove agents from
+            e.g., a SingleGrid, you can either explicitly implement your own agent.remove method or clean this up
+            near where you are calling this method.
+
+        """
+        # we need to wrap keys in a list to avoid a RunTimeError: dictionary changed size during iteration
+        for agent in list(self._agents.keys()):
+            agent.remove()

tests/test_model.py

@@ -93,3 +93,18 @@ class Sheep(Agent):
     assert model.agents_by_type[Wolf] == AgentSet([wolf], model)
     assert model.agents_by_type[Sheep] == AgentSet([sheep], model)
     assert len(model.agents_by_type) == 2
+
+
+def test_agent_remove():
+    """Test removing all agents from the model."""
+
+    class TestAgent(Agent):
+        pass
+
+    model = Model()
+    for _ in range(100):
+        TestAgent(model)
+    assert len(model.agents) == 100
+
+    model.remove_all_agents()
+    assert len(model.agents) == 0