Repository: projectmesa/mesa
Pull request #2350 (2024-10-13 08:07:34)

replace model with random in AgentSet init
This PR replaces the model argument in `AgentSet.__init__` with `random` as a keyword argument as argued in #2323.

**Motivation:**
* the only thing retrieved from model is the random number generator. Passing `model` serves no other function and there is no other use case for having model in the agentset. 
* For advanced use cases, explicitly passing random allows users to control which random number generator is being used. The most common reason for wanting this is for variance reduction through common random numbers. See #2323 and links provided there for a further elaboration on this argument.
* I cannot currently change `CellCollection.agents` to returning an agentset because CellCollection does not have a reference to model even though it does have access to its random number generator. 

Some have argued that this change creates extra load on the user. I disagree. First, In none of the examples in mesa-examples is the user creating an AgentSet object themselves. This highlights that it will be rare for users to create AgentSet themselves directly in all but the most advanced use cases. It also implies that this change is almost excusively confined to the inner workings of mesa. Second, this API is identical to the new gridspaces were we pass random explicitly instead of passing model only to get access to the central random number generator. So, this PR in fact makes the code base internally more coherent and therefore easier to use.

Another argument against this change is that we might potentially expand in the future AgentSet to use other model attributes. Again, I am not convinced by this unless clear examples of useful AgentSet functionality are given that require access to Model. We have had AgentSet for almost a year. No example has come up, I cannot come up with a convincing example, and None have not been offered.



github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔵 -3.3% [-4.0%, -2.5%] | 🔵 +0.3% [+0.2%, +0.4%] |
| BoltzmannWealth | large  | 🔵 -0.7% [-1.1%, -0.4%] | 🔵 +0.4% [-0.1%, +0.8%] |
| Schelling       | small  | 🔵 -0.1% [-0.7%, +0.5%] | 🔵 -0.5% [-0.7%, -0.3%] |
| Schelling       | large  | 🔵 -0.7% [-1.5%, +0.3%] | 🔵 -1.3% [-1.9%, -0.7%] |
| WolfSheep       | small  | 🔵 -0.7% [-0.9%, -0.5%] | 🔵 +0.2% [-0.0%, +0.4%] |
| WolfSheep       | large  | 🔵 -0.5% [-2.1%, +0.6%] | 🔵 -0.5% [-2.0%, +1.0%] |
| BoidFlockers    | small  | 🔵 -0.1% [-0.4%, +0.2%] | 🔵 +1.8% [+1.1%, +2.5%] |
| BoidFlockers    | large  | 🔵 -0.3% [-1.4%, +0.8%] | 🔵 +0.9% [+0.4%, +1.4%] |

rht: > I cannot currently change CellCollection.agents to returning an agentset because CellCollection does not have a reference to model even though it does have access to its random number generator.
> So, this PR in fact makes the code base internally more coherent and therefore easier to use.

Having access to the model everywhere is arguably also coherent and easy to use. So it could go either way.

> We have had AgentSet for almost a year. No example has come up, I cannot come up with a convincing example, and None have not been offered.

> This was the first thing I was thinking about. However, the AgentSet has been with us since 2.2, is very actively integrated and developed and so far we never needed it. So a few months back I might give this argument a lot of weight, but now I think we will perfectly manage without it.

It's because the use cases (the "experimental evidences") are limited to the existing mesa-examples, and we haven't received feedback from advanced users yet. This is in contrast to the datacollector where we received lots of requests for multi-agent data collection.

But I'm fine with this PR. It can always be changed back in the future when needed. The added bonus with this PR is that I have observed that AgentSet's utility may extend beyond ABM, as an arbitrary object DataFrame, and so tying it to an ABM model might actually be limiting it.

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index 4d28f474..2b6d2f99 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -99,14 +99,18 @@ class AgentSet(MutableSet, Sequence):
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
     """
 
-    def __init__(self, agents: Iterable[Agent], model: Model):
+    def __init__(self, agents: Iterable[Agent], random: Random | None = None):
         """Initializes the AgentSet with a collection of agents and a reference to the model.
 
         Args:
             agents (Iterable[Agent]): An iterable of Agent objects to be included in the set.
-            model (Model): The ABM model instance to which this AgentSet belongs.
+            random (Random): the random number generator
         """
-        self.model = model
+        if random is None:
+            random = (
+                Random()
+            )  # FIXME see issue 1981, how to get the central rng from model
+        self.random = random
         self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
 
     def __len__(self) -> int:
@@ -177,7 +181,7 @@ class AgentSet(MutableSet, Sequence):
 
         agents = agent_generator(filter_func, agent_type, at_most)
 
-        return AgentSet(agents, self.model) if not inplace else self._update(agents)
+        return AgentSet(agents, self.random) if not inplace else self._update(agents)
 
     def shuffle(self, inplace: bool = False) -> AgentSet:
         """Randomly shuffle the order of agents in the AgentSet.
