Repository: projectmesa/mesa
Pull request #2266 (2024-09-02 08:08:51)

AgentSet: Add `agg` method
This PR introduces the `agg` method to the `AgentSet` class, allowing users to apply aggregation functions (e.g., `min`, `max`, `sum`, `np.mean`) to attributes of agents within the `AgentSet`. This enhancement makes it easier to compute summary statistics across agent attributes directly within the `AgentSet` interface.

This will be useful in both the model operation itself as well as for future DataCollector use.

### New Method: `agg`

```python
def agg(self, attribute: str, func: Callable) -> Any:
```

**Parameters**:
- `attribute` (str): The name of the attribute to aggregate.
- `func` (Callable): The aggregation function to apply (e.g., `min`, `max`, `sum`, `np.mean`).

**Returns**:
- The result of applying the aggregation function to the attribute values of all agents in the `AgentSet`.

### Usage Examples

```python
# Get the minimum energy value
min_energy = agentset.agg("energy", min)

# Calculate the total energy of sheep in the model
total_energy = model.get_agents_of_type[Sheep].agg("energy", sum)

# Compute the average wealth of the poorest 10% using numpy
average_wealth = agentset.select(at_most=0.1).agg("wealth", np.mean)

# Custom aggregation function
def custom_func(values):
    return sum(values) / len(values)

custom_avg_energy = agentset.agg("energy", custom_func)
```
### Future work
This function could be expended by:
- Allowing multiple attribute as input for the functions
- Allowing multiple pairs of attribute-function to be inputted
- Allowing multiple pairs of multiple attributes / functions as input

However, this PRs limit the scope to a single attribute and single function, since that's the most common use case.

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]        | Run time [95% CI]       |
|-----------------|--------|---------------------------|-------------------------|
| BoltzmannWealth | small  | 🔵 +3.3% [+1.9%, +4.8%]   | 🔵 +0.0% [-0.2%, +0.2%] |
| BoltzmannWealth | large  | 🔵 +2.7% [-31.1%, +41.3%] | 🔵 +2.9% [+1.4%, +4.4%] |
| Schelling       | small  | 🔵 +0.1% [-0.1%, +0.4%]   | 🔵 +1.5% [+1.2%, +1.8%] |
| Schelling       | large  | 🔵 -0.9% [-1.9%, +0.1%]   | 🔵 +0.1% [-3.4%, +3.7%] |
| WolfSheep       | small  | 🔵 -1.1% [-2.4%, +0.3%]   | 🔵 -0.6% [-1.0%, -0.3%] |
| WolfSheep       | large  | 🔵 -2.6% [-3.2%, -2.0%]   | 🔵 -4.0% [-6.9%, -1.7%] |
| BoidFlockers    | small  | 🔵 +1.6% [+1.2%, +2.1%]   | 🔵 +0.4% [-0.2%, +1.0%] |
| BoidFlockers    | large  | 🔵 +1.7% [+0.7%, +2.6%]   | 🔵 +0.8% [-0.0%, +1.6%] |

EwoutH: While it will probably be used most of aggregation, in theory it could also be used for other things. Is it the most fitting name?
EwoutH: I just realized these are totally equivalent:
```Python
# Get the minimum energy value
min_energy = agentset.agg("energy", min)
min_energy = min(agentset.get("energy"))

# Calculate the total energy of sheep in the model
total_energy = model.get_agents_of_type[Sheep].agg("energy", sum)
total_energy = sum(model.get_agents_of_type[Sheep].agg("energy"))

# Compute the average wealth of the poorest 10% using numpy
average_wealth = agentset.select(at_most=0.1).agg("wealth", np.mean)
average_wealth = np.mean(agentset.select(at_most=0.1).agg("wealth"))
```
Is there something as too much syntactic sugar?

It does read nicely left to right though.
quaquel: Not sure about the last example, but yes, on the others, you are correct. If I read the last example correctly, you select 10% of the agents and then take their mean wealth, which is not equivalent to computing the average wealth of the poorest 10% using numpy. 
To do what you write requires something like

