Repository: projectmesa/mesa
Pull request #2220 (2024-08-17 15:23:43)

Add AgentSet.groupby
This PR adds pandas style groupby to AgentSet (a first draft for this was #2092, but this PR has stalled). It also adds, as in pandas, a GroupBy helper class to enable method chaining. The resulting API is illustrated below.


**Key Changes**

* adds a `group_by` method to AgentSet; this method enables pandas style group_by operations on the agents in the set. 
* adds a `GroupBy` helper class; this helper class is returned by the new `group_by` method and, as in pandas, has a few methods to enable method chaining in a [split, apply, combine](https://pandas.pydata.org/docs/user_guide/groupby.html) style syntax. 

**usage examples**

```python

# eppstein civil violence
# will give you a dict with the number of agents for each condition
# the result_type optional kwarg allows finer control over the group data structure
# which enhances performance if you don't need the fancy stuff in agentset
# map returns a dict with group_name and the return of whatever function is applied to each group
some_agentset.groupby("condition", result_type="list").map(len)

# random activation by type
groups = some_agentset.groupby(type)
for group, agents in groups: # we can iterate over the groups
    agents.shuffle(inplace=True).do("step")

# we can easily get the number of groups
len(groups)

#  we can get a dict with the group_name and the groups
groups.groups

# one liner random activation by type
# do returns the underlying agent set
# note how do and appply can take either a callable or a string analogous to AgentSet.do
some_agentset.groupby(type).do("shuffle", inplace=True).do("step")

```


github-actions[bot]: Performance benchmarks:

| Model        | Size   | Init time [95% CI]      | Run time [95% CI]       |
|--------------|--------|-------------------------|-------------------------|
| Schelling    | small  | 🔴 +6.1% [+5.8%, +6.5%] | 🔵 -0.3% [-0.5%, -0.1%] |
| Schelling    | large  | 🔴 +5.8% [+5.1%, +6.5%] | 🔵 +0.3% [-0.3%, +0.9%] |
| WolfSheep    | small  | 🔵 +3.5% [+2.3%, +4.7%] | 🔵 +2.4% [+2.1%, +2.8%] |
| WolfSheep    | large  | 🔴 +3.7% [+3.3%, +4.0%] | 🔵 +1.2% [+0.8%, +1.6%] |
| BoidFlockers | small  | 🔵 +2.9% [+2.2%, +3.6%] | 🔵 +1.9% [+1.3%, +2.7%] |
| BoidFlockers | large  | 🔵 +3.2% [+2.8%, +3.7%] | 🔵 +2.0% [+1.3%, +2.7%] |

EwoutH: > specifically enable user choice between agentsets and lists. Lists are faster but should be used with care

It seems that this is a recurring issue.

I have no idea if something like the following is possible in Python, but ideally you would have a fast construct that some function can accept and some functions can output aside from the agentset. As soon as you match or chain two of such functions, they recognize they can take the "fast path" and avoid compiling an recompiling an agentset.

But that's on a high level, no idea if it's possible/feasible to implement in our case in Python.
EwoutH: Apparently this is a concept called "Deferred Execution". Basically, you can do something like this:
- Instead if directly performing each operation, you chain all operations that need to be performed in a list
- Once you have gathered all operations, you check for each out if they support the fast data construct as both input and output
- If there are any matches, flip a boolean for those operations
- They you execute them as the original chain

You need some wrapper and/or decorator magic to make this work, and it adds a lot of code complexity. But then it could prevent unnecessary conversions.

It's kind of related to the lazy execution concept, but now you're able to look forward instead of only backwards.
quaquel: > It seems that this is a recurring issue.

Yes it is

> Apparently this is a concept called "Deferred Execution". Basically, you can do something like this:

I had a look but this seems a lot of trouble to get working. You need to use reflection to figure out all calls and then complex logic to figure out what do to. The short term and easier solution is to allow users to control this explicitly. With the `inplace` argument on for example `shuffle` we are already doing so. It is just a matter of documenting it carefully and making clear that it can be beneficial to profile this carefully. 
quaquel: This is ready for review. All tests are passing and the docstring is in place. The list/agentset stuff has been raised by Ewout and I believe giving users control is the least worst solution at the moment. 

What remains is the split apply combine stuff. Split/groupby is there. Apply is available on the `GroupBy` helper class. I am less sure on the need and desirability of an explicit control for combine. In pandas this is needed because everything needs to be transformed back to a DataFrame/Series. We don't have that need in MESA.
EwoutH: Thanks a lot for this. Will take a close look tomorrow. If you have the chance, could you maybe add a few more user API examples in the PR start?
EwoutH: On initial look this is a clean and solid start, especially API wise. I don't expect much work to be needed, we can extend what we find missing in later PRs.

Tomorrow I will play with some models to get a real feel for this function.
quaquel: > On initial look this is a clean and solid start, especially API wise. I don't expect much work to be needed, we can extend what we find missing in later PRs.

I came up with a couple of simple additions to `GroupBy.apply`. First, it now only takes a callable, this should be changed to a callable or string so you can call methods on the underlying agentset if desired. Next, it might be nice to control the return type to be either the return of the callable/str or the GroupBy instance itself (as we do in AgentSet in places). This allows for chaining like

```python

# random activation by type
some_agentset.groupby(type).apply("shuffle", returntype="groupby", inplace=true).apply("do")

```
EwoutH: Thanks. I think the best to keep group_by properly scoped is to have a few very clear use cases. Data collection (with aggregation) is clearly one, RandomActivationByType could be another one.

What I find weird about that scheduler by the way, is that is shuffles the order of the agent groups and the agents within the group _both_ by _default_:

https://github.com/projectmesa/mesa/blob/3cf1b76ca3a157c8a064b60130755a86e73718be/mesa/time.py#L349

Which means the same group can go twice in a row, etc. Not a weird option to have, but strange by default.

Luckily the AgentSet is way more flexible. A `shuffle_groups()` method might be useful in `GroupBy` if you want to replicate this behavior.

But all based on clear use cases.
Corvince: I think this PR adds a nice feature, but I would advice to carefully reconsider the return types. Currently the flow is like the following, if we consider all possible return types (including the ones proposed for apply)

```python
agentset.group_by("foo").apply(bar)
AgentSet -> GroupBy[str, AgentSet | list] -> Dict[str, Any] | GroupBy[str, AgentSet | list]
```

It took me quite some time to actually figure this out, because you really have to think through what can be returned by each function. Compare this to Pandas
```python
dataframe.groupby("foo").apply(bar)
DataFrame -> DataFrameGroupBy[str, DataFrame] -> DataFrame
```

I would recommend following a similar path and staying as close to our AgentSet as possible. I think defining return types is also a bit unpythonic. Instead of `agentset.groupby(foo, response_type="list")` it is more pythonic to just use `list(agentset.groupby(foo)`. 

Also please note the different spelling of `group_by` (this PR) and `groupby` (pandas)
quaquel: > I think this PR adds a nice feature, but I would advice to carefully reconsider the return types. Currently the flow is like the following, if we consider all possible return types (including the ones proposed for apply)

I have been somewhat concerned about this as well, but you drive home the importance of sorting it out. Its even more confusing then I thought. 

> I think defining return types is also a bit unpythonic. Instead of agentset.groupby(foo, response_type="list") it is more pythonic to just use list(agentset.groupby(foo).

Here I disagree somewhat. Pandas has control for returntype in various places E.g. DataFrame.apply has a way of specifying the result_type. We also do some of it in AgentSet already (i.e., AgentSet.do). The main problem is that in pandas, groupby will allways go back to a dataframe. I don't think there is such a obvious default in our case. 

One possible solution is to only add `group_by` (renamed as @Corvince  points out to groupby) and don't add the GroupBy class with the apply function. 
rht: > One possible solution is to only add group_by (renamed as @Corvince points out to groupby) and don't add the GroupBy class with the apply function.

I thought this was because Polars had decided to use `group_by` instead of `groupby`. In dplyr, it is [group_by](https://dplyr.tidyverse.org/reference/group_by.html). I'm not sure of Polars' reasoning, but it seems to be for consistency with dplyr and that Polars is a cross-language library. I'm fine with either way.
rht: Other than the typo, everything else LGTM.
EwoutH: @quaquel ready for final review, or were you still working on stuff?
quaquel: > @quaquel ready for final review, or were you still working on stuff?

There is still [this idea](https://github.com/projectmesa/mesa/pull/2220#issuecomment-2298042278). The first one (i.e. callable or string on `GroupBy.apply`), I most definately will this morning. Feedback on second point is appreciated and interacts with the point of @Corvince.


EwoutH: Thanks

> changed to a callable or string so you can call methods on the underlying agentset

That might be useful indeed.

> return of the callable/str

How would this work exactly? Can you give some usage examples?
quaquel: > How would this work exactly? Can you give some usage examples?

SImilar to `AgentSet.do` and `return_results=False`. 

```python

# one liner random activation by type
some_agentset.groupby(type).apply("shuffle", inplace=True, return_results=False).apply("step", return_results=False)

```
Corvince: I like how this evolves and I think this will be a nice building block.

However, I think providing do-like functionality under the apply name will be confusing. 

I don't know if we had this discussion already for AgentSet, but what if we split into a `apply` function and a `do` function? Where apply applies a function and returns the result. And do applies a function, but returns the AgentSet. I think this would be clearer and easier to understand in terms of data flow. 
Maybe this could also apply to the distinction of inplace? Where "do" would apply inplace and "apply" would return a new AgentSet (for functions that return an AgentSet´). Right now this would make a lot of sense in my mind, but I haven't fully thought this through yet, so I might overlook something.  
quaquel: >I don't know if we had this discussion already for AgentSet, but what if we split into a apply function and a do function?  

I have to think about this a bit. There is tradeoff between using different methods and just having a keyword to control dataflow.  Having several methods expands the API that a user has to remember. But indeed, it might be clearer because you don't have to remember the various kwargs for return types. @EwoutH @rht what are your thoughts?


EwoutH: Could you add three methods to GroupBy allow more dict-like behavior?

```Python
    def __getitem__(self, key):
        """Allow direct access to groups using the dictionary syntax"""
        return self.groups[key]

    def __len__(self):
        """Allow len() to return the number of groups"""
        return len(self.groups)

    def __getattr__(self, item):
        """
        Delegate attribute access to the internal `groups` dict, 
        allowing for all dictionary methods to be used directly on the GroupBy object.
        """
        return getattr(self.groups, item)
```
`__getattr__` is especially magic, it allows you to do things like `GroupBy.keys()`, `.update()`, `.values()` and any other [dict](https://docs.python.org/3/library/stdtypes.html#dict) method. 

This will allow users to also be able to just threat a groupby object like a dict, which is a know construct to many modellers and doesn't require to learn our API.
EwoutH: > Where apply applies a function and returns the result. And do applies a function, but returns the AgentSet.

In general, I like methods to have different names if they have different output types or structures.

Let me play with it some more.

(also maybe that's just me, but I really need usage examples to visualise this properly)
quaquel: > Could you add three methods to GroupBy allow more dict-like behavior?
> 
> ```python
>     def __getitem__(self, key):
>         """Allow direct access to groups using the dictionary syntax"""
>         return self.groups[key]
> 
>     def __len__(self):
>         """Allow len() to return the number of groups"""
>         return len(self.groups)
> 
>     def __getattr__(self, item):
>         """
>         Delegate attribute access to the internal `groups` dict, 
>         allowing for all dictionary methods to be used directly on the GroupBy object.
>         """
>         return getattr(self.groups, item)
> ```
> 
> `__getattr__` is especially magic, it allows you to do things like `GroupBy.keys()`, `.update()`, `.values()` and any other [dict](https://docs.python.org/3/library/stdtypes.html#dict) method.
> 
> This will allow users to also be able to just threat a groupby object like a dict, which is a know construct to many modellers and doesn't require to learn our API.
> 
> **Edit:** Maybe we also want `list` and maybe `del`. It would be nice if we someway just also could allow any function that calls a dict.

I deliberately did not do this. I tried to mimic pandas.AgentSetGroupBy rather than have our GroupBy be effectivley a dict. For example, if we want GroupBy to be a proper dict we would have to write

```python

for group_name, group in some_gouby_instance.items():
    # do something here

```

At the moment, because GroupBy is not a full dict, you can instead do

```python

for group_name, group in some_gouby_instance:
    # do something here

```

Which is exactly how it works in pandas. 

EwoutH: Interesting, thanks for sharing.
> ```python
> for group_name, group in some_gouby_instance:
>     # do something here
> ```

I think as long as we keep `__iter__` we can still do that right, since it overwrites the default dict behavior?

Direct key access with `__getitem__` seems very useful anyways.
quaquel: > Direct key access with __getitem__ seems very useful anyways.

In line with pandas, this is possible via `get_group`. Checking [DataFrameGroupBy](https://pandas.pydata.org/docs/reference/groupby.html), I noticed that it has a `GroupBy.groups` which returns the underlying dict. would that solve your problem?
EwoutH: Let me explain my train of thought:
- We're adding more and more stuff to our API, it's getting quite convoluted
- Users have to learn more and more specific details about our inputs, arguments and returns
- How can we simplify that?
- What if it behaved more like standard structures that people know?
- GroupBy is basically just a dict with some polish around it, right?
- What if we could let it behave as much as a dict, a beloved built-in construct basically everyone knows
- Could it be just like a superset of a dict? With a very similar API, with some additional functionality?
- That way, people could start using it like a dict and use more advanced functionality when they get to the limitations of our current approach (in line with scaffolded learning principle @tpike3 mentioned)

So from that perspective, my preferred return type would actually be a dict. But you can't call apply on a dict, so you just need a little magic around it to make a few mesa-specific function work. But personally, I was thinking of staying as close to a dict as possible.
quaquel: I follow that train of thought, but rather than let it behave like a dict, I want it to behave like pandas DataFrameGroupBy. 
Corvince: 
> (also maybe that's just me, but I really need usage examples to visualise this properly)

Adopting the example from @quaquel 

> ```python
> # one liner random activation by type
> some_agentset.groupby(type).apply("shuffle", inplace=True, return_results=False).apply("step", return_results=False)
> ```

```python
some_agentset.groupby(type).do("shuffle").do("step").apply(len)
```

This would randomly activate agents and then collect the length of each agent group, e.g. to see how many are left in each group. It would be equal to 

```python
some_agentset.do("shuffle").do("step").groupby(type).apply(len)
```

Another example
```python
some_agentset.do("select", lambda a: a.wealth > 5) # <- modifies some_agentset inline
other_agentset = some_agentset.apply("select", lambda a: a.wealth > 5) # <- returns new agentset
```
quaquel: > ```python some_agentset.do("select", lambda a: a.wealth > 5) # <- modifies some_agentset inline
>  other_agentset = some_agentset.apply("select", lambda a: a.wealth > 5) # <- returns new agentset```

This won't work. Select is a method on the AgentSet itself, while `do` applies the method/function to each agent in the set.
Corvince: I just realized that, but you beat me to editing it! 
EwoutH: > but rather than let it behave like a dict, I want it to behave like pandas DataFrameGroupBy.

I'm going to do a little bit of reading up on DataFrameGroupBy.

(but the fact that I need to do reading up on DataFrameGroupBy worries me for our general users)

What are some things you can do when it behaves like a DataFrameGroupBy object that are possible or more elegant compared that it behaves mostly like a dict with an apply function?
EwoutH: Thanks for the examples @Corvince! The aggregation use-case I'm fully behind.

Here's maybe an interesting thought exercise: What would the least complex / unfamiliar implementation (especially in user API) be to properly facilitate aggregation? Maybe we're making this way too complex.
EwoutH: For now I have found the following use cases:
- Aggregation into different groups is useful. Group, aggerate, output in some nice format (`{group: agg_value}` dict?)
- Transformation
- Doing something to a specific group: That would use `select().do()`
- Doing something to all groups: Yes, but then you would need to pass the `value` of the groupby also, instead of only the group. Otherwise you can just apply it on all agents at once.
- Doing something to all groups in a specific order of groups (random activation by type). That would indeed need a full `agents.groupby().sort_groups().do()`
  - If you have a predetermined order, you could do
    ```Python
    groups = agents.groupby()
    for group in predetermined_order:
        groups[group].do("step")
    ```
- You might want to take it further, if you want to dynamically want to order your groups on either:
  - The group key.
  - A aggregated value

<hr>

But taking a step back, and looking at the code, I think it's quite clean and a good first step. There are cases to be made for both an AgentSet or a List to be returned as each group value, but I think AgentSet is the best default, because it offers flexibility. When you need more performance you can explore list.

I would really like some syntactic sugar to allow for clean, intuitive dict-like behavior. Personally, I think `__len__` and `__getitem__` are obvious: You want to know how many groups you have or you want to select a specific group. `__getattr__` also sounds very useful, to quickly get the keys or items. How would you do that else?
quaquel: What about the following next steps?
1. I'll seperate `apply` and `do` along the line suggested by @Corvince but only in the GroupBy helper class. I do think it is a good idea also for AgentSet, but it is a breaking change (can be partially handled with a temporary warning) and might better be done in a seperate PR.
2. I still don't agree with @EwoutH on the dict behavior. First, the iteration as currently implemented is not dict like because it returns group_name and group rather than only group_name as would happen in a normal dict. Second, GroupBy in pandas is really a helper class in chaining calls (groupby, apply, combine). I'll happily add the `groups` attribute that returns the underlying dict. Third, adding dict behavior to this class will come with performance penalties and are really just passed on to the underlying dict anyway. Why not make it easy to return this dict and work on that directly?
EwoutH: What about an optional return type that returns a dict instead of a group?
quaquel: > What about an optional return type that returns a dict instead of a group?
> 
> Edit or just a `group_to_dict()` function

We just agree to seperate do and apply to get rid of a kwarg for return type. My solution is just a groups attribute (rather than a function). So

```python

some_dict: dict[Any, Agentset] = some_agentset.groupby(type).groups

```
EwoutH: > I would really like some syntactic sugar to allow for clean, intuitive dict-like behavior. Personally, I think `__len__` and `__getitem__` are obvious: You want to know how many groups you have or you want to select a specific group. `get_group()` is something new you have to learn. Just accessing an item like in a dict is obvious, especially since it already iters like a dict.
> 
> `__getattr__` also sounds very useful, to quickly get the keys or items. How would you do that else?

I'm okay with leaving `__getattr__` out of this PR. `.values()` would be weird anyway, you just get back all the agents in some weird lists. `.items()` is covered by `__iter__` and if you want to know the group name, just do `GroupBy.groups.keys()`.

How to retrieve a group we disagree on. You want to do `get_group()`, I think `__getitem__` is more elegant. Let's leave both out and discuss this separately.

I still think `__len__` would be very useful. I just grouped, but how many groups do I now have? Can be useful for statistics, for determining things like bins for plotting, etc.
quaquel: > I still think __len__ would be very useful. I just grouped, but how many groups do I now have? Can be useful for statistics, for determining things like bins for plotting, etc.

Fair enough
quaquel: I have added `__len__`, seperated `apply` and `do`, and removed `get_groups`. Unittests have also been updated. This is ready for a last review. 

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..f137ba09 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -11,7 +11,9 @@ 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
 
@@ -216,25 +218,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 +290,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]:
         """
@@ -357,7 +398,116 @@ class AgentSet(MutableSet, Sequence):
         """
         return self.model.random
 
+    def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
+        """
+        Group agents by the specified attribute or return from the callable
+
+        Args:
+            by (Callable, str): used to determine what to group agents by
+
+                                * if ``by`` is a callable, it will be called for each agent and the return is used
+                                  for grouping
+                                * if ``by`` is a str, it should refer to an attribute on the agent and the value
+                                  of this attribute will be used for grouping
+            result_type (str, optional): The datatype for the resulting groups {"agentset", "list"}
+        Returns:
+            GroupBy
+
+
+        Notes:
+        There might be performance benefits to using `result_type='list'` if you don't need the advanced functionality
+        of an AgentSet.
+
+        """
+        groups = defaultdict(list)
+
+        if isinstance(by, Callable):
+            for agent in self:
+                groups[by(agent)].append(agent)
+        else:
+            for agent in self:
+                groups[getattr(agent, by)].append(agent)
+
+        if result_type == "agentset":
+            return GroupBy(
+                {k: AgentSet(v, model=self.model) for k, v in groups.items()}
+            )
+        else:
+            return GroupBy(groups)
+
+    # consider adding for performance reasons
+    # for Sequence: __reversed__, index, and count
+    # for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
+
+
+class GroupBy:
+    """Helper class for AgentSet.groupby
+
+
+    Attributes:
+        groups (dict): A dictionary with the group_name as key and group as values
+
+    """
+
+    def __init__(self, groups: dict[Any, list | AgentSet]):
+        self.groups: dict[Any, list | AgentSet] = groups
+
+    def map(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
+        """Apply the specified callable to each group and return the results.
+
+        Args:
+            method (Callable, str): The callable to apply to each group,
+
+                                    * if ``method`` is a callable, it will be called it will be called with the group as first argument
+                                    * if ``method`` is a str, it should refer to a method on the group
+
+                                    Additional arguments and keyword arguments will be passed on to the callable.
+
+        Returns:
+            dict with group_name as key and the return of the method as value
+
+        Notes:
+            this method is useful for methods or functions that do return something. It
+            will break method chaining. For that, use ``do`` instead.
+
+        """
+        if isinstance(method, str):
+            return {
+                k: getattr(v, method)(*args, **kwargs) for k, v in self.groups.items()
+            }
+        else:
+            return {k: method(v, *args, **kwargs) for k, v in self.groups.items()}
+
+    def do(self, method: Callable | str, *args, **kwargs) -> GroupBy:
+        """Apply the specified callable to each group
+
+        Args:
+            method (Callable, str): The callable to apply to each group,
+
+                                    * if ``method`` is a callable, it will be called it will be called with the group as first argument
+                                    * if ``method`` is a str, it should refer to a method on the group
+
+                                    Additional arguments and keyword arguments will be passed on to the callable.
+
+        Returns:
+            the original GroupBy instance
+
+        Notes:
+            this method is useful for methods or functions that don't return anything and/or
+            if you want to chain multiple do calls
+
+        """
+        if isinstance(method, str):
+            for v in self.groups.values():
+                getattr(v, method)(*args, **kwargs)
+        else:
+            for v in self.groups.values():
+                method(v, *args, **kwargs)
+
+        return self
+
+    def __iter__(self):
+        return iter(self.groups.items())
 
-# consider adding for performance reasons
-# for Sequence: __reversed__, index, and count
-# for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
+    def __len__(self):
+        return len(self.groups)


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index f4e64ce5..516902f9 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)]
@@ -348,3 +371,45 @@ def test_agentset_shuffle():
     agentset = AgentSet(test_agents, model=model)
     agentset.shuffle(inplace=True)
     assert not all(a1 == a2 for a1, a2 in zip(test_agents, agentset))
+
+
+def test_agentset_groupby():
+    class TestAgent(Agent):
+        def __init__(self, unique_id, model):
+            super().__init__(unique_id, model)
+            self.even = self.unique_id % 2 == 0
+
+        def get_unique_identifier(self):
+            return self.unique_id
+
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    groups = agentset.groupby("even")
+    assert len(groups.groups[True]) == 5
+    assert len(groups.groups[False]) == 5
+
+    groups = agentset.groupby(lambda a: a.unique_id % 2 == 0)
+    assert len(groups.groups[True]) == 5
+    assert len(groups.groups[False]) == 5
+    assert len(groups) == 2
+
+    for group_name, group in groups:
+        assert len(group) == 5
+        assert group_name in {True, False}
+
+    sizes = agentset.groupby("even", result_type="list").map(len)
+    assert sizes == {True: 5, False: 5}
+
+    attributes = agentset.groupby("even", result_type="agentset").map("get", "even")
+    for group_name, group in attributes.items():
+        assert all(group_name == entry for entry in group)
+
+    groups = agentset.groupby("even", result_type="agentset")
+    another_ref_to_groups = groups.do("do", "step")
+    assert groups == another_ref_to_groups
+
+    groups = agentset.groupby("even", result_type="agentset")
+    another_ref_to_groups = groups.do(lambda x: x.do("step"))
+    assert groups == another_ref_to_groups