Skip to content

Commit e42b570

Browse files
authored
speed up mesa readme_plot script (#26)
* speed up mesa, add perfplot to dev dependencies * changed to mesa's deafult random generator
1 parent 23fe6f1 commit e42b570

File tree

6 files changed

+69
-34
lines changed

6 files changed

+69
-34
lines changed

README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ DataFrames are optimized for simultaneous operations through [SIMD processing](h
1111

1212
The following is a performance graph showing execution time using mesa and mesa-frames for the [Boltzmann Wealth model](https://mesa.readthedocs.io/en/stable/tutorials/intro_tutorial.html).
1313

14-
![Performance Graph](https://github.com/adamamer20/mesa_frames/blob/main/docs/images/readme_plot.png)
14+
![Performance Graph](https://github.com/adamamer20/mesa_frames/blob/main/docs/images/readme_plot_0.png)
1515

1616
![Performance Graph without Mesa](https://github.com/adamamer20/mesa_frames/blob/main/docs/images/readme_plot_1.png)
1717

docs/images/readme_plot.png

-32.3 KB
Binary file not shown.

docs/images/readme_plot_0.png

51.8 KB
Loading

docs/images/readme_plot_1.png

6.54 KB
Loading

docs/scripts/readme_plot.py

+65-32
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def __init__(self, unique_id, model):
2828
def step(self):
2929
# Verify agent has some wealth
3030
if self.wealth > 0:
31-
other_agent = self.random.choice(self.model.schedule.agents)
31+
other_agent = self.random.choice(self.model.agents)
3232
if other_agent is not None:
3333
other_agent.wealth += 1
3434
self.wealth -= 1
@@ -41,19 +41,13 @@ def __init__(self, N):
4141
super().__init__()
4242
self.num_agents = N
4343
# Create scheduler and assign it to the model
44-
self.schedule = mesa.time.RandomActivation(self)
45-
46-
# Create agents
47-
for i in range(self.num_agents):
48-
a = MoneyAgent(i, self)
49-
# Add the agent to the scheduler
50-
self.schedule.add(a)
44+
self.agents = [MoneyAgent(i, self) for i in range(self.num_agents)]
5145

5246
def step(self):
5347
"""Advance the model by one step."""
54-
55-
# The model's step will go here for now this will call the step method of each agent and print the agent's unique_id
56-
self.schedule.step()
48+
self.random.shuffle(self.agents)
49+
for agent in self.agents:
50+
agent.step()
5751

5852
def run_model(self, n_steps) -> None:
5953
for _ in range(n_steps):
@@ -176,7 +170,9 @@ def __init__(self, n: int, model: ModelDF) -> None:
176170
# 2. Adding the dataframe with add
177171
# self.add(pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)}))
178172
# 3. Adding the dataframe with __iadd__
179-
self += pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)})
173+
self += pd.DataFrame(
174+
{"unique_id": np.arange(n, dtype="int64"), "wealth": np.ones(n)}
175+
)
180176

181177
def step(self) -> None:
182178
# The give_money method is called
@@ -212,7 +208,9 @@ class MoneyAgentPandasNative(AgentSetPandas):
212208
def __init__(self, n: int, model: ModelDF) -> None:
213209
super().__init__(model)
214210
## Adding the agents to the agent set
215-
self += pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)})
211+
self += pd.DataFrame(
212+
{"unique_id": np.arange(n, dtype="int64"), "wealth": np.ones(n)}
213+
)
216214

217215
def step(self) -> None:
218216
# The give_money method is called
@@ -274,33 +272,68 @@ def mesa_frames_pandas_native(n_agents: int) -> None:
274272
model.run_model(100)
275273

276274

275+
def plot_and_print_benchmark(labels, kernels, n_range, title, image_path):
276+
out = perfplot.bench(
277+
setup=lambda n: n,
278+
kernels=kernels,
279+
labels=labels,
280+
n_range=n_range,
281+
xlabel="Number of agents",
282+
equality_check=None,
283+
title=title,
284+
)
285+
286+
plt.ylabel("Execution time (s)")
287+
out.save(image_path)
288+
289+
print("\nExecution times:")
290+
for i, label in enumerate(labels):
291+
print(f"---------------\n{label}:")
292+
for n, t in zip(out.n_range, out.timings_s[i]):
293+
print(f" Number of agents: {n}, Time: {t:.2f} seconds")
294+
print("---------------")
295+
296+
277297
def main():
278298
sns.set_theme(style="whitegrid")
279299

280-
labels = [
281-
# "mesa",
300+
labels_0 = [
301+
"mesa",
282302
"mesa-frames (pl concise)",
283303
"mesa-frames (pl native)",
284304
"mesa-frames (pd concise)",
285305
"mesa-frames (pd native)",
286306
]
287-
out = perfplot.bench(
288-
setup=lambda n: n,
289-
kernels=[
290-
# mesa_implementation,
291-
mesa_frames_polars_concise,
292-
mesa_frames_polars_native,
293-
mesa_frames_pandas_concise,
294-
mesa_frames_pandas_native,
295-
],
296-
labels=labels,
297-
n_range=[k for k in range(100, 10000, 1000)],
298-
xlabel="Number of agents",
299-
equality_check=None,
300-
title="100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels),
301-
)
302-
plt.ylabel("Execution time (s)")
303-
out.save("docs/images/readme_plot_2.png")
307+
kernels_0 = [
308+
mesa_implementation,
309+
mesa_frames_polars_concise,
310+
mesa_frames_polars_native,
311+
mesa_frames_pandas_concise,
312+
mesa_frames_pandas_native,
313+
]
314+
n_range_0 = [k for k in range(0, 100001, 10000)]
315+
title_0 = "100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels_0)
316+
image_path_0 = "docs/images/readme_plot_0.png"
317+
318+
plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0)
319+
320+
labels_1 = [
321+
"mesa-frames (pl concise)",
322+
"mesa-frames (pl native)",
323+
"mesa-frames (pd concise)",
324+
"mesa-frames (pd native)",
325+
]
326+
kernels_1 = [
327+
mesa_frames_polars_concise,
328+
mesa_frames_polars_native,
329+
mesa_frames_pandas_concise,
330+
mesa_frames_pandas_native,
331+
]
332+
n_range_1 = [k for k in range(100000, 1000001, 100000)]
333+
title_1 = "100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels_1)
334+
image_path_1 = "docs/images/readme_plot_1.png"
335+
336+
plot_and_print_benchmark(labels_1, kernels_1, n_range_1, title_1, image_path_1)
304337

305338

306339
if __name__ == "__main__":

pyproject.toml

+3-1
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,12 @@ pandas = [
2121
"pyarrow",
2222
]
2323
polars = [
24-
"polars>=1.0.0", #polars._typing (see mesa_frames.types) added in 1.0.0
24+
"polars>=1.0.0", #polars._typing (see mesa_frames.types_) added in 1.0.0
2525
]
2626
dev = [
2727
"mesa_frames[pandas,polars]",
28+
"perfplot", #readme_script
29+
"seaborn", #readme_script
2830
"pytest",
2931
"pytest-cov",
3032
"typeguard",

0 commit comments

Comments
 (0)