```python

agentset.sort("wealth").select(at_most=0.1).agg("wealth", np.mean)
np.mean(agentset.sort("wealth").select(at_most=0.1))
```

So for just one attribute and one aggregation function, we now have 2 ways to do it if we merge this PR. I think the real use case is when you extend this to multipel attributes/and or multiple callables.


Patch:
diff --git a/benchmarks/WolfSheep/wolf_sheep.py b/benchmarks/WolfSheep/wolf_sheep.py
index ecb78594..655d51bb 100644
--- a/benchmarks/WolfSheep/wolf_sheep.py
+++ b/benchmarks/WolfSheep/wolf_sheep.py
@@ -227,8 +227,8 @@ class WolfSheep(Model):
             patch.move_to(cell)
 
     def step(self):
-        self.get_agents_of_type(Sheep).shuffle(inplace=True).do("step")
-        self.get_agents_of_type(Wolf).shuffle(inplace=True).do("step")
+        self.agents_by_type[Sheep].shuffle(inplace=True).do("step")
+        self.agents_by_type[Wolf].shuffle(inplace=True).do("step")
 
 
 if __name__ == "__main__":
diff --git a/mesa/agent.py b/mesa/agent.py
index 9a4dd120..8d04df9e 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -304,6 +304,20 @@ class AgentSet(MutableSet, Sequence):
 
         return res
 
+    def agg(self, attribute: str, func: Callable) -> Any:
+        """
+        Aggregate an attribute of all agents in the AgentSet using a specified function.
+
+        Args:
+            attribute (str): The name of the attribute to aggregate.
+            func (Callable): The function to apply to the attribute values (e.g., min, max, sum, np.mean).
+
+        Returns:
+            Any: The result of applying the function to the attribute values. Often a single value.
+        """
+        values = self.get(attribute)
+        return func(values)
+
     def get(self, attr_names: str | list[str]) -> list[Any]:
         """
         Retrieve the specified attribute(s) from each agent in the AgentSet.
diff --git a/mesa/experimental/devs/examples/wolf_sheep.py b/mesa/experimental/devs/examples/wolf_sheep.py
index 9fa6b3d9..99fde6a9 100644
--- a/mesa/experimental/devs/examples/wolf_sheep.py
+++ b/mesa/experimental/devs/examples/wolf_sheep.py
@@ -231,8 +231,8 @@ class WolfSheep(mesa.Model):
             self.grid.place_agent(patch, pos)
 
     def step(self):
-        self.get_agents_of_type(Sheep).shuffle(inplace=True).do("step")
-        self.get_agents_of_type(Wolf).shuffle(inplace=True).do("step")
+        self.agents_by_type[Sheep].shuffle(inplace=True).do("step")
+        self.agents_by_type[Wolf].shuffle(inplace=True).do("step")
 
 
 if __name__ == "__main__":
diff --git a/mesa/model.py b/mesa/model.py
index 9387e6d1..f276cc82 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -33,11 +33,13 @@ class Model:
     Properties:
         agents: An AgentSet containing all agents in the model
         agent_types: A list of different agent types present in the model.
+        agents_by_type: A dictionary where the keys are agent types and the values are the corresponding AgentSets.
         steps: An integer representing the number of steps the model has taken.
                It increases automatically at the start of each step() call.
 
     Methods:
         get_agents_of_type: Returns an AgentSet of agents of the specified type.
+            Deprecated: Use agents_by_type[agenttype] instead.
         run_model: Runs the model's simulation until a defined end condition is reached.
         step: Executes a single step of the model's simulation process.
         next_id: Generates and returns the next unique identifier for an agent.
@@ -106,24 +108,26 @@ class Model:
         """Return a list of all unique agent types registered with the model."""
         return list(self._agents_by_type.keys())
 
-    def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
-        """Retrieves an AgentSet containing all agents of the specified type.
-
-        Args:
-            agenttype: The type of agent to retrieve.
-
-        Raises:
-            KeyError: If agenttype does not exist
-
+    @property
+    def agents_by_type(self) -> dict[type[Agent], AgentSet]:
+        """A dictionary where the keys are agent types and the values are the corresponding AgentSets."""
+        return self._agents_by_type
 
