Allow selecting a fraction of agents in the AgentSet (#2253)

This PR updates the `select` method in the `AgentSet` class by replacing the `n` parameter with a more versatile `at_most` parameter. The `at_most` parameter allows for selecting either a specific number of agents or a fraction of the total agents when provided as an integer or a float, respectively. Additionally, backward compatibility is maintained by supporting the deprecated `n` parameter, which will trigger a warning when used.

### Motive
Previously, the `select` method only allowed users to specify a fixed number of agents (`n`) to be selected. The new `at_most` parameter extends this functionality by enabling the selection of agents based on a proportion of the total set, which is particularly useful in scenarios where relative selection is desired over absolute selection.

### Implementation
- **`at_most` Parameter:** 
  - Accepts either an integer (to select a fixed number of agents) or a float between 0.0 and 1.0 (to select a fraction of the total agents).
  - `at_most=1` selects one agent, while `at_most=1.0` selects all agents.
  - If a float is provided, it determines the maximum fraction of agents to be selected from the total set. It rounds down to the nearest number of whole agents.
- **Backward Compatibility:**
  - The deprecated `n` parameter is still supported, but it now serves as a fallback for `at_most` and triggers a deprecation warning.
- **Behavior Notes:**
  - `at_most` serves as an upper limit on the number of selected agents. If additional filtering criteria are provided, the final selection may include fewer agents.
  - For random sampling, users should shuffle the `AgentSet` before applying `at_most`.

Author: EwoutH

mesa/agent.py

@@ -114,39 +114,58 @@ 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, at most the first number of matching agents are selected.
+              - If a float between 0 and 1, at most that fraction of original the agents are selected.
             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, the number of agents returned 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