Repository: projectmesa/mesa
Pull request #2304 (2024-09-20 07:10:02)

Solara viz: use_task for non-threaded continous play
This is an alternative to https://github.com/projectmesa/mesa/pull/2303 

While investigating #2303 I realized that `use_thread` is deprecated and recommended to replace with `use_task`, which also allows not using a thread at all, but a coroutine. This allows this version to also work in environments where threading can be problematic (colab) or isn't available at all (Pyodide, WASM, etc.). It also avoids possible issues with threads running in the background forever. 

It also runs blazingly fast!

[screen-capture (1).webm](https://github.com/user-attachments/assets/857e4ea7-519a-491c-bf5e-52fc87a7c40c)


github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔴 +6.8% [+5.0%, +8.6%] | 🔴 +8.0% [+6.8%, +9.2%] |
| BoltzmannWealth | large  | 🔵 -1.4% [-2.9%, -0.5%] | 🔵 -0.4% [-0.8%, +0.0%] |
| Schelling       | small  | 🔵 -0.2% [-0.4%, +0.0%] | 🔵 -0.2% [-0.4%, -0.1%] |
| Schelling       | large  | 🔵 -0.4% [-0.5%, -0.1%] | 🔵 -0.0% [-0.3%, +0.2%] |
| WolfSheep       | small  | 🔵 -0.4% [-0.7%, -0.2%] | 🔵 -0.2% [-0.4%, -0.1%] |
| WolfSheep       | large  | 🔵 -0.4% [-0.8%, +0.0%] | 🔵 -0.4% [-0.9%, +0.0%] |
| BoidFlockers    | small  | 🔵 -0.5% [-0.9%, -0.2%] | 🔵 +0.5% [-0.2%, +1.2%] |
| BoidFlockers    | large  | 🔵 -0.2% [-0.6%, +0.1%] | 🔵 -0.4% [-1.3%, +0.5%] |


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())
 
diff --git a/mesa/visualization/solara_viz.py b/mesa/visualization/solara_viz.py
index 129a9486..ed495732 100644
--- a/mesa/visualization/solara_viz.py
+++ b/mesa/visualization/solara_viz.py
@@ -23,8 +23,8 @@ See the Visualization Tutorial and example models for more details.
 
 from __future__ import annotations
 
+import asyncio
 import copy
-import time
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Literal
 
@@ -235,12 +235,12 @@ def ModelController(model: solara.Reactive[Model], play_interval=100):
 
     solara.use_effect(save_initial_model, [model.value])
 
-    def step():
+    async def step():
         while playing.value:
-            time.sleep(play_interval / 1000)
+            await asyncio.sleep(play_interval / 1000)
             do_step()
 
-    solara.use_thread(step, [playing.value])
+    solara.lab.use_task(step, dependencies=[playing.value], prefer_threaded=False)
 
     def do_step():
         """Advance the model by one step."""


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.