Improve performance of AgentSet and iter_cell_list_contents

[Corvince]

I fell into a rabbit hole this weekend that I didn't want to fall in anymore, but here we are: I was looking at mesa performance and noticed that the new AgentSet feature took quite some burden on run times. Here are the results on my local machine, manually running 3 model runs for each (so not completely accurate, but indicative enough). I also included the results for this PR

Mesa Version	Schelling 100x100 for 100 steps
Mesa 2.1.5	~11.8 sec
Mesa 2.2.0	~18.2 sec
Mesa-PR	~8.9 sec

So I think AgentSet had some noticable effect, but through some modifications to AgentSet and also some stuff in Space.py I was able to increase overall performance.

I think the changes are non-intrusive, but here they are

Agentset

I first inherit the weakref from a normal dict instead of adding all agents sequentially
I also added a short-cut for the select function. If you call AgentSet.select() you immediately receive a copy the agentset, without triggering any logic for filter-func, etc. The scheduler does this, so it applies to all models using the scheduler.

Space

I modified iter_cell_list_contents to retrieve the default_val of the space only once. I also rewrote the generator expression to a normal generator, because perfomance is about the same, but it is easier to understand.
Additionally iter_neighbors now returns the content directly instead of the call to iter_cell_list_contents, which is annotated by accept_tuple_arguments, which costed some performance even though we knew its a list of positions. The difference might appear small, but in my model run this generator gets called 6 million times, so even a microsecond faster execution translates to a whole second.

Observation: Weakref

With these modifications the use of Weakref became the new bottleneck. Since weakKeyDictionaries are implemented in pure python and not in C, they considerable slower than Agent dictionaries. But I think it is still manageable.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (9819927) 79.45% compared to head (1bd198d) 79.39%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1964      +/-   ##
==========================================
- Coverage   79.45%   79.39%   -0.07%     
==========================================
  Files          15       15              
  Lines        1285     1286       +1     
  Branches      285      246      -39     
==========================================
  Hits         1021     1021              
- Misses        225      227       +2     
+ Partials       39       38       -1     

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

[EwoutH]

Thanks a lot for this effort!

I also added a short-cut for the select function. If you call AgentSet.select() you immediately receive a copy the agentset, without triggering any logic for filter-func, etc. The scheduler does this, so it applies to all models using the scheduler.

I honestly thought we did this already. This will probably be the largest speedup, right?

[rht]

This seems to be identical to agent_generator. And if there is a way to deduplicate the definition.

[Corvince]

whoopsie, you are correct, this shouldn't be here anymore. Removed it now

[rht]

While I can see caching default_val() as a speedup, how does this line speed up the code?

[Corvince]

This just chains the iterable from the parent function - which is already cache-optimized. So now we don't have to call iter_cell_list_contents, which is wrapped in an annotation

[rht]

Mesa-PR ~8.9 sec

I'd be curious to know if the PR can be broken down into Mesa-PR-1st-commit and Mesa-PR-2nd-commit. But only if you have the time. It seems to me that the AgentSet commit is the one that provides the most speedup, but not sure.

[quaquel]

Thanks for this. I saw this rabbit hole last week as well. I am glad you went down it and came up with some solutions. I'll try to find time for a closer look at the code later today.

[Corvince]

Mesa-PR ~8.9 sec

I'd be curious to know if the PR can be broken down into Mesa-PR-1st-commit and Mesa-PR-2nd-commit. But only if you have the time. It seems to me that the AgentSet commit is the one that provides the most speedup, but not sure.

Tried it just now, but my laptop seems to be doing some stuff in the background, the results are very flaky. So just from my memory I started out with the modifications to space.py and that imroved things by about 2 seconds. The rest was indeed from Agentset, with dict update vs dict set about 1 second differene and the rest coming from the shortcut and the updated agent_generator

Will try to run again later to get better results

[EwoutH]

I wil try to create a benchmark script that runs some of the most used models in a semi-reproducible way. Thesis has some priority, but I will try to do it later this week (see #1931 (reply in thread)).

[quaquel]

I looked through the code, and all changes seem sensible and non-intrusive.

[EwoutH]

Approving based on @quaquel's insights.

@rht could you do final review and merge if satisfied?