mesa-frames: faster mesa with data-frame based operations

[adamamer20]

Hello,
I am excited to introduce a new extension for the Mesa simulation library, named mesa-frames. This extension is designed to enhance the performance of models with thousands of agents by leveraging pandas dataframes for storing agent attributes and utilizing vectorized operations. In cases where vectorized operations are not feasible, operations can be performed by row by using .apply. This approach has demonstrated greater efficiency than parallel processing the steps of the agents in my project.
The goal of mesa-frames is to facilitate an easy transition from the base Mesa library, ensuring compatibility with existing Mesa operations while introducing the capability to 'vectorize' a model using a decorator. Currently, the extension is in its initial proof-of-concept stage.
I warmly welcome feedback, suggestions, and contributions to this project.
Thank you.

[Corvince]

Wow that sounds super interesting and the performance chart looks very promising.

I am curious what kind of problems can be modeled this way and how we can facilitate the interoperability with mesa.

Congratulations on the launch

[EwoutH]

Looks awesome, really excited to see where it’s going!

If you need any (specific) input, feel free to tag me.

[rht]

Ran the benchmark for small number of agents.

The upfront baseline runtime of 0.25 s is worth the tradeoff of being slower than the current AgentSet implementation. I think we can mold the API of mesa-frames to have the methods AgentSet (select, shuffle, sort, do). Or alternatively, replace the data structure backend of AgentSet from weakref.WeakKeyDictionary() to a pandas DF.

Edit: clarification

[rht]

Using pola.rs

with the following DF code

class MoneyModelDF(ModelDF):
    def __init__(self, N):
        super().__init__()
        self.num_agents = N
        self.agents = pl.DataFrame({"wealth": [1 for i in range(N)]})

    def step(self, merged_mro=True):
        wealthy_agents = self.agents.filter(pl.col("wealth") > 0)
        n_wealthy = len(wealthy_agents)
        if n_wealthy > 0:
            self.agents = self.agents.with_columns(wealth=pl.when(pl.col("wealth") > 0).then(pl.col("wealth") - 1).otherwise("wealth"))

            rows_to_update = np.random.choice(self.agents.height, size=n_wealthy, replace=False)
            # Could be wrong, but this is not going to be much slower than the previous with_columns operation anyway
            self.agents = self.agents.with_columns(wealth=pl.when(pl.col("wealth").is_in(rows_to_update)).then(pl.col("wealth") + 1).otherwise("wealth"))

[rht]

With Polars, we can get a huge speedup (~10x over pandas) not just for the AgentSet, but also on agent step(). However, Polars native expression API is hard to translate to how users would usually define their agent step(). There are definitely lots of use case of ABM that is not tied to geopandas (e.g. economic models). As such, I think it is worth implementing a variant of mesa-frames using Polars.

[adamamer20]

The speedup is amazing! I initially started with subclassing Dataframes but now i have switched to using a custom accessor to avoid the complexity of constructors. Since also polars has a similar api for extensions, i think it wouldn't be much more time consuming doing both and it is definetely worth it.

[rht]

I didn't believe it at first. Too good to be true. But the earlier versions of my code resembled idiomatic pandas, and was slightly slower than the pandas version. There were warnings saying that it wasn't using Polars native expression API, and I fixed them one by one until I got to the version I posted above.

[rht]

Context: #1561, cc: @mrgoodfish check out the Polars vs pandas benchmark result above.

[adamamer20]

That's true, in the case when it is unavoidable to do RandomActivation. I suppose having a do method that applies on a shuffled rows sequentially, would be convenient, for feature completeness sake, albeit slow. I think you could try PyArrow as the backend to pandas instead of the NumPy backend, which may be even faster.

Sure, I will try to test both.

Wondering if any kind of queuing or pipelining can help speed this up. You can shuffle once, and then the order and the operations are known, only the exact input data isn’t (since previous agent actions may have changed states).

Edit: with the exception of agents being removed by other agents.

if the operation relies on the states of x previous agents one could use pandas rolling method. In this case, the implementation should still be left to the user. Otherwise, if the operation relies on external states or, as you mentioned, agents are removed, I do not see any other way of doing this then iterating sequentially through the rows.

I will try to implement AgentSet using dataframes soon and let you know.

[rht]

See also https://pola.rs/ for performance.

[rht]

@adamamer20 is it correct that in order to test for the PyArrow backend, in https://github.com/adamamer20/mesa-frames/blob/9bee978dc29a9751d9c7fcc3b6c2ca614a4b2da6/mesa_frames/model.py#L262, I added "[pyarrow]"?

[adamamer20]

@adamamer20 is it correct that in order to test for the PyArrow backend, in https://github.com/adamamer20/mesa-frames/blob/9bee978dc29a9751d9c7fcc3b6c2ca614a4b2da6/mesa_frames/model.py#L262, I added "[pyarrow]"?

