Repository: projectmesa/mesa
Pull request #2219 (2024-08-17 15:18:25)

Allow AgentSet.do() to take Callable function
This PR enhances `AgentSet.do` to take a callable or str. Currently, AgentSet.do takes a `str` which maps to a method on the agents in the set. This PR makes it possible to also use a `Callable` instead.

This callable will be called with the `agent` as the first argument.

### ⚠️ Breaking change ⚠️ 
A small breaking change is introduced here: the `method_name` parameter is renamed to `method`. For models that use this as a keyword argument this is a breaking change, and need to replace `do(method_name="something")` with `do(method="something")`.

### Usage examples
Here are a few examples of how to use `do()` with `Callable` functions.

Suppose we want to move each agent by a fixed offset:

```python
def move_agent(agent, dx, dy):
    # Assume each agent has a 'pos' attribute which is a tuple (x, y)
    if agent.pos:
        agent.pos = (agent.pos[0] + dx, agent.pos[1] + dy)

# Usage with do()
agentset.do(move_agent, 1, -1)
```

You might want to increase the energy level of all agents by a fixed amount:

```python
agentset.do(lambda agent: setattr(agent, 'energy', agent.energy + 10))
```

You can apply a function conditionally to agents that meet a certain criterion:

```python
def heal_agent(agent):
    if agent.health < 50:
        agent.health += 20

# Apply the healing only to agents with health below 50
agentset.do(lambda agent: heal_agent(agent) if agent.health < 50 else None)
```

You may want to gather results from applying a function to each agent:

```python
def get_agent_status(agent):
    return f"Agent {agent.unique_id}: {agent.status}"

statuses = agentset.do(get_agent_status, return_results=True)
```

If your callable needs additional arguments, you can pass them directly via `do()`:

```python
def update_wealth(agent, amount, tax_rate):
    agent.wealth += amount * (1 - tax_rate)

agentset.do(update_wealth, 100, 0.2)
```

Suppose you want to update each agent's position based on their velocity:

```python
def update_position(agent):
    if agent.pos and hasattr(agent, 'velocity'):
        agent.pos = (agent.pos[0] + agent.velocity[0], agent.pos[1] + agent.velocity[1])

agentset.do(update_position)
```

And of course the current behavior of passing a string as method is still supported:

```python
agentset.do("step")
```

github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]      | Run time [95% CI]       |
|--------------|--------|-------------------------|-------------------------|
| Schelling    | small  | 🔵 +3.0% [+2.7%, +3.3%] | 🔵 +1.2% [+1.0%, +1.4%] |
| Schelling    | large  | 🔴 +3.5% [+3.2%, +3.8%] | 🔵 +0.4% [-0.2%, +1.0%] |
| WolfSheep    | small  | 🔵 +2.5% [+1.7%, +3.4%] | 🔵 +2.7% [+2.4%, +3.0%] |
| WolfSheep    | large  | 🔵 +2.7% [+2.5%, +2.9%] | 🔵 +1.4% [+0.9%, +1.9%] |
| BoidFlockers | small  | 🔵 +1.4% [+0.7%, +2.0%] | 🔵 +0.7% [+0.1%, +1.3%] |
| BoidFlockers | large  | 🔵 +1.4% [+1.0%, +1.9%] | 🔵 -0.4% [-1.1%, +0.3%] |


Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index 2f0e842d..e65b25f6 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -11,9 +11,7 @@ from __future__ import annotations
 import contextlib
 import copy
 import operator
-import warnings
 import weakref
-from collections import defaultdict
 from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
@@ -52,17 +50,11 @@ class Agent:
         # register agent
         try:
             self.model.agents_[type(self)][self] = None
-        except AttributeError:
+        except AttributeError as err:
             # model super has not been called
-            self.model.agents_ = defaultdict(dict)
-            self.model.agents_[type(self)][self] = None
-            self.model.agentset_experimental_warning_given = False
-
-            warnings.warn(
-                "The Mesa Model class was not initialized. In the future, you need to explicitly initialize the Model by calling super().__init__() on initialization.",
-                FutureWarning,
-                stacklevel=2,
-            )
+            raise RuntimeError(
+                "The Mesa Model class was not initialized. You must explicitly initialize the Model by calling super().__init__() on initialization."
+            ) from err
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
@@ -100,8 +92,6 @@ class AgentSet(MutableSet, Sequence):
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
     """
 
-    agentset_experimental_warning_given = False
-
     def __init__(self, agents: Iterable[Agent], model: Model):
         """
         Initializes the AgentSet with a collection of agents and a reference to the model.
