Add optimized shuffle_do() method to AgentSet

[EwoutH]

This PR introduces a new shuffle_do method to the AgentSet class, optimizing the process of shuffling agents and applying a method to them in a single operation.

Motive

The current approach of shuffling agents and then applying a method (shuffle().do()) is inefficient, especially for large agent sets. This new method aims to significantly improve performance by combining these operations, reducing memory allocations and iterations.

Implementation

• Added a new shuffle_do method to the AgentSet class in mesa/agent.py.
• The method takes a method parameter (either a string or callable) and additional args/kwargs.
• It shuffles the agents in-place.
• The specified method is then applied to each agent in the shuffled order.
• Added corresponding unit tests in tests/test_agent.py.
• Updated the benchmarks (BoltzmannWealth and Wolf-sheep).

Usage Examples

# Using a string method name
model.agents.shuffle_do("step")

# Using a lambda function
model.agents.shuffle_do(lambda agent: agent.move())

Before/After Performance Comparison:

Configuration	shuffle().do()	shuffle(inplace=True).do()	shuffle_do()
10,000 agents, 1,000 steps	8.27 s	3.65 s	3.06 s
100,000 agents, 100 steps	13.71 s	6.31 s	3.71 s
1,000,000 agents, 10 steps	18.36 s	9.44 s	5.75 s

As shown, shuffle_do() provides significant performance improvements, especially for larger agent sets.

Additional Notes

• This change is backwards compatible and doesn't affect existing code using shuffle().do().
• The performance gain is more pronounced for larger agent sets and more frequent shuffling operations.
• Examples and docs can be updated (other PR).

Edit: Example models are updated in projectmesa/mesa-examples#201.

[github-actions]

Performance benchmarks:

Model	Size	Init time [95% CI]	Run time [95% CI]
BoltzmannWealth	small	🔵 -2.1% [-2.6%, -1.6%]	🔵 -0.5% [-0.6%, -0.4%]
BoltzmannWealth	large	🔵 -9.1% [-23.6%, -1.3%]	🔵 -2.7% [-7.5%, +3.2%]
Schelling	small	🔵 -2.0% [-2.3%, -1.6%]	🔵 -3.2% [-3.8%, -2.8%]
Schelling	large	🔵 -2.2% [-3.1%, -1.1%]	🟢 -9.1% [-12.0%, -6.2%]
WolfSheep	small	🔵 -0.7% [-2.0%, +0.6%]	🔵 -1.3% [-1.7%, -0.8%]
WolfSheep	large	🔵 -1.9% [-2.8%, -1.0%]	🔵 -4.3% [-6.9%, -1.7%]
BoidFlockers	small	🔵 +2.2% [+1.8%, +2.6%]	🔵 +0.2% [-0.4%, +0.8%]
BoidFlockers	large	🔵 +2.6% [+2.1%, +3.2%]	🔵 -0.4% [-1.0%, +0.3%]

[EwoutH]

Small benchmark:

from mesa import Agent, Model

class MyAgent(Agent):
    def __init__(self, model):
        super().__init__(model)

    def step(self):
        pass

class MyModel1(Model):
    def __init__(self):
        super().__init__()
        # Create 10000
        for _ in range(10000):
            MyAgent(self)

    def step(self):
        self.agents.shuffle().do("step")


class MyModel2(Model):
    def __init__(self):
        super().__init__()
        # Create 10000
        for _ in range(10000):
            MyAgent(self)

    def step(self):
        self.agents.shuffle(inplace=True).do("step")


class MyModel3(Model):
    def __init__(self):
        super().__init__()
        # Create 10000
        for _ in range(10000):
            MyAgent(self)

    def step(self):
        self.agents.shuffle_do("step")


# Benchmark the three models
from time import time

models = [MyModel1(), MyModel2(), MyModel3()]
times = []
print("Running models...")
for model in models:
    print(f"Running {model}")
    start = time()
    for _ in range(1000):
        model.step()
    times.append(time() - start)

print(f"Results: {times}")

Configuration	shuffle().do()	shuffle(inplace=True).do()	shuffle_do()
10,000 agents, 1,000 steps	8.27 s	3.65 s	3.06 s
100,000 agents, 100 steps	13.71 s	6.31 s	3.71 s
1,000,000 agents, 10 steps	18.36 s	9.44 s	5.75 s

Based on this benchmark, such a function might be useful. I would propose do_random().

[EwoutH]

A bit more radical approach: #2284

[EwoutH]

I was thinking of adding a shuffle keyword argument to some functions, starting with do(). So you can do:

model.agents.do("step", shuffle=True)

I think this is explicit and allows us to also apply this to other performance critical functions.

What does everybody think about this as an API? It should have the same performance as shuffle_do.

[quaquel]

I am indifferent between the two. So pick which one you prefer and get this done. The performance benefits are clear and shuffling followed by do is probably the most common sequence of operations done with an agentset. This last point might be a reason to have it as a separate method rather than via a keyword argument.

[quaquel]

the pr text itself is not correct:

It shuffles the agents in-place using the model's random number generator.

It uses self.random, which can be any rng. So just remove the last part from using onward.

Otherwise fine.

[github-actions]

Performance benchmarks:

Model	Size	Init time [95% CI]	Run time [95% CI]
BoltzmannWealth	small	🔵 -0.3% [-1.5%, +0.9%]	🟢 -9.8% [-10.4%, -9.2%]
BoltzmannWealth	large	🔵 +0.9% [-1.0%, +2.7%]	🟢 -8.8% [-9.3%, -8.4%]
Schelling	small	🔵 -1.2% [-1.9%, -0.4%]	🔵 -1.4% [-1.9%, -0.9%]
Schelling	large	🔵 -0.8% [-2.1%, +0.4%]	🔵 -0.4% [-1.2%, +0.3%]
WolfSheep	small	🔵 -1.1% [-1.9%, -0.4%]	🔵 -2.2% [-5.3%, +1.1%]
WolfSheep	large	🔵 -2.3% [-3.3%, -1.3%]	🟢 -6.6% [-7.6%, -5.5%]
BoidFlockers	small	🔵 -0.2% [-1.1%, +0.8%]	🔵 -2.0% [-2.7%, -1.3%]
BoidFlockers	large	🔵 -0.8% [-1.7%, -0.0%]	🔵 -3.4% [-4.1%, -2.6%]

[EwoutH]

Thanks for the quick reviews!

It uses self.random, which can be any rng.

Good catch, fixed.

~10% on BoltzmannWealth and ~6.5% on WolfSheep runtime reduction. That's quite serious for a single optimized function!

[quaquel]

yes, it would be nice to do (for ourselves first) a benchmark comparison between 2.1.4 and an alpha 3.0 release. 2.1.4 is the latest version used in the ABM performance comparison.