Fix: AgentSet initialization should not require explicit random number generator (#2789)

This PR modifies the AgentSet initialization to automatically use the random number generator from the model of the first agent when no explicit random parameter is provided, eliminating unnecessary UserWarnings in typical usage.

Previously, AgentSet initialization would raise a UserWarning whenever random=None, even though in most cases the agents being added to the set already have access to a seeded random number generator through their model. This resulted in users needing to explicitly pass random=model.random in every AgentSet creation, creating unnecessary boilerplate.

Co-authored-by: Jan Kwakkel <[email protected]>
Co-authored-by: Ewout ter Hoeven <[email protected]>

Author: 3 people

mesa/agent.py

@@ -162,30 +162,39 @@ class AgentSet(MutableSet, Sequence):
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
 
     Notes:
-        A `UserWarning` is issued if `random=None`. You can resolve this warning by explicitly
-        passing a random number generator. In most cases, this will be the seeded random number
-        generator in the model. So, you would do `random=self.random` in a `Model` or `Agent` instance.
+        If random is None then the random number generator in the model of the first agent is used.
+        If the agents list is empty and random is also None a user warning is issued and the AgentSet
+        is an empty list and a default random number generator.  This can make models non-reproducible.
+        If your code may create an AgentSet with no agents please pass a random number generator explicitly.
 
     """
 
-    def __init__(self, agents: Iterable[Agent], random: Random | None = None):
+    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.
-            random (Random): the random number generator
+            random (Random | np.random.Generator | None): the random number generator
         """
-        if random is None:
+        self._agents = weakref.WeakKeyDictionary(dict.fromkeys(agents))
+        if (len(self._agents) == 0) and random is None:
             warnings.warn(
-                "Random number generator not specified, this can make models non-reproducible. Please pass a random number generator explicitly",
+                "No Agents specified in creation of AgentSet and no random number generator specified. "
+                "This can make models non-reproducible. Please pass a random number generator explicitly",
                 UserWarning,
                 stacklevel=2,
             )
-            random = (
-                Random()
-            )  # FIXME see issue 1981, how to get the central rng from model
-        self.random = random
-        self._agents = weakref.WeakKeyDictionary(dict.fromkeys(agents))
+            random = Random()
+
+        if random is not None:
+            self.random = random
+        else:
+            # all agents in an AgentSet should share the same model, just take it from first
+            self.random = self._agents.keys().__next__().model.random
 
     def __len__(self) -> int:
         """Return the number of agents in the AgentSet."""

tests/test_agent.py

@@ -62,7 +62,7 @@ def test_agentset():
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
 
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     assert agents[0] in agentset
     assert len(agentset) == len(agents)
@@ -118,7 +118,7 @@ def test_function(agent):
 
     # because AgentSet uses weakrefs, we need hard refs as well....
     other_agents, another_set = pickle.loads(  # noqa: S301
-        pickle.dumps([agents, AgentSet(agents, random=model.random)])
+        pickle.dumps([agents, AgentSet(agents)])
     )
     assert all(
         a1.unique_id == a2.unique_id for a1, a2 in zip(another_set, other_agents)
@@ -131,17 +131,33 @@ def test_agentset_initialization():
     model = Model()
     empty_agentset = AgentSet([], random=model.random)
     assert len(empty_agentset) == 0
+    with pytest.warns(UserWarning):
+        empty_agentset2 = AgentSet([])
+    assert len(empty_agentset2) == 0
 
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     assert len(agentset) == 10
 
 
+def test_agentset_initialization_w_random():
+    """Test agentset initialization."""
+    model = Model()
+    empty_agentset = AgentSet([], random=model.random)
+    assert len(empty_agentset) == 0
+    assert empty_agentset.random == model.random
+
+    agents = [AgentTest(model) for _ in range(10)]
+    agentset = AgentSet(agents)
+    assert len(agentset) == 10
+    assert agentset.random == model.random
+
+
 def test_agentset_serialization():
     """Test pickleability of agentset."""
     model = Model()
     agents = [AgentTest(model) for _ in range(5)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     serialized = pickle.dumps(agentset)
     deserialized = pickle.loads(serialized)  # noqa: S301
@@ -156,7 +172,7 @@ def test_agent_membership():
     """Test agent membership in AgentSet."""
     model = Model()
     agents = [AgentTest(model) for _ in range(5)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     assert agents[0] in agentset
     assert AgentTest(model) not in agentset
@@ -218,7 +234,7 @@ def test_agentset_get_item():
     """Test integer based access to AgentSet."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     assert agentset[0] == agents[0]
     assert agentset[-1] == agents[-1]
@@ -232,7 +248,7 @@ def test_agentset_do_str():
     """Test AgentSet.do with str."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     with pytest.raises(AttributeError):
         agentset.do("non_existing_method")
@@ -245,7 +261,7 @@ def test_agentset_do_str():
     n = 10
     model = Model()
     agents = [AgentDoTest(model) for _ in range(n)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -255,7 +271,7 @@ def test_agentset_do_str():
     # setup
     model = Model()
     agents = [AgentDoTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -267,7 +283,7 @@ def test_agentset_do_callable():
     """Test AgentSet.do with callable."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     # Test callable with non-existent function
     with pytest.raises(AttributeError):
@@ -281,7 +297,7 @@ def test_agentset_do_callable():
     n = 10
     model = Model()
     agents = [AgentDoTest(model) for _ in range(n)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -292,7 +308,7 @@ def test_agentset_do_callable():
     # setup again for lambda function tests
     model = Model()
     agents = [AgentDoTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -310,7 +326,7 @@ def remove_function(agent):
     # setup again for actual function tests
     model = Model()
     agents = [AgentDoTest(model) for _ in range(n)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -321,7 +337,7 @@ def remove_function(agent):
     # setup again for actual function tests
     model = Model()
     agents = [AgentDoTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
     for agent in agents:
         agent.agent_set = agentset
 
@@ -386,7 +402,7 @@ def test_agentset_agg():
         agent.energy = i + 1
         agent.wealth = 10 * (i + 1)
 
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     # Test min aggregation
     min_energy = agentset.agg("energy", min)
@@ -435,7 +451,7 @@ def __init__(self, model, age=None):
 
     model = Model()
     agents = [TestAgentWithAttribute(model, age=i) for i in range(5)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     # Set a new attribute "health" and an existing attribute "age" for all agents
     agentset.set("health", 100).set("age", 50).set("status", "active")
@@ -454,7 +470,7 @@ def test_agentset_map_str():
     """Test AgentSet.map with strings."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     with pytest.raises(AttributeError):
         agentset.do("non_existing_method")
@@ -467,7 +483,7 @@ def test_agentset_map_callable():
     """Test AgentSet.map with callable."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     # Test callable with non-existent function
     with pytest.raises(AttributeError):
@@ -494,7 +510,7 @@ def test_method(self):
             self.called = True
 
     agents = [TestAgentShuffleDo(model) for _ in range(100)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     # Test shuffle_do with a string method name
     agentset.shuffle_do("test_method")
@@ -544,7 +560,7 @@ def test_agentset_get_attribute():
     """Test AgentSet.get for attributes."""
     model = Model()
     agents = [AgentTest(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     unique_ids = agentset.get("unique_id")
     assert unique_ids == [agent.unique_id for agent in agents]
@@ -558,7 +574,7 @@ def test_agentset_get_attribute():
         agent = AgentTest(model)
         agent.i = i**2
         agents.append(agent)
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     values = agentset.get(["unique_id", "i"])
 
@@ -634,7 +650,7 @@ def get_unique_identifier(self):
 
     model = Model()
     agents = [TestAgent(model) for _ in range(10)]
-    agentset = AgentSet(agents, random=model.random)
+    agentset = AgentSet(agents)
 
     groups = agentset.groupby("even")
     assert len(groups.groups[True]) == 5