AgentSet: Allow selecting a fraction of agents in the AgentSet

mesa/agent.py

@@ -121,39 +121,59 @@ def __contains__(self, agent: Agent) -> bool:
     def select(
         self,
         filter_func: Callable[[Agent], bool] | None = None,
-        n: int = 0,
+        at_most: int | float = float("inf"),
         inplace: bool = False,
         agent_type: type[Agent] | None = None,
+        n: int | None = None,
     ) -> AgentSet:
         """
         Select a subset of agents from the AgentSet based on a filter function and/or quantity limit.
 
         Args:
             filter_func (Callable[[Agent], bool], optional): A function that takes an Agent and returns True if the
                 agent should be included in the result. Defaults to None, meaning no filtering is applied.
-            n (int, optional): The number of agents to select. If 0, all matching agents are selected. Defaults to 0.
+            at_most (int | float, optional): The maximum amount of agents to select. Defaults to infinity.
+              - If an integer of 1 or larger, the first n matching agents are selected.
+              - If a float between 0 and 1, at most that fraction of original the agents are selected.
+            fraction (float, optional): The fraction of agents to select. If 0.2, it selects 20% of Agents in the AgentSet. Defaults to 1.0.
             inplace (bool, optional): If True, modifies the current AgentSet; otherwise, returns a new AgentSet. Defaults to False.
             agent_type (type[Agent], optional): The class type of the agents to select. Defaults to None, meaning no type filtering is applied.
 
         Returns:
             AgentSet: A new AgentSet containing the selected agents, unless inplace is True, in which case the current AgentSet is updated.
+
+        Notes:
+            - at_most just return the first n or fraction of agents. To take a random sample, shuffle() beforehand.
+            - at_most is an upper limit. When specifying other criteria it can be smaller.
         """
+        if n is not None:
+            warnings.warn(
+                "The parameter 'n' is deprecated. Use 'at_most' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            at_most = n
 
-        if filter_func is None and agent_type is None and n == 0:
+        inf = float("inf")
+        if filter_func is None and agent_type is None and at_most == inf:
             return self if inplace else copy.copy(self)
 
-        def agent_generator(filter_func=None, agent_type=None, n=0):
+        # Check if at_most is of type float
+        if at_most <= 1.0 and isinstance(at_most, float):
+            at_most = int(len(self) * at_most)  # Note that it rounds down (floor)
+
+        def agent_generator(filter_func, agent_type, at_most):
             count = 0
             for agent in self:
+                if count >= at_most:
+                    break
                 if (not filter_func or filter_func(agent)) and (
                     not agent_type or isinstance(agent, agent_type)
                 ):
                     yield agent
                     count += 1
-                    if 0 < n <= count:
-                        break
 
-        agents = agent_generator(filter_func, agent_type, n)
+        agents = agent_generator(filter_func, agent_type, at_most)
 
         return AgentSet(agents, self.model) if not inplace else self._update(agents)
 

tests/test_agent.py

@@ -59,6 +59,12 @@ def test_agentset():
     def test_function(agent):
         return agent.unique_id > 5
 
+    assert len(agentset.select(at_most=0.2)) == 2  # Select 20% of agents
+    assert len(agentset.select(at_most=0.549)) == 5  # Select 50% of agents
+    assert len(agentset.select(at_most=0.09)) == 0  # Select 0% of agents
+    assert len(agentset.select(at_most=1.0)) == 10  # Select 100% agents
+    assert len(agentset.select(at_most=1)) == 1  # Select 1 agent
+
     assert len(agentset.select(test_function)) == 5
     assert len(agentset.select(test_function, n=2)) == 2
     assert len(agentset.select(test_function, inplace=True)) == 5