From 449d230e1eaef86a3b49b7d424238cbbed9aa194 Mon Sep 17 00:00:00 2001
From: Jan Kwakkel <j.h.kwakkel@tudelft.nl>
Date: Sat, 27 Jan 2024 22:41:58 +0100
Subject: [PATCH] perf: Speed up Agentset.shuffle (#2010)

* Make RandomActivationByType.agents_by_type backward compatible

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* updated warning

* Update .gitignore

* Update global_benchmark.py

* pep8 rename

* Update schelling.py

* Update schelling.py

* Update schelling.py

* Update schelling.py

* Squashed commit of the following:

commit c0de4a1c4dbb586f26084764d6f34818ce1e5dfb
Author: Ewout ter Hoeven <E.M.terHoeven@student.tudelft.nl>
Date:   Sun Jan 21 14:43:08 2024 +0100

    Add CI workflow for performance benchmarks (#1983)

    This commit introduces a new GitHub Actions workflow for performance benchmarking. The workflow is triggered on `pull_request_target` events and can be triggered with a "/rerun-benchmarks" comment in the PR. It should be compatible with PRs from both forks and the main repository. It includes steps for setting up the environment, checking out the PR and main branches, installing dependencies, and running benchmarks on both branches. The results are then compared, encoded, and posted as a comment on the PR.

* Revert "Squashed commit of the following:"

This reverts commit 948508754eaabb2312cef6863c467ea905069551.

* Update .gitignore

* further updates

* Update benchmarks/WolfSheep/__init__.py

* shuffle update

* switch to walrus operator in do

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* revert inplace to False

* remove old code

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 mesa/agent.py       | 26 ++++++++++++++++----------
 tests/test_agent.py |  2 +-
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/mesa/agent.py b/mesa/agent.py
index 6155284ba42..243a187337b 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -188,15 +188,21 @@ def shuffle(self, inplace: bool = False) -> AgentSet:
 
         Returns:
             AgentSet: A shuffled AgentSet. Returns the current AgentSet if inplace is True.
-        """
-        shuffled_agents = list(self)
-        self.random.shuffle(shuffled_agents)
 
-        return (
-            AgentSet(shuffled_agents, self.model)
-            if not inplace
-            else self._update(shuffled_agents)
-        )
+        Note:
+            Using inplace = True is more performant
+
+        """
+        weakrefs = list(self._agents.keyrefs())
+        self.random.shuffle(weakrefs)
+
+        if inplace:
+            self._agents.data = {entry: None for entry in weakrefs}
+            return self
+        else:
+            return AgentSet(
+                (agent for ref in weakrefs if (agent := ref()) is not None), self.model
+            )
 
     def sort(
         self,
@@ -251,9 +257,9 @@ def do(
         """
         # we iterate over the actual weakref keys and check if weakref is alive before calling the method
         res = [
-            getattr(agentref(), method_name)(*args, **kwargs)
+            getattr(agent, method_name)(*args, **kwargs)
             for agentref in self._agents.keyrefs()
-            if agentref()
+            if (agent := agentref()) is not None
         ]
 
         return res if return_results else self
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 60a0eec15fa..7ad538eba27 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -66,7 +66,7 @@ def test_function(agent):
     assert all(a1 == a2 for a1, a2 in zip(agentset.select(), agentset))
     assert all(a1 == a2 for a1, a2 in zip(agentset.select(n=5), agentset[:5]))
 
-    assert len(agentset.shuffle().select(n=5)) == 5
+    assert len(agentset.shuffle(inplace=False).select(n=5)) == 5
 
     def test_function(agent):
         return agent.unique_id