From d90b0f72cd65d5875eb9e4012990699cfe886c5c Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Sun, 21 Jan 2024 18:57:57 +0100
Subject: [PATCH 01/25] further updates

---
 benchmarks/WolfSheep/__init__.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/benchmarks/WolfSheep/__init__.py b/benchmarks/WolfSheep/__init__.py
index e69de29bb2d..18b86ab19ba 100644
--- a/benchmarks/WolfSheep/__init__.py
+++ b/benchmarks/WolfSheep/__init__.py
@@ -0,0 +1,14 @@
+from wolf_sheep import WolfSheep
+
+if __name__ == "__main__":
+    # for profiling this benchmark model
+    import time
+
+    # model = WolfSheep(15, 25, 25, 60, 40, 0.2, 0.1, 20)
+    model = WolfSheep(15, 100, 100, 1000, 500, 0.4, 0.2, 20)
+
+    start_time = time.perf_counter()
+    for _ in range(100):
+        model.step()
+
+    print(time.perf_counter() - start_time)

From 95864900d9e7981c8ea7240e34672f3fe792d681 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Sun, 21 Jan 2024 19:13:00 +0100
Subject: [PATCH 02/25] Update benchmarks/WolfSheep/__init__.py

---
 benchmarks/WolfSheep/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmarks/WolfSheep/__init__.py b/benchmarks/WolfSheep/__init__.py
index 18b86ab19ba..98b1e9fdfed 100644
--- a/benchmarks/WolfSheep/__init__.py
+++ b/benchmarks/WolfSheep/__init__.py
@@ -1,4 +1,4 @@
-from wolf_sheep import WolfSheep
+from .wolf_sheep import WolfSheep
 
 if __name__ == "__main__":
     # for profiling this benchmark model

From 96cd80abe4c1ec92fcc6c8accee74623113a1cb8 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Sat, 17 Aug 2024 17:08:37 +0200
Subject: [PATCH 03/25] add groupby

---
 mesa/agent.py | 67 ++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 64 insertions(+), 3 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index e65b25f69c2..f03804ca2b0 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -357,7 +357,68 @@ def random(self) -> Random:
         """
         return self.model.random
 
+    def group_by(
+            self, by: Callable | str,
+            result_type: str = "agentset"
+    ) -> BaseGroupBy:
+        """
+        Group agents by the specified attribute
+
+        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): The datatype for the resutling groups {"agentset", "list"}
+        Returns:
+            AgentSetGroupBy
+
+        """
+        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 AgentSetGroupBy(groups)
+        else:
+            return ListGroupBy(groups)
+
+    # consider adding for performance reasons
+    # for Sequence: __reversed__, index, and count
+    # for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
+
+class BaseGroupBy:
+    def __init__(self, groups: dict[Any, list | AgentSet]):
+        self.groups: dict[Any, list|AgentSet] = groups
+
+    def get_group(self, name: Any):
+        # return group for specified name
+        return self.groups[name]
+
+    def apply(self, callable: Callable):
+        # fixme, we have callable over the entire group and callable on each group member
+        # apply callable to each group and return dict {group_name, return of callable for group}
+        return {k: callable(v) for k, v in self.groups}
+
+    def __iter__(self):
+        return iter(self.groups.items())
+
+class AgentSetGroupBy(BaseGroupBy):
+    # Helper class to enable pandas style split, apply, combine syntax
+
+    def __init__(self, groups: dict[Any, list]):
+        super().__init__({k:AgentSet(v) for k,v in groups.items()})
+    def do(self, method: str | Callable, *args, **kwargs):
+        # fixme what about return type here and enabling method chaining?
+        return {k: v.do(method, *args, **kwargs) for k, v in self.groups}
 
-# consider adding for performance reasons
-# for Sequence: __reversed__, index, and count
-# for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
+class ListGroupBy(BaseGroupBy):
+    # Helper class to enable pandas style split, apply, combine syntax
+    pass

From 07467f3484e8990f60c8defac43f636a9f99e879 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 17 Aug 2024 15:23:48 +0000
Subject: [PATCH 04/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/agent.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index f03804ca2b0..08dbe4ce623 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -358,8 +358,7 @@ def random(self) -> Random:
         return self.model.random
 
     def group_by(
-            self, by: Callable | str,
-            result_type: str = "agentset"
+        self, by: Callable | str, result_type: str = "agentset"
     ) -> BaseGroupBy:
         """
         Group agents by the specified attribute
