Merge remote-tracking branch 'upstream/main'

Author: quaquel

.github/PULL_REQUEST_TEMPLATE/bug.md

@@ -0,0 +1,16 @@
+## Summary
+<!-- Provide a brief summary of the bug and its impact. -->
+
+## Bug / Issue
+<!-- Link to the related issue(s) and describe the bug. Include details like the context, what was expected, and what actually happened. -->
+
+## Implementation
+<!-- Describe the changes made to resolve the issue. Highlight any important parts of the code that were modified. -->
+
+## Testing
+<!-- Detail the testing performed to verify the fix. Include information on test cases, steps taken, and any relevant results.
+
+If you're fixing the visualisation, add before/after screenshots. -->
+
+## Additional Notes
+<!-- Add any additional information that may be relevant for the reviewers, such as potential side effects, dependencies, or related work.

.github/PULL_REQUEST_TEMPLATE/feature.md

@@ -0,0 +1,16 @@
+## Summary
+<!-- Provide a concise summary of the feature and its purpose. -->
+
+## Motive
+<!-- Explain the reasoning behind this feature. Include details on the problem it addresses or the enhancement it provides. -->
+
+## Implementation
+<!-- Describe how the feature was implemented. Include details on the approach taken, important decisions made, and code changes. -->
+
+## Usage Examples
+<!-- Provide code snippets or examples demonstrating how to use the new feature. Highlight key scenarios where this feature will be beneficial.
+
+If you're modifying the visualisation, add before/after screenshots. -->
+
+## Additional Notes
+<!-- Add any additional information that may be relevant for the reviewers, such as potential side effects, dependencies, or related work. -->

mesa/agent.py

@@ -11,7 +11,9 @@
 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 @@ def _update(self, agents: Iterable[Agent]):
         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 @@ def do(
                 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 @@ def random(self) -> Random:
         """
         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)

mesa/visualization/solara_viz.py

@@ -158,7 +158,11 @@ def do_reseed():
         """Update the random seed for the model."""
         reactive_seed.value = model.random.random()
 
-    dependencies = [current_step.value, reactive_seed.value]
+    dependencies = [
+        *list(model_parameters.values()),
+        current_step.value,
+        reactive_seed.value,
+    ]
 
     # if space drawer is disabled, do not include it
     layout_types = [{"Space": "default"}] if space_drawer else []

tests/test_agent.py

@@ -83,12 +83,6 @@ def test_function(agent):
     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 remove_function(agent):
     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