@@ -227,26 +217,37 @@ class AgentSet(MutableSet, Sequence):
         return self
 
     def do(
-        self, method_name: str, *args, return_results: bool = False, **kwargs
+        self, method: str | Callable, *args, return_results: bool = False, **kwargs
     ) -> AgentSet | list[Any]:
         """
-        Invoke a method on each agent in the AgentSet.
+        Invoke a method or function on each agent in the AgentSet.
 
         Args:
-            method_name (str): The name of the method to call on each agent.
+            method (str, callable): the callable to do on each agents
+
+                                        * in case of str, the name of the method to call on each agent.
+                                        * in case of callable, the function to be called with each agent as first argument
+
             return_results (bool, optional): If True, returns the results of the method calls; otherwise, returns the AgentSet itself. Defaults to False, so you can chain method calls.
-            *args: Variable length argument list passed to the method being called.
-            **kwargs: Arbitrary keyword arguments passed to the method being called.
+            *args: Variable length argument list passed to the callable being called.
+            **kwargs: Arbitrary keyword arguments passed to the callable being called.
 
         Returns:
-            AgentSet | list[Any]: The results of the method calls if return_results is True, otherwise the AgentSet itself.
+            AgentSet | list[Any]: The results of the callable calls if return_results is True, otherwise the AgentSet itself.
         """
         # we iterate over the actual weakref keys and check if weakref is alive before calling the method
-        res = [
-            getattr(agent, method_name)(*args, **kwargs)
-            for agentref in self._agents.keyrefs()
-            if (agent := agentref()) is not None
-        ]
+        if isinstance(method, str):
+            res = [
+                getattr(agent, method)(*args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
+        else:
+            res = [
+                method(agent, *args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
 
         return res if return_results else self
 
diff --git a/mesa/space.py b/mesa/space.py
index a6d070f0..39433cb3 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -586,7 +586,7 @@ class PropertyLayer:
         aggregate_property(operation): Performs an aggregate operation over all cells.
     """
 
-    agentset_experimental_warning_given = False
+    propertylayer_experimental_warning_given = False
 
     def __init__(
         self, name: str, width: int, height: int, default_value, dtype=np.float64
@@ -633,14 +633,14 @@ class PropertyLayer:
 
         self.data = np.full((width, height), default_value, dtype=dtype)
 
-        if not self.__class__.agentset_experimental_warning_given:
+        if not self.__class__.propertylayer_experimental_warning_given:
             warnings.warn(
                 "The new PropertyLayer and _PropertyGrid classes experimental. It may be changed or removed in any and all future releases, including patch releases.\n"
                 "We would love to hear what you think about this new feature. If you have any thoughts, share them with us here: https://github.com/projectmesa/mesa/discussions/1932",
                 FutureWarning,
                 stacklevel=2,
             )
-            self.__class__.agentset_experimental_warning_given = True
+            self.__class__.propertylayer_experimental_warning_given = True
 
     def set_cell(self, position: Coordinate, value):
         """
diff --git a/mesa/visualization/solara_viz.py b/mesa/visualization/solara_viz.py
index 4f159884..fba73207 100644
--- a/mesa/visualization/solara_viz.py
+++ b/mesa/visualization/solara_viz.py
@@ -70,7 +70,7 @@ def Card(
                 )
             elif space_drawer:
                 # if specified, draw agent space with an alternate renderer
-                space_drawer(model, agent_portrayal)
+                space_drawer(model, agent_portrayal, dependencies=dependencies)
         elif "Measure" in layout_type:
             rv.CardTitle(children=["Measure"])
             measure = measures[layout_type["Measure"]]


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 0cd21112..f4e64ce5 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -176,7 +176,7 @@ def test_agentset_get_item():
         _ = agentset[20]
 
 
-def test_agentset_do_method():
+def test_agentset_do_str():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]
     agentset = AgentSet(agents, model)
@@ -210,6 +210,72 @@ def test_agentset_do_method():
     assert len(agentset) == 0
 
 
+def test_agentset_do_callable():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    # Test callable with non-existent function
+    with pytest.raises(AttributeError):
+        agentset.do(lambda agent: agent.non_existing_method())
+
+    # tests for addition and removal in do using callables
+    # do iterates, so no error should be raised to change size while iterating
+    # related to issue #1595
+
+    # setup for lambda function tests
+    n = 10
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Lambda for addition
+    agentset.do(lambda agent: agent.do_add())
+    assert len(agentset) == 2 * n
+
+    # setup again for lambda function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Lambda for removal
+    agentset.do(lambda agent: agent.do_remove())
+    assert len(agentset) == 0
+
+    # setup for actual function tests
+    def add_function(agent):
+        agent.do_add()
+
+    def remove_function(agent):
+        agent.do_remove()
+
+    # setup again for actual function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Actual function for addition
+    agentset.do(add_function)
+    assert len(agentset) == 2 * n
+
+    # setup again for actual function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Actual function for removal
+    agentset.do(remove_function)
+    assert len(agentset) == 0
+
+
 def test_agentset_get_attribute():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]
diff --git a/tests/test_solara_viz.py b/tests/test_solara_viz.py
index f3422b87..5fef7e0b 100644
--- a/tests/test_solara_viz.py
+++ b/tests/test_solara_viz.py
@@ -136,7 +136,9 @@ def test_call_space_drawer(mocker):
             space_drawer=altspace_drawer,
         )
     )
-    altspace_drawer.assert_called_with(model, agent_portrayal)
+    altspace_drawer.assert_called_with(
+        model, agent_portrayal, dependencies=dependencies
+    )
 
 
 def test_slider():