Skip to content

Commit

Permalink
Allow selecting a fraction of agents in the AgentSet (projectmesa#2253)
Browse files Browse the repository at this point in the history
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`.
  • Loading branch information
EwoutH committed Sep 24, 2024
1 parent d721576 commit 7354c4a
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 7 deletions.
33 changes: 26 additions & 7 deletions mesa/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,39 +129,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)

Expand Down
6 changes: 6 additions & 0 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 7354c4a

Please sign in to comment.