Repository: projectmesa/mesa
Pull request #2237 (2024-08-21 17:43:50)

Split AgentSet into map and do to seperate return types
# Motivation
In #2220 @corvince suggested to seperate `AgentSet.do` into 2 seperate methods based on their return type.  `AgentSet.do` executes the callable/method/function and returns the original AgentSet. `AgentSet.map`executes the callable/method/function and returns the results. At the moment both behaviors are in AgentSet.do and can be controlled with the `return_results` keyword argument. 

# Implementation details.
The code is straightforward. I wrap all kwargs to `AgentSet.do` into a dict. Check if `return_results` is in it. If it is, a DepractionWarning is issued. The rest of the code proceeds as normal. 


# Usage

```python

# some form of staged activation
# because do returns the agentset we can chain calls
some_agentset.do("step1").do("step2").do("step3")

# map returns the result of the method so here a list of unique_ids is returned
# no further method chaining of AgentSet methods is possible. 
some_agentset.map(lambda x: x.unique_id)

```

github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]      | Run time [95% CI]       |
|--------------|--------|-------------------------|-------------------------|
| Schelling    | small  | 🔵 +0.2% [-0.2%, +0.5%] | 🔵 -1.1% [-1.2%, -0.9%] |
| Schelling    | large  | 🔵 +0.3% [-0.2%, +0.8%] | 🔵 -1.0% [-2.6%, +0.4%] |
| WolfSheep    | small  | 🔵 +0.3% [-0.8%, +1.4%] | 🔵 -0.5% [-0.8%, -0.1%] |
| WolfSheep    | large  | 🔵 +0.2% [-0.0%, +0.5%] | 🔵 -0.6% [-1.6%, +0.4%] |
| BoidFlockers | small  | 🔵 -3.2% [-3.8%, -2.6%] | 🔵 -3.1% [-3.6%, -2.6%] |
| BoidFlockers | large  | 🟢 -4.0% [-4.7%, -3.3%] | 🟢 -4.4% [-5.3%, -3.6%] |

github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]      | Run time [95% CI]       |
|--------------|--------|-------------------------|-------------------------|
| Schelling    | small  | 🔵 -1.1% [-1.5%, -0.8%] | 🔵 +1.1% [+0.8%, +1.3%] |
| Schelling    | large  | 🔵 +0.1% [-0.4%, +0.7%] | 🔵 +5.6% [+2.9%, +8.3%] |
| WolfSheep    | small  | 🔵 -0.2% [-1.5%, +1.0%] | 🔵 +0.4% [+0.1%, +0.7%] |
| WolfSheep    | large  | 🔵 +0.6% [+0.2%, +1.1%] | 🔵 +2.0% [+0.8%, +3.2%] |
| BoidFlockers | small  | 🔵 +1.6% [+0.8%, +2.3%] | 🔵 +0.9% [+0.1%, +1.6%] |
| BoidFlockers | large  | 🔵 +0.8% [+0.3%, +1.3%] | 🔵 +0.4% [-0.1%, +0.9%] |

