Repository: projectmesa/mesa
Pull request #2290 (2024-09-10 08:02:17)

GroupBy: Add `count` and `agg` methods
Added two new methods to the `GroupBy` class:

- `count`: Returns the count of agents in each group.
- `agg`: Performs aggregation on a specific attribute across groups, applying a function like `sum`, `min`, `max`, or `mean`.

### Usage examples
1. **`count()`**: 
   - Returns the count of agents in each group.
   - **Usage Example:**
     ```python
     grouped_agents.count()
     ```
     **Output:**
     ```python
     {
         "group_1": 10,  # 10 agents in group_1
         "group_2": 8,   # 8 agents in group_2
     }
     ```

2. **`agg(attr_name, func)`**:
   - Aggregates the values of a specific attribute across each group using a provided function (e.g., `sum`, `min`, `max`, `mean`).
   - **Usage Example:**
     ```python
     # Sum the wealth of all agents in each group
     grouped_agents.agg("wealth", sum)
     ```
     **Output:**
     ```python
     {
         "group_1": 15000,  # Total wealth of agents in group_1
         "group_2": 12000,  # Total wealth of agents in group_2
     }
     ```

<hr>

I'm a bit worried about our AgentSet and now GroupBy exploding our code base.
1. Is there any built-in functionality or libraries for this? It seems we're implementing very obvious behavior ourselves here.
2. Do we need to split into multiple files?
3. How do we keep this properly documented? Do we need guides for this?

(obviously tests are still needed)

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]          |
|-----------------|--------|-------------------------|----------------------------|
| BoltzmannWealth | small  | 🔵 -1.7% [-2.5%, -0.9%] | 🔵 -0.4% [-0.5%, -0.3%]    |
| BoltzmannWealth | large  | 🔵 +0.4% [-0.6%, +1.4%] | 🔵 +0.0% [-2.9%, +3.3%]    |
| Schelling       | small  | 🔵 +0.2% [-0.1%, +0.4%] | 🔵 -0.3% [-0.5%, -0.0%]    |
| Schelling       | large  | 🔵 -0.4% [-1.1%, +0.3%] | 🔵 -2.4% [-3.7%, -1.1%]    |
| WolfSheep       | small  | 🔵 -2.0% [-2.9%, -1.2%] | 🔵 -2.2% [-2.6%, -1.9%]    |
| WolfSheep       | large  | 🟢 -6.0% [-7.3%, -4.9%] | 🟢 -12.7% [-14.6%, -10.9%] |
| BoidFlockers    | small  | 🔵 -0.3% [-0.7%, +0.3%] | 🔵 +1.2% [+0.6%, +1.8%]    |
| BoidFlockers    | large  | 🔵 +0.1% [-0.8%, +1.1%] | 🔵 +1.2% [+0.5%, +1.8%]    |

quaquel: 1. Not sure you need `count`. You can already do it via `some_agentset.groupby(type).map(len)`. So, if you are worried about the number of lines of code in the code base, you could leave it out. I personally would not worry too much about the number of lines of code and am open to adding some shorthands for frequently used patterns (like getting the number of agents of a particular type or that meet a given criterion). 
2. `agg` makes good sense and is a useful addition. I think I even mentioned adding it to groupby while adding it to agentset.
EwoutH: Added tests and cleaned up commits, ready for final review.

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index 5badc825..2d098c35 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -15,7 +15,7 @@ import operator
 import warnings
 import weakref
 from collections import defaultdict
-from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
+from collections.abc import Callable, Hashable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
 # mypy
@@ -611,6 +611,29 @@ class GroupBy:
 
         return self
 
+    def count(self) -> dict[Any, int]:
+        """Return the count of agents in each group.
+
+        Returns:
+            dict: A dictionary mapping group names to the number of agents in each group.
+        """
+        return {k: len(v) for k, v in self.groups.items()}
+
+    def agg(self, attr_name: str, func: Callable) -> dict[Hashable, Any]:
+        """Aggregate the values of a specific attribute across each group using the provided function.
+
+        Args:
+            attr_name (str): The name of the attribute to aggregate.
+            func (Callable): The function to apply (e.g., sum, min, max, mean).
+
+        Returns:
+            dict[Hashable, Any]: A dictionary mapping group names to the result of applying the aggregation function.
+        """
+        return {
+            group_name: func([getattr(agent, attr_name) for agent in group])
+            for group_name, group in self.groups.items()
+        }
+
     def __iter__(self):  # noqa: D105
         return iter(self.groups.items())
 


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 769be4ec..f43a80c8 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -524,6 +524,7 @@ def test_agentset_groupby():
         def __init__(self, model):
             super().__init__(model)
             self.even = self.unique_id % 2 == 0
+            self.value = self.unique_id * 10
 
         def get_unique_identifier(self):
             return self.unique_id
@@ -560,6 +561,37 @@ def test_agentset_groupby():
     another_ref_to_groups = groups.do(lambda x: x.do("step"))
     assert groups == another_ref_to_groups
 
+    # New tests for count() method
+    groups = agentset.groupby("even")
+    count_result = groups.count()
+    assert count_result == {True: 5, False: 5}
+
+    # New tests for agg() method
+    groups = agentset.groupby("even")
+    sum_result = groups.agg("value", sum)
+    assert sum_result[True] == sum(agent.value for agent in agents if agent.even)
+    assert sum_result[False] == sum(agent.value for agent in agents if not agent.even)
+
+    max_result = groups.agg("value", max)
+    assert max_result[True] == max(agent.value for agent in agents if agent.even)
+    assert max_result[False] == max(agent.value for agent in agents if not agent.even)
+
+    min_result = groups.agg("value", min)
+    assert min_result[True] == min(agent.value for agent in agents if agent.even)
+    assert min_result[False] == min(agent.value for agent in agents if not agent.even)
+
+    # Test with a custom aggregation function
+    def custom_agg(values):
+        return sum(values) / len(values) if values else 0
+
+    custom_result = groups.agg("value", custom_agg)
+    assert custom_result[True] == custom_agg(
+        [agent.value for agent in agents if agent.even]
+    )
+    assert custom_result[False] == custom_agg(
+        [agent.value for agent in agents if not agent.even]
+    )
+
 
 def test_oldstyle_agent_instantiation():
     """Old behavior of Agent creation with unique_id and model as positional arguments.