Introduce AgentSet class

[EwoutH]

This PR is based on the discussion in #1912 (comment) for the initial discussion.

This PR introduces the AgentSet class in the Mesa agent-based modeling framework, along with significant changes to the agent management process. The AgentSet class is designed to encapsulate and manage a collection of agents, providing methods for efficient selection, sorting, shuffling, and applying actions to groups of agents. This addition aims to enhance the framework's scalability and flexibility in handling agent operations.

Key Changes

1. Agent Class Modifications: The Agent class has been updated to manage agent registration and removal within the model's _agents attribute directly. This approach eliminates the need for separate registration and removal methods.
2. Model Class Enhancements: The Model class now utilizes the AgentSet class. The agents property returns an AgentSet instance, representing all agents in the model, thus streamlining agent access and manipulation.
3. AgentSet Functionality: The AgentSet class includes methods like select, shuffle, sort, and do for more intuitive operations on agent collections. These methods facilitate tasks like filtering agents, randomizing their order, sorting, and applying actions to each agent.
4. Efficient Agent Management: AgentSet maintains weak references to agents, enabling efficient management of agent lifecycles and garbage collection. It uses a WeakKeyDictionary to store agents.

Usage examples

1. Selecting and Filtering Agents:

   # Select agents with a specific attribute
   selected_agents = model.agents.select(lambda agent: agent.attribute > threshold)

2. Shuffling Agents:

   # Randomize agent order
   shuffled_agents = model.agents.shuffle()

3. Sorting Agents:

   # Sort agents by wealth, high to low
   sorted_agents = model.agents.sort("wealth")

4. Applying Actions to Agents:

   # Apply a method to each agent
   model.agents.do('method_name')

5. Retrieving Agent Attributes:

   # Get a specific attribute from all agents
   attributes = model.agents.get('attr_name')

Advanced Chaining examples

1. Filter, Shuffle, and Apply an Action:
   Chaining select, shuffle, and do allows for filtering, randomizing, and then performing an action on the selected agents.

   # Select agents meeting a condition, shuffle them, and then apply a method
   model.agents.select(lambda a: a.condition).shuffle().do('action_method')

2. Sort, Select Top N, and Get Attribute:
   This chain sorts agents, selects the top N agents based on the sort, and retrieves a specific attribute from these agents.

   # Sort agents, select the top N, and get an attribute
   top_agents_attributes = model.agents.sort(key=lambda a: a.rank).select(n=5).get('attribute')

3. Complex Filtering and Action Application:
   Combining multiple filters and then applying an action. This chain can be particularly useful for complex conditional operations.

   # Apply a method to agents that meet multiple conditions
   model.agents.select(lambda a: a.condition1).select(lambda a: a.condition2).do('method')

4. Combined Set Operations and Action Application:
   If set-like operations (union, intersection, difference) are implemented, they can be combined with action application.

   # Combine two AgentSets and apply a method to the combined set
   (agentset1 + agentset2).do('collaborative_method')

5. Shuffle, Sort, and Sequential Action Execution:
   This chain first randomizes the order of agents, then sorts them based on an attribute, and finally applies a method sequentially.

   # Shuffle, sort, and then apply a method in sequence
   model.agents.shuffle().sort(key=lambda a: a.priority).do('sequential_method')

6. Nested Chain Operations:
   Combining chaining with nested operations for more sophisticated scenarios.

   # Nested operation: Select agents, then within that selection, shuffle and apply an action
   model.agents.select(lambda a: a.group == 'X').shuffle().do('group_specific_method')

Replacing schedulers

You can use the new AgentSet functionality to replacer the current schedulers and create entirely new ones.

1. BaseScheduler

The BaseScheduler activates each agent once per step, in the order they were added.

class MyModel(Model):
    def step(self):
        self.agents.do('step')

2. RandomActivation

RandomActivation activates each agent once per step, in random order.

class MyModel(Model):
    def step(self):
        self.agents.shuffle().do('step')

3. SimultaneousActivation

