Repository: projectmesa/mesa
Pull request #2254 (2024-08-28 10:16:21)

AgentSet: Add `set` method
This PR adds a `set` method to the `AgentSet` class that allows setting a specified attribute of all agents within the set to a given value. This method provides a streamlined way to update agent attributes in bulk, without having to resort to `lambda` or other callable functions.

### Motive
The motivation behind this feature is to simplify and expedite the process of modifying attributes across all agents in an `AgentSet`. Previously, this required iterating manually over the agents or using a `lambda` or other callable function. This new method makes the operation more elegant.

### Implementation
The `set` method was added to the `AgentSet` class. It iterates through all agents in the set and applies the `setattr` function to assign the specified value to the desired attribute. The method operates in place, modifying the existing agents directly.

The (now updated) AgentSet itself is returned, which allows for chaining.

You can both set existing Agent variables or set new ones.

### Usage Examples
```python
# Set the 'has_license' attribute to True for all agents in the AgentSet
agents.set('has_license', True)
```
Note that it doesn't matter if `has_license` already exists in the Agent, both work.

Together with #2253, you can now set a value for a fraction of your AgentSet:
```Python
# Select 40% of the agents from the AgentSet
model.agents.select(p=0.4).set('has_license', True)
```
It's also possible to set multiple things by chaining the `set()` commands:
```python
agents.set('has_license', True).set('age', 45).set('continent', 'Europe')
```
### Additional Notes
It feels logical that if you have a `get` method, you also should have a `set` method.

**Edit:**
- Modified to return an AgentSet, so you can keep the output or keep chaining.
- It modified the Agents itself, like with `do()`, so inplace isn't needed.
- We don't allow dicts, those can be added later if needed.


github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔴 +4.7% [+3.2%, +6.2%] | 🔵 +0.6% [+0.5%, +0.8%] |
| BoltzmannWealth | large  | 🔵 -0.1% [-0.9%, +0.6%] | 🔵 +4.3% [+2.8%, +5.8%] |
| Schelling       | small  | 🔵 +0.8% [+0.4%, +1.2%] | 🔵 -0.0% [-0.3%, +0.2%] |
| Schelling       | large  | 🔵 +0.6% [-0.2%, +1.4%] | 🔴 +4.1% [+3.1%, +5.0%] |
| WolfSheep       | small  | 🔵 +0.5% [-0.9%, +2.0%] | 🔵 +1.2% [+0.9%, +1.4%] |
| WolfSheep       | large  | 🔵 +0.8% [+0.4%, +1.2%] | 🔴 +6.7% [+5.4%, +7.9%] |
| BoidFlockers    | small  | 🔵 +3.0% [+2.2%, +3.8%] | 🔵 +0.7% [-0.4%, +1.8%] |
| BoidFlockers    | large  | 🔵 +1.7% [+1.1%, +2.4%] | 🔵 +1.5% [+0.7%, +2.2%] |

quaquel: good idea, and probably an inplace keyword should be included preferably in this PR. It keeps the API clear and consistent across methods. 
EwoutH: Can we think of cases where you want to continue the chain and return an AgentSet? In that case, should a default `return=False` which you can switch to `True` be also included?

Maybe I'm now mixing things up with in place.
EwoutH: I see three scenarios here:
1. I want to modify my agentset inplace
2. I want to modify my agentset inplace and return that updated agentset (to assign or continue chaining)
3. I want to create a new agentset (inplace) and return that new agentset (to assign or continue chaining)

Of course you can just ignore the return for 1. So return is not really needed, just return it and the user can decide if the y use it or not. Okay, I get it.

So we probably need both

Now we're in a bit of a pickle, because with `.set()` `inplace=True` seems like a good default. But most other functions have it `False` by default.

I think the most similar function is `.do()`. That also only does it only in place. This is basically a `.do()` with less steps.

To summarize:
- Always return the new AgentSet
- We need to pick one of two options:
  1. Like `do()`, and just always do it "in place"
  2. Add an `inplace` defaulting to `True`
quaquel: good analysis. There is no point to return a new agent set, just as with do. 
EwoutH: Thanks, I will add a few unittest tonight and merge.
EwoutH: What I just now realized and what's quite powerful is that you also can set completely new variables this way.

Punching above it's weight!

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index f137ba09..39abb77d 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -316,6 +316,21 @@ class AgentSet(MutableSet, Sequence):
                 for agent in self._agents
             ]
 
+    def set(self, attr_name: str, value: Any) -> AgentSet:
+        """
+        Set a specified attribute to a given value for all agents in the AgentSet.
+
+        Args:
+            attr_name (str): The name of the attribute to set.
+            value (Any): The value to set the attribute to.
+
+        Returns:
+            AgentSet: The AgentSet instance itself, after setting the attribute.
+        """
+        for agent in self:
+            setattr(agent, attr_name, value)
+        return self
+
     def __getitem__(self, item: int | slice) -> Agent:
         """
         Retrieve an agent or a slice of agents from the AgentSet.


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 516902f9..8c5bf351 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -270,6 +270,30 @@ def test_agentset_do_callable():
     assert len(agentset) == 0
 
 
+def test_agentset_set_method():
+    # Initialize the model and agents with and without existing attributes
+    class TestAgentWithAttribute(Agent):
+        def __init__(self, unique_id, model, age=None):
+            super().__init__(unique_id, model)
+            self.age = age
+
+    model = Model()
+    agents = [TestAgentWithAttribute(model.next_id(), model, age=i) for i in range(5)]
+    agentset = AgentSet(agents, model)
+
+    # Set a new attribute "health" and an existing attribute "age" for all agents
+    agentset.set("health", 100).set("age", 50).set("status", "active")
+
+    # Check if all agents have the "health", "age", and "status" attributes correctly set
+    for agent in agentset:
+        assert hasattr(agent, "health")
+        assert agent.health == 100
+        assert hasattr(agent, "age")
+        assert agent.age == 50
+        assert hasattr(agent, "status")
+        assert agent.status == "active"
+
+
 def test_agentset_map_str():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]