@@ -386,7 +385,7 @@ def group_by(
                 groups[getattr(agent, by)].append(agent)
 
         if result_type == "agentset":
-           return AgentSetGroupBy(groups)
+            return AgentSetGroupBy(groups)
         else:
             return ListGroupBy(groups)
 
@@ -394,9 +393,10 @@ def group_by(
     # for Sequence: __reversed__, index, and count
     # for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
 
+
 class BaseGroupBy:
     def __init__(self, groups: dict[Any, list | AgentSet]):
-        self.groups: dict[Any, list|AgentSet] = groups
+        self.groups: dict[Any, list | AgentSet] = groups
 
     def get_group(self, name: Any):
         # return group for specified name
@@ -410,15 +410,18 @@ def apply(self, callable: Callable):
     def __iter__(self):
         return iter(self.groups.items())
 
+
 class AgentSetGroupBy(BaseGroupBy):
     # Helper class to enable pandas style split, apply, combine syntax
 
     def __init__(self, groups: dict[Any, list]):
-        super().__init__({k:AgentSet(v) for k,v in groups.items()})
+        super().__init__({k: AgentSet(v) for k, v in groups.items()})
+
     def do(self, method: str | Callable, *args, **kwargs):
         # fixme what about return type here and enabling method chaining?
         return {k: v.do(method, *args, **kwargs) for k, v in self.groups}
 
+
 class ListGroupBy(BaseGroupBy):
     # Helper class to enable pandas style split, apply, combine syntax
     pass

From b04099fc3074af119ede728fee16fe56e7cde82b Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Sun, 18 Aug 2024 16:37:36 +0200
Subject: [PATCH 05/25] Update agent.py

---
 mesa/agent.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mesa/agent.py b/mesa/agent.py
index 08dbe4ce623..96f1f7603c1 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -13,6 +13,7 @@
 import operator
 import weakref
 from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
+from collections import defaultdict
 from random import Random
 
 # mypy

From 6faaf9e0db7ba3102bafc55b34aa5ee1989e976d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 19 Aug 2024 09:31:50 +0000
Subject: [PATCH 06/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/agent.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 96f1f7603c1..371aa7ac88c 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -12,8 +12,8 @@
 import copy
 import operator
 import weakref
-from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
 from collections import defaultdict
+from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
 # mypy

From 7eed1db43b2e164cbe98c6be1491758a28beeb3e Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Mon, 19 Aug 2024 11:41:05 +0200
Subject: [PATCH 07/25] add tests and simplify groupby objects

---
 mesa/agent.py       | 22 ++++------------------
 tests/test_agent.py | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 371aa7ac88c..67e1c0ef593 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -386,16 +386,16 @@ def group_by(
                 groups[getattr(agent, by)].append(agent)
 
         if result_type == "agentset":
-            return AgentSetGroupBy(groups)
+            return GroupBy({k:AgentSet(v, model=self.model) for k,v in groups.items()})
         else:
-            return ListGroupBy(groups)
+            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 BaseGroupBy:
+class GroupBy:
     def __init__(self, groups: dict[Any, list | AgentSet]):
         self.groups: dict[Any, list | AgentSet] = groups
 
@@ -406,23 +406,9 @@ def get_group(self, name: Any):
     def apply(self, callable: Callable):
         # fixme, we have callable over the entire group and callable on each group member
         # apply callable to each group and return dict {group_name, return of callable for group}
-        return {k: callable(v) for k, v in self.groups}
+        return {k: callable(v) for k, v in self.groups.items()}
 
     def __iter__(self):
         return iter(self.groups.items())
 
 
-class AgentSetGroupBy(BaseGroupBy):
-    # Helper class to enable pandas style split, apply, combine syntax
-
-    def __init__(self, groups: dict[Any, list]):
-        super().__init__({k: AgentSet(v) for k, v in groups.items()})
-
-    def do(self, method: str | Callable, *args, **kwargs):
-        # fixme what about return type here and enabling method chaining?
-        return {k: v.do(method, *args, **kwargs) for k, v in self.groups}
-
-
-class ListGroupBy(BaseGroupBy):
-    # Helper class to enable pandas style split, apply, combine syntax
-    pass
diff --git a/tests/test_agent.py b/tests/test_agent.py
index f4e64ce5338..e94b8a63632 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -348,3 +348,26 @@ 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.group_by("even")
+    assert len(groups.get_group(True)) == 5
+    assert len(groups.get_group(False)) == 5
+
+    sizes = agentset.group_by("even", result_type="list").apply(len)
+    assert sizes == {True: 5, False: 5}
+

From e9be4ccacf3cfc47241641aafec3ad7cabe78ad5 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Mon, 19 Aug 2024 11:41:55 +0200
Subject: [PATCH 08/25] Update test_agent.py

---
 tests/test_agent.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/test_agent.py b/tests/test_agent.py
index e94b8a63632..1b1586c3053 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -368,6 +368,11 @@ def get_unique_identifier(self):
     assert len(groups.get_group(True)) == 5
     assert len(groups.get_group(False)) == 5
 
+    for group_name, group in groups:
+        assert len(group) == 5
+
     sizes = agentset.group_by("even", result_type="list").apply(len)
     assert sizes == {True: 5, False: 5}
 
+
+

From 6a329bd304050a15688405523b05752e6a96b367 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 19 Aug 2024 09:42:20 +0000
Subject: [PATCH 09/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/agent.py       | 6 +++---
 tests/test_agent.py | 8 ++------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 67e1c0ef593..0f00d602a6f 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -386,7 +386,9 @@ def group_by(
                 groups[getattr(agent, by)].append(agent)
 
         if result_type == "agentset":
-            return GroupBy({k:AgentSet(v, model=self.model) for k,v in groups.items()})
+            return GroupBy(
+                {k: AgentSet(v, model=self.model) for k, v in groups.items()}
+            )
         else:
             return GroupBy(groups)
 
@@ -410,5 +412,3 @@ def apply(self, callable: Callable):
 
     def __iter__(self):
         return iter(self.groups.items())
-
-
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 1b1586c3053..b2cca908f2f 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -352,14 +352,13 @@ def test_agentset_shuffle():
 
 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
+            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)
@@ -373,6 +372,3 @@ def get_unique_identifier(self):
 
     sizes = agentset.group_by("even", result_type="list").apply(len)
     assert sizes == {True: 5, False: 5}
-
-
-

From 16d75c5ad9c1e5e9c36e75e1019887d073b134cb Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Mon, 19 Aug 2024 15:46:57 +0200
Subject: [PATCH 10/25] ruff fixes

---
 mesa/agent.py       | 4 ++--
 tests/test_agent.py | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 0f00d602a6f..1a7c24d42af 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -360,7 +360,7 @@ def random(self) -> Random:
 
     def group_by(
         self, by: Callable | str, result_type: str = "agentset"
-    ) -> BaseGroupBy:
+    ) -> GroupBy:
         """
         Group agents by the specified attribute
 
@@ -373,7 +373,7 @@ def group_by(
                                   of this attribute will be used for grouping
             result_type (str): The datatype for the resutling groups {"agentset", "list"}
         Returns:
-            AgentSetGroupBy
+            GroupBy
 
         """
         groups = defaultdict(list)
diff --git a/tests/test_agent.py b/tests/test_agent.py
index b2cca908f2f..496e0104f7e 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -369,6 +369,7 @@ def get_unique_identifier(self):
 
     for group_name, group in groups:
         assert len(group) == 5
+        assert group_name in {True, False}
 
     sizes = agentset.group_by("even", result_type="list").apply(len)
     assert sizes == {True: 5, False: 5}

From e09f5a2ec2ebc343dfcb17ba224fefec320280be Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 19 Aug 2024 13:47:05 +0000
Subject: [PATCH 11/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/agent.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 1a7c24d42af..acc40cc9b1e 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -358,9 +358,7 @@ def random(self) -> Random:
         """
         return self.model.random
 
-    def group_by(
-        self, by: Callable | str, result_type: str = "agentset"
-    ) -> GroupBy:
+    def group_by(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
         """
         Group agents by the specified attribute
 

From a4d95592dba31eed6dca81176a95350e9208cfdc Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Mon, 19 Aug 2024 15:51:02 +0200
Subject: [PATCH 12/25] additional docs

---
 mesa/agent.py | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index acc40cc9b1e..bf6dc9c0fd0 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -369,7 +369,7 @@ def group_by(self, by: Callable | str, result_type: str = "agentset") -> GroupBy
                                   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): The datatype for the resutling groups {"agentset", "list"}
+            result_type (str, optional): The datatype for the resulting groups {"agentset", "list"}
         Returns:
             GroupBy
 
@@ -396,17 +396,31 @@ def group_by(self, by: Callable | str, result_type: str = "agentset") -> GroupBy
 
 
 class GroupBy:
+    """Helper class for AgentSet.groupby"""
+
     def __init__(self, groups: dict[Any, list | AgentSet]):
         self.groups: dict[Any, list | AgentSet] = groups
 
     def get_group(self, name: Any):
+        """
+        Retrieve the specified group by name.
+        """
+
         # return group for specified name
         return self.groups[name]
 
-    def apply(self, callable: Callable):
-        # fixme, we have callable over the entire group and callable on each group member
-        # apply callable to each group and return dict {group_name, return of callable for group}
-        return {k: callable(v) for k, v in self.groups.items()}
+    def apply(self, callable: Callable, *args, **kwargs):
+        """Apply the specified callable to each group
+
+        Args:
+            callable (Callable): The callable to apply to each group, it will be called with the group as first argument
+                                 Additional arguments and keyword arguments will be passed on to the callable.
+
+        Returns:
+            dict with group_name as key and the return of the callable as value
+
+        """
+        return {k: callable(v, *args, **kwargs) for k, v in self.groups.items()}
 
     def __iter__(self):
         return iter(self.groups.items())

From 7db02b86d2e2727dc9f93a741659828d499e8a78 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Mon, 19 Aug 2024 20:56:51 +0200
Subject: [PATCH 13/25] add test for callable

---
 tests/test_agent.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/test_agent.py b/tests/test_agent.py
index 496e0104f7e..ec370b93e6c 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -367,6 +367,12 @@ def get_unique_identifier(self):
     assert len(groups.get_group(True)) == 5
     assert len(groups.get_group(False)) == 5
 
+
+    groups = agentset.group_by(lambda a: a.unique_id % 2 == 0)
+    assert len(groups.get_group(True)) == 5
+    assert len(groups.get_group(False)) == 5
+
+
     for group_name, group in groups:
         assert len(group) == 5
         assert group_name in {True, False}

From 8aa56dacace0c525daa4366698834aedd1352ac6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 19 Aug 2024 18:58:01 +0000
Subject: [PATCH 14/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_agent.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/test_agent.py b/tests/test_agent.py
index ec370b93e6c..585901b2232 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -367,12 +367,10 @@ def get_unique_identifier(self):
     assert len(groups.get_group(True)) == 5
     assert len(groups.get_group(False)) == 5
 
-
     groups = agentset.group_by(lambda a: a.unique_id % 2 == 0)
     assert len(groups.get_group(True)) == 5
     assert len(groups.get_group(False)) == 5
 
-
     for group_name, group in groups:
         assert len(group) == 5
         assert group_name in {True, False}

From 943a907d14a484b94aed28c4d3210fa5292fa080 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Tue, 20 Aug 2024 17:09:06 +0200
Subject: [PATCH 15/25] rename groupby and additional docs

---
 mesa/agent.py       | 10 ++++++++--
 tests/test_agent.py |  6 +++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index bf6dc9c0fd0..dcc384275fc 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -358,9 +358,9 @@ def random(self) -> Random:
         """
         return self.model.random
 
-    def group_by(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
+    def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
         """
-        Group agents by the specified attribute
+        Group agents by the specified attribute or return from the callable
 
         Args:
             by (Callable, str): used to determine what to group agents by
@@ -373,6 +373,12 @@ def group_by(self, by: Callable | str, result_type: str = "agentset") -> GroupBy
         Returns:
             GroupBy
 
+
+        Nptes
+        -----
+        There might be performance benefits to using `result_type='list'` if you don't need the advanced functionality
+        of an AgentSet.
+
         """
         groups = defaultdict(list)
 
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 585901b2232..a842bb6f511 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -363,11 +363,11 @@ def get_unique_identifier(self):
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]
     agentset = AgentSet(agents, model)
 
-    groups = agentset.group_by("even")
+    groups = agentset.groupby("even")
     assert len(groups.get_group(True)) == 5
     assert len(groups.get_group(False)) == 5
 
-    groups = agentset.group_by(lambda a: a.unique_id % 2 == 0)
+    groups = agentset.groupby(lambda a: a.unique_id % 2 == 0)
     assert len(groups.get_group(True)) == 5
     assert len(groups.get_group(False)) == 5
 
@@ -375,5 +375,5 @@ def get_unique_identifier(self):
         assert len(group) == 5
         assert group_name in {True, False}
 
-    sizes = agentset.group_by("even", result_type="list").apply(len)
+    sizes = agentset.groupby("even", result_type="list").apply(len)
     assert sizes == {True: 5, False: 5}

From 9fc7963b0a60cdb681d231b8725aff35615f7322 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Tue, 20 Aug 2024 17:09:26 +0200
Subject: [PATCH 16/25] Update agent.py

---
 mesa/agent.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index dcc384275fc..04fc550dbcc 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -374,8 +374,7 @@ def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
             GroupBy
 
 
-        Nptes
-        -----
+        Nptes:
         There might be performance benefits to using `result_type='list'` if you don't need the advanced functionality
         of an AgentSet.
 

From c4812815b26fd5a34b853db57341a7965558df5d Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 21 Aug 2024 09:41:50 +0200
Subject: [PATCH 17/25] add str option to GroupBy.apply

---
 mesa/agent.py       | 19 +++++++++++++------
 tests/test_agent.py | 12 ++++++++----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 04fc550dbcc..8ed43c7bfd2 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -374,7 +374,7 @@ def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
             GroupBy
 
 
-        Nptes:
+        Notes:
         There might be performance benefits to using `result_type='list'` if you don't need the advanced functionality
         of an AgentSet.
 
@@ -414,18 +414,25 @@ def get_group(self, name: Any):
         # return group for specified name
         return self.groups[name]
 
-    def apply(self, callable: Callable, *args, **kwargs):
+    def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
         """Apply the specified callable to each group
 
         Args:
-            callable (Callable): The callable to apply to each group, it will be called with the group as first argument
-                                 Additional arguments and keyword arguments will be passed on to the callable.
+            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 callable as value
+            dict with group_name as key and the return of the method as value
 
         """
-        return {k: callable(v, *args, **kwargs) for k, v in self.groups.items()}
+        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 __iter__(self):
         return iter(self.groups.items())
diff --git a/tests/test_agent.py b/tests/test_agent.py
index a842bb6f511..ea97d013425 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -13,9 +13,9 @@ def get_unique_identifier(self):
 
 class TestAgentDo(Agent):
     def __init__(
-        self,
-        unique_id,
-        model,
+            self,
+            unique_id,
+            model,
     ):
         super().__init__(unique_id, model)
         self.agent_set = None
@@ -291,7 +291,7 @@ def test_agentset_get_attribute():
     agents = []
     for i in range(10):
         agent = TestAgent(model.next_id(), model)
-        agent.i = i**2
+        agent.i = i ** 2
         agents.append(agent)
     agentset = AgentSet(agents, model)
 
@@ -377,3 +377,7 @@ def get_unique_identifier(self):
 
     sizes = agentset.groupby("even", result_type="list").apply(len)
     assert sizes == {True: 5, False: 5}
+
+    attributes = agentset.groupby("even", result_type="agentset").apply("get", "even")
+    for group_name, group in attributes.items():
+        assert all([group_name == entry for entry in group])

From 585025b0d9ffa3d88ed3e2f8b85694fff7b4be40 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Aug 2024 07:41:59 +0000
Subject: [PATCH 18/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/agent.py       | 4 +++-
 tests/test_agent.py | 8 ++++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 8ed43c7bfd2..d46b51bb835 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -430,7 +430,9 @@ def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
 
         """
         if isinstance(method, str):
-            return {k: getattr(v, method)(*args, **kwargs) for k, v in self.groups.items()}
+            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()}
 
diff --git a/tests/test_agent.py b/tests/test_agent.py
index ea97d013425..7a289d8490c 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -13,9 +13,9 @@ def get_unique_identifier(self):
 
 class TestAgentDo(Agent):
     def __init__(
-            self,
-            unique_id,
-            model,
+        self,
+        unique_id,
+        model,
     ):
         super().__init__(unique_id, model)
         self.agent_set = None
@@ -291,7 +291,7 @@ def test_agentset_get_attribute():
     agents = []
     for i in range(10):
         agent = TestAgent(model.next_id(), model)
-        agent.i = i ** 2
+        agent.i = i**2
         agents.append(agent)
     agentset = AgentSet(agents, model)
 

From f0dccf6a3adbd320307ed8d58538dd6d28f8e1ef Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 21 Aug 2024 10:38:22 +0200
Subject: [PATCH 19/25] remove unnecessary list comprehension

---
 tests/test_agent.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_agent.py b/tests/test_agent.py
index 7a289d8490c..f2bbd201c7e 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -380,4 +380,4 @@ def get_unique_identifier(self):
 
     attributes = agentset.groupby("even", result_type="agentset").apply("get", "even")
     for group_name, group in attributes.items():
-        assert all([group_name == entry for entry in group])
+        assert all(group_name == entry for entry in group)

From eaf0a5bf4a70a2684aedbf624d46dbb85b9ffcf3 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 21 Aug 2024 15:46:21 +0200
Subject: [PATCH 20/25] seperate apply and do

---
 mesa/agent.py | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index d46b51bb835..06e38946b93 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -401,7 +401,13 @@ def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
 
 
 class GroupBy:
-    """Helper class for AgentSet.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
@@ -436,5 +442,32 @@ def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
         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:
+            GroupBy
+
+        """
+        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())
+
+    def __len__(self):
+        return len(self.groups)

From 7c3a78f002a2b88b74ac30e7da254a6bd9b21865 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 21 Aug 2024 15:48:47 +0200
Subject: [PATCH 21/25] remove unintened update to __init__.py for wolfsheep
 benchmark

---
 benchmarks/WolfSheep/__init__.py | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/benchmarks/WolfSheep/__init__.py b/benchmarks/WolfSheep/__init__.py
index 98b1e9fdfed..e69de29bb2d 100644
--- a/benchmarks/WolfSheep/__init__.py
+++ b/benchmarks/WolfSheep/__init__.py
@@ -1,14 +0,0 @@
-from .wolf_sheep import WolfSheep
-
-if __name__ == "__main__":
-    # for profiling this benchmark model
-    import time
-
-    # model = WolfSheep(15, 25, 25, 60, 40, 0.2, 0.1, 20)
-    model = WolfSheep(15, 100, 100, 1000, 500, 0.4, 0.2, 20)
-
-    start_time = time.perf_counter()
-    for _ in range(100):
-        model.step()
-
-    print(time.perf_counter() - start_time)

From a25dd06d48cafd8c534931428b8a92d28faff9c4 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Wed, 21 Aug 2024 18:55:39 +0200
Subject: [PATCH 22/25] additional tests

---
 mesa/agent.py       | 44 ++++++++++++++++++++++----------------------
 tests/test_agent.py | 18 ++++++++++++++----
 2 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 06e38946b93..79f1bef1cfe 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -118,11 +118,11 @@ def __contains__(self, agent: Agent) -> bool:
         return agent in self._agents
 
     def select(
-        self,
-        filter_func: Callable[[Agent], bool] | None = None,
-        n: int = 0,
-        inplace: bool = False,
-        agent_type: type[Agent] | None = None,
+            self,
+            filter_func: Callable[[Agent], bool] | None = None,
+            n: int = 0,
+            inplace: bool = False,
+            agent_type: type[Agent] | None = None,
     ) -> AgentSet:
         """
         Select a subset of agents from the AgentSet based on a filter function and/or quantity limit.
@@ -145,7 +145,7 @@ def agent_generator(filter_func=None, agent_type=None, n=0):
             count = 0
             for agent in self:
                 if (not filter_func or filter_func(agent)) and (
-                    not agent_type or isinstance(agent, agent_type)
+                        not agent_type or isinstance(agent, agent_type)
                 ):
                     yield agent
                     count += 1
@@ -182,10 +182,10 @@ def shuffle(self, inplace: bool = False) -> AgentSet:
             )
 
     def sort(
-        self,
-        key: Callable[[Agent], Any] | str,
-        ascending: bool = False,
-        inplace: bool = False,
+            self,
+            key: Callable[[Agent], Any] | str,
+            ascending: bool = False,
+            inplace: bool = False,
     ) -> AgentSet:
         """
         Sort the agents in the AgentSet based on a specified attribute or custom function.
@@ -218,7 +218,7 @@ def _update(self, agents: Iterable[Agent]):
         return self
 
     def do(
-        self, method: str | Callable, *args, return_results: bool = False, **kwargs
+            self, method: str | Callable, *args, return_results: bool = False, **kwargs
     ) -> AgentSet | list[Any]:
         """
         Invoke a method or function on each agent in the AgentSet.
@@ -412,14 +412,6 @@ class GroupBy:
     def __init__(self, groups: dict[Any, list | AgentSet]):
         self.groups: dict[Any, list | AgentSet] = groups
 
-    def get_group(self, name: Any):
-        """
-        Retrieve the specified group by name.
-        """
-
-        # return group for specified name
-        return self.groups[name]
-
     def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
         """Apply the specified callable to each group
 
@@ -434,6 +426,10 @@ def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
         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 {
@@ -454,15 +450,19 @@ def do(self, method: Callable | str, *args, **kwargs) -> GroupBy:
                                     Additional arguments and keyword arguments will be passed on to the callable.
 
         Returns:
-            GroupBy
+            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)
+            for v in self.groups.values():
+                method(v, *args, **kwargs)
 
         return self
 
diff --git a/tests/test_agent.py b/tests/test_agent.py
index f2bbd201c7e..f5eaefba55d 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -364,12 +364,13 @@ def get_unique_identifier(self):
     agentset = AgentSet(agents, model)
 
     groups = agentset.groupby("even")
-    assert len(groups.get_group(True)) == 5
-    assert len(groups.get_group(False)) == 5
+    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.get_group(True)) == 5
-    assert len(groups.get_group(False)) == 5
+    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
@@ -381,3 +382,12 @@ def get_unique_identifier(self):
     attributes = agentset.groupby("even", result_type="agentset").apply("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
+

From 5544cba121297bcf749e486880cd69f5027fe398 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Aug 2024 16:56:23 +0000
Subject: [PATCH 23/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mesa/agent.py       | 24 ++++++++++++------------
 tests/test_agent.py |  1 -
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 79f1bef1cfe..25fb2db754b 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -118,11 +118,11 @@ def __contains__(self, agent: Agent) -> bool:
         return agent in self._agents
 
     def select(
-            self,
-            filter_func: Callable[[Agent], bool] | None = None,
-            n: int = 0,
-            inplace: bool = False,
-            agent_type: type[Agent] | None = None,
+        self,
+        filter_func: Callable[[Agent], bool] | None = None,
+        n: int = 0,
+        inplace: bool = False,
+        agent_type: type[Agent] | None = None,
     ) -> AgentSet:
         """
         Select a subset of agents from the AgentSet based on a filter function and/or quantity limit.
@@ -145,7 +145,7 @@ def agent_generator(filter_func=None, agent_type=None, n=0):
             count = 0
             for agent in self:
                 if (not filter_func or filter_func(agent)) and (
-                        not agent_type or isinstance(agent, agent_type)
+                    not agent_type or isinstance(agent, agent_type)
                 ):
                     yield agent
                     count += 1
@@ -182,10 +182,10 @@ def shuffle(self, inplace: bool = False) -> AgentSet:
             )
 
     def sort(
-            self,
-            key: Callable[[Agent], Any] | str,
-            ascending: bool = False,
-            inplace: bool = False,
+        self,
+        key: Callable[[Agent], Any] | str,
+        ascending: bool = False,
+        inplace: bool = False,
     ) -> AgentSet:
         """
         Sort the agents in the AgentSet based on a specified attribute or custom function.
@@ -218,7 +218,7 @@ def _update(self, agents: Iterable[Agent]):
         return self
 
     def do(
-            self, method: str | Callable, *args, return_results: bool = False, **kwargs
+        self, method: str | Callable, *args, return_results: bool = False, **kwargs
     ) -> AgentSet | list[Any]:
         """
         Invoke a method or function on each agent in the AgentSet.
@@ -428,7 +428,7 @@ def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
 
         Notes:
             this method is useful for methods or functions that do return something. It
-            will break method chaining. For that, use ``do`` instead. 
+            will break method chaining. For that, use ``do`` instead.
 
         """
         if isinstance(method, str):
diff --git a/tests/test_agent.py b/tests/test_agent.py
index f5eaefba55d..3420ab637e3 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -390,4 +390,3 @@ def get_unique_identifier(self):
     groups = agentset.groupby("even", result_type="agentset")
     another_ref_to_groups = groups.do(lambda x: x.do("step"))
     assert groups == another_ref_to_groups
-

From e075bcb5de6b497a599066f6543cd287e7a66ef0 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Thu, 22 Aug 2024 10:41:36 +0200
Subject: [PATCH 24/25] Change apply to map

---
 mesa/agent.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 25fb2db754b..4267f345f18 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -412,8 +412,8 @@ class GroupBy:
     def __init__(self, groups: dict[Any, list | AgentSet]):
         self.groups: dict[Any, list | AgentSet] = groups
 
-    def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
-        """Apply the specified callable to each group
+    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,

From c72aa37b578b76c4494ec7c9e9163ce97e664832 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Thu, 22 Aug 2024 10:44:22 +0200
Subject: [PATCH 25/25] also update tests

---
 tests/test_agent.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_agent.py b/tests/test_agent.py
index 3420ab637e3..5e3a5f2e7cb 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -376,10 +376,10 @@ def get_unique_identifier(self):
         assert len(group) == 5
         assert group_name in {True, False}
 
-    sizes = agentset.groupby("even", result_type="list").apply(len)
+    sizes = agentset.groupby("even", result_type="list").map(len)
     assert sizes == {True: 5, False: 5}
 
-    attributes = agentset.groupby("even", result_type="agentset").apply("get", "even")
+    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)