In SimultaneousActivation, all agents are activated simultaneously. This can be simulated by first collecting all agents' decisions and then applying them.

class MyModel(Model):
    def step(self):
        self.agents.do("step")
        self.agents.do("advance")

4. StagedActivation

StagedActivation allows for multiple stages per step. For instance, agents first "move", then "eat", then "sleep".

class MyModel(Model):
    def step(self):
        self.agents.do('move')
        self.agents.do('eat')
        self.agents.do('sleep')

5. RandomActivationByType

RandomActivationByType activates each type of agent in random order.

class MyModel(Model):
    def step(self):
        # Assuming AgentType1 and AgentType2 are subclasses of Agent
        self.agents.select(agent_type=AgentType1).shuffle().do('step')
        self.agents.select(agent_type=AgentType2).shuffle().do('step')

These examples demonstrate how the AgentSet class can be used to replicate the behavior of different schedulers in Mesa. The AgentSet class provides a more flexible and concise way to define the activation and interaction patterns among agents within the model.

Future work

• Adding more Sequence and MutableSet methods (e.g., __reversed__, index, count, clear, pop, __ior__) for enhanced performance and functionality.
• Rewriting current schedulers using AgentSets.

[codecov]

Codecov Report

Attention: 1 lines in your changes are missing coverage. Please review.

Comparison is base (9495a5a) 77.53% compared to head (22b99cd) 79.10%.

Files	Patch %	Lines
mesa/model.py	91.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1916      +/-   ##
==========================================
+ Coverage   77.53%   79.10%   +1.56%     
==========================================
  Files          15       15              
  Lines        1015     1096      +81     
  Branches      221      236      +15     
==========================================
+ Hits          787      867      +80     
- Misses        197      198       +1     
  Partials       31       31              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[quaquel]

some quick suggestions

[EwoutH]

I won’t be working on it between now and out meeting tomorrow, so take it as far as you like. Maybe there is also still some stuff from #1911 that can be ported over.

[quaquel]

I'll see what I can do tomorrow morning. I just addressed my comments on your first submission and added the inplace keyword argument.

[Corvince]

Awesome PR

[EwoutH]

Thanks for working on It. I will try to do an initial review tomorrow between 9 and 10 CET, otherwise around noon (feel free to continue regardless).

[EwoutH]

I’m leaning towards just documenting properly that you can use agentset.copy() or copy(agentset) instead of the whole inplace thing everywhere. It just seems to add unnecessary complexity to every method while copy is a general solution.

[quaquel]

I disagree with the copy thing. First, it adds another method call into any chain. Second, the API now follows Pandas, so users of Mesa are likely to be familiar with it. Third, in particular in the case of a large number of agents, there are substantial memory advantages to using inplace=True.. Fourth, the internal readability of code is relevant but not at the expense of the usability and cleanliness of the API.

[EwoutH]

Thanks for your perspective. Sound arguments, let me further think about it. Also curious what others think.