rht: While this PR simplifies the typing and type checking, what about consistency with [pandas](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.apply.html)/[Polars](https://docs.pola.rs/api/python/version/0.18/reference/dataframe/api/polars.DataFrame.apply.html) behavior? It seems that Polars' is simplest in that it only allows returning a DF. This PR's `do` and `apply` combine into pandas' `apply`, because the latter has multiple return types.
EwoutH: In general I'm in favor of splitting methods when they have different outputs. I think in this case I also support it.

> `some_agentset.apply(lambda x: x.unique_id)`

I almost feel like AgentSet should be the input to some function here. Like `f(agentset) = ...`, instead of `agentset.apply(f) = ...`

<hr>

On a totally different note, I'm taking a few days of Mesa until Monday. Don't let my absence block anything.
quaquel: > While this PR simplifies the typing and type checking, what about consistency with [pandas](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.apply.html)/[Polars](https://docs.pola.rs/api/python/version/0.18/reference/dataframe/api/polars.DataFrame.apply.html) behavior? It seems that Polars' is simplest in that it only allows returning a DF. This PR's `do` and `apply` combine into pandas' `apply` behavior, because the latter has multiple return types.

MESA is not some DataFrame library but we can draw inspiration from their API. While working with dataframes, it makes sense that each operation returns a new dataframe. You are manipulating data. In MESA, sometimes we want to manipulate agents (so we use `do` in this PR), sometimes we want to get the return from this manipulation (so we use `apply` in this PR). 

For a further motivation for this PR, please check [this comment](https://github.com/projectmesa/mesa/pull/2220#issuecomment-2298521797) which gives a clear example of how confusing method chaining can become. So rather than controlling the return type with a keyword argument, it seems easier to split the two return types into two seperate methods. This makes it easier to see what is going on. 

Note that I also use the exact same seperation between `apply` and `do` in the [groupby PR](https://github.com/projectmesa/mesa/pull/2220#issue-2471564909). This makes the MESA API internally consistent.

Patch:
diff --git a/.github/workflows/build_lint.yml b/.github/workflows/build_lint.yml
index 3a16a605..ebd0258b 100644
--- a/.github/workflows/build_lint.yml
+++ b/.github/workflows/build_lint.yml
@@ -33,6 +33,8 @@ jobs:
         os: [windows, ubuntu, macos]
         python-version: ["3.12"]
         include:
+          - os: ubuntu
+            python-version: "3.13"
           - os: ubuntu
             python-version: "3.11"
           - os: ubuntu
@@ -47,12 +49,16 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+        allow-prereleases: true
         cache: 'pip'
     - name: Install uv
       run: pip install uv
-    - name: Install Mesa
-      # See https://github.com/astral-sh/uv/issues/1945
+    - name: Install Mesa with uv pip
       run: uv pip install --system .[dev]
+      if: matrix.python-version != '3.13'
+    - name: Install Mesa with pip
+      run: pip install .[dev]
+      if: matrix.python-version == '3.13'
     - name: Test with pytest
       run: pytest --durations=10 --cov=mesa tests/ --cov-report=xml
     - if: matrix.os == 'ubuntu'
diff --git a/HISTORY.md b/HISTORY.md
index 73dd5834..91d4c703 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -1,6 +1,37 @@
 ---
 title: Release History
 ---
+# 3.0.0a2 (2024-08-21)
+## Highlights
+In Mesa 3.0 alpha 2 (`v3.0.0a2`) we've done more clean-up in preparation for Mesa 3.0. We now [require](https://github.com/projectmesa/mesa/pull/2218) `super().__init__()`  to run on initializing a Mesa model subclass, `Model.agents` is now fully reserved for the Model's internal AgentSet and we fixed a bug in our Solara space_drawer.
+
+A new feature was added in [#2219](https://github.com/projectmesa/mesa/pull/2219), which now also allows `AgentSet.do()` to take any callable function, instead of only a string referencing to an Agent method. The argument name was changed from `method_name` to `method`.
+
+The new Solara visualisation now allows portraying agents in different shapes, checkout some examples in [#2214](https://github.com/projectmesa/mesa/pull/2214).
+
+We're also working hard on our [example models](https://github.com/projectmesa/mesa-examples). All model warnings were [resolved](https://github.com/projectmesa/mesa-examples/pull/153) and we've [replaced](https://github.com/projectmesa/mesa-examples/pull/161) a lot of schedulers with simpler and more flexible AgentSet functionality. Checkout our [open issues](https://github.com/projectmesa/mesa-examples/issues) if you want to help improve our example models further!
+
+## What's Changed
+### ⚠️ Breaking changes
+* breaking: Add dependencies argument to custom space_drawer by @rht in https://github.com/projectmesa/mesa/pull/2209
+* Require Mesa models to be initialized with `super().__init__()` by @EwoutH in https://github.com/projectmesa/mesa/pull/2218
+* Allow AgentSet.do() to take Callable function by @quaquel in https://github.com/projectmesa/mesa/pull/2219
+* Change warning when setting model.agents to error by @EwoutH in https://github.com/projectmesa/mesa/pull/2225
+### 🧪 Experimental features
+* devs/eventlist: Add repr method to print EventList pretty by @EwoutH in https://github.com/projectmesa/mesa/pull/2195
+### 🛠 Enhancements made
+* Visualisation: Allow specifying Agent shapes in agent_portrayal by @rmhopkins4 in https://github.com/projectmesa/mesa/pull/2214
+### 📜 Documentation improvements
+* docs/conf.py: Use modern `intersphinx_mapping` format by @EwoutH in https://github.com/projectmesa/mesa/pull/2206
+* docs: Update Readme and tutorials to mention Mesa 3.0 pre-releases by @EwoutH in https://github.com/projectmesa/mesa/pull/2203
+### 🔧 Maintenance
+* CI: Let pytest treat warnings as errors for examples by @EwoutH in https://github.com/projectmesa/mesa/pull/2204
+
+## New Contributors
+* @rmhopkins4 made their first contribution in https://github.com/projectmesa/mesa/pull/2214
+
+**Full Changelog**: https://github.com/projectmesa/mesa/compare/v3.0.0a1...v3.0.0a2
+
 # 3.0.0a1 (2024-08-01)
 ## Highlights
 Mesa 3.0 alpha 1 (`v3.0.0a1`) is another step towards our next major version. This release introduces a name change from JupyterViz (jupyter_viz) to SolaraViz (solara_viz), to better represent the tech stack being used. It also includes two bugfixes also present in 2.3.2.
diff --git a/mesa/__init__.py b/mesa/__init__.py
index 1c362f45..50d6b300 100644
--- a/mesa/__init__.py
+++ b/mesa/__init__.py
@@ -24,7 +24,7 @@ __all__ = [
 ]
 
 __title__ = "mesa"
-__version__ = "3.0.0a1"
+__version__ = "3.0.0a2"
 __license__ = "Apache 2.0"
 _this_year = datetime.datetime.now(tz=datetime.timezone.utc).date().year
 __copyright__ = f"Copyright {_this_year} Project Mesa Team"
diff --git a/mesa/agent.py b/mesa/agent.py
index e65b25f6..88056731 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -11,6 +11,7 @@ from __future__ import annotations
 import contextlib
 import copy
 import operator
+import warnings
 import weakref
 from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
@@ -216,25 +217,64 @@ class AgentSet(MutableSet, Sequence):
         self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
         return self
 
-    def do(
-        self, method: str | Callable, *args, return_results: bool = False, **kwargs
-    ) -> AgentSet | list[Any]:
+    def do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
         """
         Invoke a method or function on each agent in the AgentSet.
 
         Args:
-            method (str, callable): the callable to do on each agents
+            method (str, callable): the callable to do on each agent
 
                                         * 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 callable being called.
             **kwargs: Arbitrary keyword arguments passed to the callable being called.
 
         Returns:
             AgentSet | list[Any]: The results of the callable calls if return_results is True, otherwise the AgentSet itself.
         """
+        try:
+            return_results = kwargs.pop("return_results")
+        except KeyError:
+            return_results = False
+        else:
+            warnings.warn(
+                "Using return_results is deprecated. Use AgenSet.do in case of return_results=False, and "
+                "AgentSet.map in case of return_results=True",
+                stacklevel=2,
+            )
+
+        if return_results:
+            return self.map(method, *args, **kwargs)
+
+        # we iterate over the actual weakref keys and check if weakref is alive before calling the method
+        if isinstance(method, str):
+            for agentref in self._agents.keyrefs():
+                if (agent := agentref()) is not None:
+                    getattr(agent, method)(*args, **kwargs)
+        else:
+            for agentref in self._agents.keyrefs():
+                if (agent := agentref()) is not None:
+                    method(agent, *args, **kwargs)
+
+        return self
+
+    def map(self, method: str | Callable, *args, **kwargs) -> list[Any]:
+        """
+        Invoke a method or function on each agent in the AgentSet and return the results.
+
+        Args:
+            method (str, callable): the callable to apply on each agent
+
+                                        * 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
+
+            *args: Variable length argument list passed to the callable being called.
+            **kwargs: Arbitrary keyword arguments passed to the callable being called.
+
+        Returns:
+           list[Any]: The results of the callable calls
+        """
         # we iterate over the actual weakref keys and check if weakref is alive before calling the method
         if isinstance(method, str):
             res = [
@@ -249,7 +289,7 @@ class AgentSet(MutableSet, Sequence):
                 if (agent := agentref()) is not None
             ]
 
-        return res if return_results else self
+        return res
 
     def get(self, attr_names: str | list[str]) -> list[Any]:
         """


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index f4e64ce5..1541a46b 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -83,12 +83,6 @@ def test_agentset():
     assert all(
         a1 == a2.unique_id for a1, a2 in zip(agentset.get("unique_id"), agentset)
     )
-    assert all(
-        a1 == a2.unique_id
-        for a1, a2 in zip(
-            agentset.do("get_unique_identifier", return_results=True), agentset
-        )
-    )
     assert agentset == agentset.do("get_unique_identifier")
 
     agentset.discard(agents[0])
@@ -276,6 +270,35 @@ def test_agentset_do_callable():
     assert len(agentset) == 0
 
 
+def test_agentset_map_str():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    with pytest.raises(AttributeError):
+        agentset.do("non_existing_method")
+
+    results = agentset.map("get_unique_identifier")
+    assert all(i == entry for i, entry in zip(results, range(1, 11)))
+
+
+def test_agentset_map_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.map(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
+
+    results = agentset.map(lambda agent: agent.unique_id)
+    assert all(i == entry for i, entry in zip(results, range(1, 11)))
+
+
 def test_agentset_get_attribute():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]