-        """
-        return self._agents_by_type[agenttype]
+    def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
+        """Deprecated: Retrieves an AgentSet containing all agents of the specified type."""
+        warnings.warn(
+            f"Model.get_agents_of_type() is deprecated, please replace get_agents_of_type({agenttype})"
+            f"with the property agents_by_type[{agenttype}].",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.agents_by_type[agenttype]
 
     def _setup_agent_registration(self):
         """helper method to initialize the agent registration datastructures"""
         self._agents = {}  # the hard references to all agents in the model
         self._agents_by_type: dict[
-            type, AgentSet
+            type[Agent], AgentSet
         ] = {}  # a dict with an agentset for each class of agents
         self._all_agents = AgentSet([], self)  # an agenset with all agents
 


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 972fb237..bb61ae1c 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -1,5 +1,6 @@
 import pickle
 
+import numpy as np
 import pytest
 
 from mesa.agent import Agent, AgentSet
@@ -276,6 +277,41 @@ def test_agentset_do_callable():
     assert len(agentset) == 0
 
 
+def test_agentset_agg():
+    model = Model()
+    agents = [TestAgent(i, model) for i in range(10)]
+
+    # Assign some values to attributes
+    for i, agent in enumerate(agents):
+        agent.energy = i + 1
+        agent.wealth = 10 * (i + 1)
+
+    agentset = AgentSet(agents, model)
+
+    # Test min aggregation
+    min_energy = agentset.agg("energy", min)
+    assert min_energy == 1
+
+    # Test max aggregation
+    max_energy = agentset.agg("energy", max)
+    assert max_energy == 10
+
+    # Test sum aggregation
+    total_energy = agentset.agg("energy", sum)
+    assert total_energy == sum(range(1, 11))
+
+    # Test mean aggregation using numpy
+    avg_wealth = agentset.agg("wealth", np.mean)
+    assert avg_wealth == 55.0
+
+    # Test aggregation with a custom function
+    def custom_func(values):
+        return sum(values) / len(values)
+
+    custom_avg_energy = agentset.agg("energy", custom_func)
+    assert custom_avg_energy == 5.5
+
+
 def test_agentset_set_method():
     # Initialize the model and agents with and without existing attributes
     class TestAgentWithAttribute(Agent):
diff --git a/tests/test_model.py b/tests/test_model.py
index 837d6b00..c92fcb2e 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1,4 +1,4 @@
-from mesa.agent import Agent
+from mesa.agent import Agent, AgentSet
 from mesa.model import Model
 
 
@@ -53,3 +53,19 @@ def test_agent_types():
     test_agent = TestAgent(model.next_id(), model)
     assert test_agent in model.agents
     assert type(test_agent) in model.agent_types
+
+
+def test_agents_by_type():
+    class Wolf(Agent):
+        pass
+
+    class Sheep(Agent):
+        pass
+
+    model = Model()
+    wolf = Wolf(1, model)
+    sheep = Sheep(2, model)
+
+    assert model.agents_by_type[Wolf] == AgentSet([wolf], model)
+    assert model.agents_by_type[Sheep] == AgentSet([sheep], model)
+    assert len(model.agents_by_type) == 2
diff --git a/tests/test_time.py b/tests/test_time.py
index a9d8e2e6..4203ca5d 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -320,7 +320,7 @@ class TestRandomActivationByType(TestCase):
         agent_types = model.agent_types
         for agent_type in agent_types:
             assert model.schedule.get_type_count(agent_type) == len(
-                model.get_agents_of_type(agent_type)
+                model.agents_by_type[agent_type]
             )
 
     # def test_add_non_unique_ids(self):