@@ -200,7 +204,7 @@ class AgentSet(MutableSet, Sequence):
             return self
         else:
             return AgentSet(
-                (agent for ref in weakrefs if (agent := ref()) is not None), self.model
+                (agent for ref in weakrefs if (agent := ref()) is not None), self.random
             )
 
     def sort(
@@ -225,7 +229,7 @@ class AgentSet(MutableSet, Sequence):
         sorted_agents = sorted(self._agents.keys(), key=key, reverse=not ascending)
 
         return (
-            AgentSet(sorted_agents, self.model)
+            AgentSet(sorted_agents, self.random)
             if not inplace
             else self._update(sorted_agents)
         )
@@ -477,7 +481,7 @@ class AgentSet(MutableSet, Sequence):
         Returns:
             dict: A dictionary representing the state of the AgentSet.
         """
-        return {"agents": list(self._agents.keys()), "model": self.model}
+        return {"agents": list(self._agents.keys()), "random": self.random}
 
     def __setstate__(self, state):
         """Set the state of the AgentSet during deserialization.
@@ -485,18 +489,9 @@ class AgentSet(MutableSet, Sequence):
         Args:
             state (dict): A dictionary representing the state to restore.
         """
-        self.model = state["model"]
+        self.random = state["random"]
         self._update(state["agents"])
 
-    @property
-    def random(self) -> Random:
-        """Provide access to the model's random number generator.
-
-        Returns:
-            Random: The random number generator associated with the model.
-        """
-        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.
 
@@ -529,7 +524,7 @@ class AgentSet(MutableSet, Sequence):
 
         if result_type == "agentset":
             return GroupBy(
-                {k: AgentSet(v, model=self.model) for k, v in groups.items()}
+                {k: AgentSet(v, random=self.random) for k, v in groups.items()}
             )
         else:
             return GroupBy(groups)
diff --git a/mesa/model.py b/mesa/model.py
index e4f70793..b6500eb4 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -52,8 +52,6 @@ class Model:
         self.running = True
         self.steps: int = 0
 
-        self._setup_agent_registration()
-
         self._seed = seed
         if self._seed is None:
             # We explicitly specify the seed here so that we know its value in
@@ -65,6 +63,9 @@ class Model:
         self._user_step = self.step
         self.step = self._wrapped_step
 
+        # setup agent registration data structures
+        self._setup_agent_registration()
+
     def _wrapped_step(self, *args: Any, **kwargs: Any) -> None:
         """Automatically increments time and steps after calling the user's step method."""
         # Automatically increment time and step counters
@@ -119,7 +120,9 @@ class Model:
         self._agents_by_type: dict[
             type[Agent], AgentSet
         ] = {}  # a dict with an agentset for each class of agents
-        self._all_agents = AgentSet([], self)  # an agenset with all agents
+        self._all_agents = AgentSet(
+            [], random=self.random
+        )  # an agenset with all agents
 
     def register_agent(self, agent):
         """Register the agent with the model.
@@ -153,7 +156,7 @@ class Model:
                 [
                     agent,
                 ],
-                self,
+                random=self.random,
             )
 
         self._all_agents.add(agent)


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 73aab897..aa5ff935 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -62,7 +62,7 @@ def test_agentset():
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
 
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     assert agents[0] in agentset
     assert len(agentset) == len(agents)
@@ -118,7 +118,7 @@ def test_agentset():
 
     # because AgentSet uses weakrefs, we need hard refs as well....
     other_agents, another_set = pickle.loads(  # noqa: S301
-        pickle.dumps([agents, AgentSet(agents, model)])
+        pickle.dumps([agents, AgentSet(agents, random=model.random)])
     )
     assert all(
         a1.unique_id == a2.unique_id for a1, a2 in zip(another_set, other_agents)
@@ -129,11 +129,11 @@ def test_agentset():
 def test_agentset_initialization():
     """Test agentset initialization."""
     model = Model()
-    empty_agentset = AgentSet([], model)
+    empty_agentset = AgentSet([], random=model.random)
     assert len(empty_agentset) == 0
 
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     assert len(agentset) == 10
 
 
@@ -141,7 +141,7 @@ def test_agentset_serialization():
     """Test pickleability of agentset."""
     model = Model()
     agents = [AgentTest(model) for _ in range(5)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     serialized = pickle.dumps(agentset)
     deserialized = pickle.loads(serialized)  # noqa: S301
@@ -156,7 +156,7 @@ def test_agent_membership():
     """Test agent membership in AgentSet."""
     model = Model()
     agents = [AgentTest(model) for _ in range(5)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     assert agents[0] in agentset
     assert AgentTest(model) not in agentset
@@ -166,7 +166,7 @@ def test_agent_add_remove_discard():
     """Test adding, removing and discarding agents from AgentSet."""
     model = Model()
     agent = AgentTest(model)
-    agentset = AgentSet([], model)
+    agentset = AgentSet([], random=model.random)
 
     agentset.add(agent)
     assert agent in agentset
@@ -186,7 +186,7 @@ def test_agentset_get_item():
     """Test integer based access to AgentSet."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     assert agentset[0] == agents[0]
     assert agentset[-1] == agents[-1]
@@ -200,7 +200,7 @@ def test_agentset_do_str():
     """Test AgentSet.do with str."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     with pytest.raises(AttributeError):
         agentset.do("non_existing_method")
@@ -213,7 +213,7 @@ def test_agentset_do_str():
     n = 10
     model = Model()
     agents = [AgentDoTest(model) for _ in range(n)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -223,7 +223,7 @@ def test_agentset_do_str():
     # setup
     model = Model()
     agents = [AgentDoTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -235,7 +235,7 @@ def test_agentset_do_callable():
     """Test AgentSet.do with callable."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     # Test callable with non-existent function
     with pytest.raises(AttributeError):
@@ -249,7 +249,7 @@ def test_agentset_do_callable():
     n = 10
     model = Model()
     agents = [AgentDoTest(model) for _ in range(n)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -260,7 +260,7 @@ def test_agentset_do_callable():
     # setup again for lambda function tests
     model = Model()
     agents = [AgentDoTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -278,7 +278,7 @@ def test_agentset_do_callable():
     # setup again for actual function tests
     model = Model()
     agents = [AgentDoTest(model) for _ in range(n)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -289,7 +289,7 @@ def test_agentset_do_callable():
     # setup again for actual function tests
     model = Model()
     agents = [AgentDoTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -354,7 +354,7 @@ def test_agentset_agg():
         agent.energy = i + 1
         agent.wealth = 10 * (i + 1)
 
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     # Test min aggregation
     min_energy = agentset.agg("energy", min)
@@ -391,7 +391,7 @@ def test_agentset_set_method():
 
     model = Model()
     agents = [TestAgentWithAttribute(model, age=i) for i in range(5)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     # Set a new attribute "health" and an existing attribute "age" for all agents
     agentset.set("health", 100).set("age", 50).set("status", "active")
@@ -410,7 +410,7 @@ def test_agentset_map_str():
     """Test AgentSet.map with strings."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     with pytest.raises(AttributeError):
         agentset.do("non_existing_method")
@@ -423,7 +423,7 @@ def test_agentset_map_callable():
     """Test AgentSet.map with callable."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     # Test callable with non-existent function
     with pytest.raises(AttributeError):
@@ -450,7 +450,7 @@ def test_agentset_shuffle_do():
             self.called = True
 
     agents = [TestAgentShuffleDo(model) for _ in range(100)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     # Test shuffle_do with a string method name
     agentset.shuffle_do("test_method")
@@ -477,7 +477,7 @@ def test_agentset_get_attribute():
     """Test AgentSet.get for attributes."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     unique_ids = agentset.get("unique_id")
     assert unique_ids == [agent.unique_id for agent in agents]
@@ -491,7 +491,7 @@ def test_agentset_get_attribute():
         agent = AgentTest(model)
         agent.i = i**2
         agents.append(agent)
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     values = agentset.get(["unique_id", "i"])
 
@@ -521,7 +521,7 @@ def test_agentset_select_by_type():
 
     # Combine the two types of agents
     mixed_agents = test_agents + other_agents
-    agentset = AgentSet(mixed_agents, model)
+    agentset = AgentSet(mixed_agents, random=model.random)
 
     # Test selection by type
     selected_test_agents = agentset.select(agent_type=AgentTest)
@@ -544,11 +544,11 @@ def test_agentset_shuffle():
     model = Model()
     test_agents = [AgentTest(model) for _ in range(12)]
 
-    agentset = AgentSet(test_agents, model=model)
+    agentset = AgentSet(test_agents, random=model.random)
     agentset = agentset.shuffle()
     assert not all(a1 == a2 for a1, a2 in zip(test_agents, agentset))
 
-    agentset = AgentSet(test_agents, model=model)
+    agentset = AgentSet(test_agents, random=model.random)
     agentset.shuffle(inplace=True)
     assert not all(a1 == a2 for a1, a2 in zip(test_agents, agentset))
 
@@ -567,7 +567,7 @@ def test_agentset_groupby():
 
     model = Model()
     agents = [TestAgent(model) for _ in range(10)]
-    agentset = AgentSet(agents, model)
+    agentset = AgentSet(agents, random=model.random)
 
     groups = agentset.groupby("even")
     assert len(groups.groups[True]) == 5