I would propose all the "getting data back" things (like get_each) to be done in a future PR. That could be done together with the redesign of the data collector (see #348).

[Corvince]

Thanks for your perspective. Sound arguments, let me further think about it. Also curious what others think.

I would propose all the "getting data back" things (like get_each) to be done in a future PR. That could be done together with the redesign of the data collector (see #348).

I concur with keeping inplace in the API, since I think copying should be the default. We will see how that turns out performance wise, but we shouldnt worry about that just yet.

I also agree we should keep this PR small at first and add further functions later. This is a very solid basis and can/should be used for lots of things so not overcomplicate things now.

[quaquel]

I also agree we should keep this PR small at first and add further functions later. This is a very solid basis and can/should be used for lots of things so not overcomplicate things now.

Keeping the PR focused is the right way to go. However, it should also be feature-complete, and the basic structure should be solid. Otherwise, the code base keeps changing repeatedly to reflect new ideas. My suggestion would be to focus on making sure AgentSet has the following things:

• Complete behavior of an ordered set
• inclusion of shuffle, select, and sort
• ensure that you can call methods and retrieve attributes on the agents in the set. So, for me, do_each and get_each are within scope, although I am unsure about their naming. It also makes sense to ensure do_each has a return. This makes do_each much more useful across the board.

What can be discussed or kept for later, in my view, are:

• support for listlike indexing. AgentSet is ordered, so, conceptually, it makes sense to be able to retrieve by index. It also would mean there is no need for a separate AgentList class (contra AgentPy). It is possible to add this by implementing __getitem__. There are, however, some tricky questions because of the use of weak references. Should the index just refer to whatever is the fifth item in the ordered set, or should it return None or raise an exception if the agent has been removed by the garbage collector?
• the sophistication of select. The present implementation is bare bones because it only accepts a boolean callable. So selecting the first n items is not supported through select at the moment. It can however be done through slicing agents[0:5], if indexing is supported
• Adding methods like to_dataframe or as_dataframe.

[EwoutH]

Inheriting from MutableSet is an interesting approach. Since we're using it as a mixin, do we comply with these notes?

Also remarkably, there is a reference to an OrderedSet recipe in the official Python docs.

[EwoutH]

If either of you has anything to add to the Python discussion on OrderedSets, please do so. I'm trying to make the case but your probably have more specialized knowledge about this: https://discuss.python.org/t/add-orderedset-to-stdlib/12730/92

[quaquel]

We follow 1, but I am not sure we need to worry about 2 (it involves subset and superset comparisons as far as I can tell); 3 is likewise unnecessary.

While googling this, I came across some discussions on pickling the set. Is this something we want to support?

[Corvince]

I also agree we should keep this PR small at first and add further functions later. This is a very solid basis and can/should be used for lots of things so not overcomplicate things now.

Keeping the PR focused is the right way to go. However, it should also be feature-complete, and the basic structure should be solid. Otherwise, the code base keeps changing repeatedly to reflect new ideas. My suggestion would be to focus on making sure AgentSet has the following things:

• Complete behavior of an ordered set
• inclusion of shuffle, select, and sort
• ensure that you can call methods and retrieve attributes on the agents in the set. So, for me, do_each and get_each are within scope, although I am unsure about their naming. It also makes sense to ensure do_each has a return. This makes do_each much more useful across the board.

The question is should do_each return the return values the called function or return the AgentSet itself. The latter would allow to keep on chaining, for example:

agents.do_each(step).select(n=10).do_each(move)

But getting the return values could also be valueable, so maybe we can also add agents.apply(func), which returns whatever the function returns. Just thinking out loud here.

[quaquel]

Unit tests have been added, and it seems pretty complete.

However, there are several outstanding issues:

1. select is very barebones because it only takes a callable -> bool. Do we want to support other arguments, such as n=5 for the first five items?
2. sort likewise depends on a callable, so you cannot do a simple sort on attributes. In my view, this is such a common use case that we should support it. It is straightforward to do. if key is a string, use getattr with the key.
3. The current class is not pickle-able. This can be an issue if you want to run in parallel and have created AgentSets somewhere during the initialization of the model.
4. Naming for do_each invokes the given method on each agent in the set, and get_each retrieves the given attribute from each agent in the set.

[Corvince]

Just two more comments from my side, otherwise I think we are good to go!

[quaquel]

Just two more comments from my side, otherwise I think we are good to go!

both have been done.

[Corvince]

I approved, but would like to have 2 more eyes on this before I feel comfortable merging. But I think @EwoutH you can start to adapt this in your other prs

[EwoutH]

Thanks, I agree on at least another reviewer. No other PRs I have currently open have to be modified based on this PR (we can modify the schedulers, datacollector, and maybe grid in the future).

[tpike3]

@EwoutH This is great!!!

Thanks @Corvince and @quaquel this is truly an amazing effort and incredible discussion.

[ankitk50]

@EwoutH , I get this error while running one of the examples which looks related to this change:

Traceback (most recent call last):
  File "/Users/ankitkumar/mesa/venvMesa/lib/python3.11/site-packages/reacton/core.py", line 1647, in _render
    root_element = el.component.f(*el.args, **el.kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ankitkumar/mesa/mesa/mesa/experimental/jupyter_viz.py", line 57, in JupyterViz
    reset_counter = solara.use_reactive(0)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ankitkumar/mesa/venvMesa/lib/python3.11/site-packages/reacton/core.py", line 828, in use_memo
    return rc.use_memo(f, dependencies, debug_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ankitkumar/mesa/venvMesa/lib/python3.11/site-packages/reacton/core.py", line 1193, in use_memo
    value = f()
            ^^^
  File "/Users/ankitkumar/mesa/mesa/mesa/experimental/jupyter_viz.py", line 52, in make_model
    def make_model():
                ^^^^^^
  File "/Users/ankitkumar/mesa/mesa-examples/examples/boltzmann_wealth_model_experimental/model.py", line 29, in __init__
    a = MoneyAgent(i, self)
        ^^^^^^^^^^^^^^^^^^^
  File "/Users/ankitkumar/mesa/mesa-examples/examples/boltzmann_wealth_model_experimental/model.py", line 53, in __init__
    super().__init__(unique_id, model)
  File "/Users/ankitkumar/mesa/mesa/mesa/agent.py", line 45, in __init__
    self.model.agents[type(self)][self] = None
    ^^^^^^^^^^^^^^^^^
AttributeError: 'BoltzmannWealthModel' object has no attribute 'agents'

[EwoutH]

Thanks for reporting, that's because the example (for some reason) doesn't run super().__init__() in the BoltzmannWealthModel init, and thus the Mesa Model initialization is never ran, and thus model.agents never initialized.

It should be fixed in #1928.

[quaquel]

I used the BoltzmannWealthModel when testing #1928 so I know that one works fine with #1928.

[rht]

(Still catching up with the influx of exciting constructive developments with the new APIs.)

Modifying slightly the example:

# Get a specific attribute from all agents
attributes = model.agents.select(...).get('attr_name')
# Or
agent_attrs_df = model.agents.select(...).get('attr_name1', 'attr_name2').to_df()
# Or
agent_attrs_df = to_df(model.agents.select(...).get('attr_name1', 'attr_name2'))

As far as I understand, this overlaps with the declarative data collecting specification in #348 (comment)? I find it excessive to have 2 APIs, where one is for agent dynamics description and another for data collection. I am wondering if the 2 APIs are inevitable. The constraint seems to be to make it easy to build complex scheduler and data collection, while to minimize the amount of documentation needed.

[rht]

How do I select model.agents by agent type in a way that is more performant than model.agents.select(lambda agent: isinstance(agent, AgentType))?

[quaquel]

1. I agree that any revision of data collection should built on AgentSet to avoid as much as possible having 2 different API's. I have been thinking about this in light also of Multi-Agent data collection #348 and Add system state tracking #1933. In short, I believe all data collection can be simplified to retrieval of one or more attributes from an object (or collection of objects) and optionally applying a function to it. This is consistent with @Corvince proposal in Multi-Agent data collection #348. So, target in this proposal should also accept an AgentSet (and SystemState from #1933 ). In addition, any new approach to data collection should avoid retrieving the same data multiple times. For example, in the Boltzman wealth model, currently wealthis retrieved once as agent reporter and once for theGini` model reporter.

Currently, AgentSet.get does not accept multiple attributes (so (model.agents.select(...).get('attr_name1', 'attr_name2') does not work). I believe this is something that would be good to add because it avoids having to iterate repeatedly over the AgentSet for the different attributes you want to retrieve.

3. There is a seperate method on the mode: get_agents_of_type. model stores the agents internally (i.e., _agents) in a dict with type as key and AgentSet as value.

[EwoutH]

Just thought of this now, but their might be one edge-case in which we might break user models: If model.agents is already defined in their models. Some model might use agents to track a number of agents or some set of them.