yes, that's correct. i believe that in that case there's no need to convert dtype=int to dtype=float because I remember reading somewhere that PyArrow supports missing values with int.

As for polars, it is definitely worth to try because it allows also to use a dataset larger than the RAM (useful if you have millions of agents). The only two potential issues was a slightly different interface than pandas and the lack of pandas extensions (for example, it's missing the equivalent of geopandas which is necessary for mesa-geo).

[jackiekazil]

This is exciting stuff!

[adamamer20]

@rht @Corvince @EwoutH @jackiekazil let me know what you think of the implementation so far in adamamer20/mesa_frames#8. If you like the API, I will implement the backend with polars.

[rht]

@adamamer20 check out #1997 and https://github.com/projectmesa/mesa/wiki/Projects-for-Google-Summer-of-Code-2024.

[adamamer20]

@tpike3 @rht thank you very much for the project proposal!

[rht]

We put in for the GSOC 2024 if we get accepted then you apply to be a contributor to it.

@adamamer20 Mesa has been accepted: https://summerofcode.withgoogle.com/programs/2024/organizations/project-mesa. The application opens on March 18.

[adamamer20]

@rht wonderful news, thank you. I will apply as soon as it opens.

[rht]

@adamamer20 we haven't heard back from you. Let's continue the discussion on https://matrix.to/#/#project-mesa:matrix.org.

[EwoutH]

@adamamer20 Are you still working on this? I might want to run some simulations with around ~100.000 agents so I might become very interested - and invested - in this.

[adamamer20]

@EwoutH Sure! I was waiting on the results and feedback from the GSoC proposal before going back into development. If you want to provide some guidance or contribute, it would be more than welcome. What are your priorities?

[EwoutH]

Great to hear! I'm going to dive deeper in it in the coming week. I will be simulating a city of a hundred thousand to a million inhabitants, but with limited interactivity and relative simple behaviour.

I would like mesa-frames to be a clean subset of mesa. Maybe we can integrate it in the mesa organization, like we have done with Mesa-geo. I will take this up with the maintainers.

Could you check out #2066? One of the things I would like to be able to do, is efficiently simulate a day, in which agents do only one or a few times something, while still having a relatively high temporal resolution (minutes). Discrete event scheduling would be perfect for that. But that might be relatively difficult with frames (or not, it's a gut feeling).

As of GSoC, we would have to look unbiased of course, but we certainly consider it. In any case I hope we can keep developing mesa-frames.

[adamamer20]

That's the direction I intended for mesa-frames too!
I think supporting event scheduling shouldn't be too difficult. I have to dive deeper into the code of #2066 but adding a column for events and masking active_agents on that column should work. I will try to add something like it in the coming weeks.
Regardless of the GSoC outcome, I'm commited to mesa-frames but the majority of the development will still take place during summer. With courses and exams, the time I can dedicate to it is limited.

[adamamer20]

@jackiekazil @tpike3 @rht @wang-boyu @quaquel @Corvince @EwoutH
Thank you all for accepting my proposal!
In the next weeks I'll develop a more concrete plan but I plan on closing adamamer20/mesa-frames#8 soon and work on adamamer20/mesa-frames#9. I'll also familiarize myself with the latest 3.0.0 developments. I'll use this as the main discussion thread.

[rht]

Congratulations, @adamamer20. We will also have to think of a stable API so that maintaining mesa-frames is more manageable. And visualization integration.

[adamamer20]

For the API, I was experimenting with dunder methods to align to pandas/polars syntax and make it more user-friendly. I think they could also be added to mesa.AgentSet.

[rht]

Regarding with the naming, I'm wondering if it is easier to remember to name it mesa-frame instead of mesa-frames? I don't see a rationale for making it plural.

[adamamer20]

I'm not opposed to that. I was also dubious about the hyphen because python beginners may find it confusing (you need to install mesa-frame but import mesa_frame)

[jackiekazil]

I thought it was frames, because it involved dataframes (plural). I don't care enough about "s" or no "s". So I defer to you all.

RE: hyphen versus underscore - this is long been a thing in python of inconsistency. You can see this happening in our community as well -- names with hyphens (mesa-geo), names which are two words together (cropyield).

Some notes on this --

• package name uses hyphen if needed, because SEO, I believe - https://www.stanventures.com/blog/underscore-in-url/#:~:text=Google%20recommends%20using%20hyphens%20rather,find%20in%20Google's%20official%20documentation.
• underscores, while possible are discouraged according to pep8 - https://peps.python.org/pep-0008/#package-and-module-names

These are the things that end up being something that we (engineers) can never agree on, but spend a lot of time talking about.
Now that I have gotten to the end of my comment, I am not sure any of it was really helpful.