Repository: projectmesa/mesa
Pull request #1985 (2024-01-21 10:22:37)

Minor edits to benchmarking code
* fix for < python 3.12
* Insert the current repo of mesa into sys.path to ensure you are using the correct version
* pep8 name changes of files
* added code to run each benchmark in isolation (useful for profiling and line profiling, used it extensively already for https://github.com/projectmesa/mesa/issues/1900#issuecomment-1902643678)


Patch:
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
new file mode 100644
index 00000000..09eb1f95
--- /dev/null
+++ b/.github/workflows/benchmarks.yml
@@ -0,0 +1,77 @@
+name: Performance benchmarks
+
+on:
+  pull_request_target:
+    types: [opened, ready_for_review]
+    branches:
+      - main
+  issue_comment:
+    types: [created]
+
+permissions:
+  issues: write
+  pull-requests: write
+
+jobs:
+  run-benchmarks:
+    if: >
+      github.event_name == 'pull_request_target' ||
+      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '/rerun-benchmarks'))
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout PR branch
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+      - name: Add project directory to PYTHONPATH
+        run: echo "PYTHONPATH=$PYTHONPATH:$(pwd)" >> $GITHUB_ENV
+      - name: Install dependencies
+        run: pip install numpy pandas tqdm tabulate
+      - name: Checkout main branch
+        uses: actions/checkout@v4
+        with:
+          ref: main
+          repository: projectmesa/mesa
+      - name: Run benchmarks on main branch
+        working-directory: benchmarks
+        run: python global_benchmark.py
+      - name: Upload benchmark results
+        uses: actions/upload-artifact@v4
+        with:
+          name: timings-main
+          path: benchmarks/timings_1.pickle
+      - name: Checkout PR branch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.head.ref }}
+          repository: ${{ github.event.pull_request.head.repo.full_name }}
+          fetch-depth: 0
+          persist-credentials: false
+          clean: false
+      - name: Download benchmark results
+        uses: actions/download-artifact@v4
+        with:
+          name: timings-main
+          path: benchmarks
+      - name: Run benchmarks on PR branch
+        working-directory: benchmarks
+        run: python global_benchmark.py
+      - name: Run compare timings and encode output
+        working-directory: benchmarks
+        run: |
+          TIMING_COMPARISON=$(python compare_timings.py | base64 -w 0)  # Base64 encode the output
+          echo "TIMING_COMPARISON=$TIMING_COMPARISON" >> $GITHUB_ENV
+      - name: Comment PR
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const output = Buffer.from(process.env.TIMING_COMPARISON, 'base64').toString('utf-8');
+            const issue_number = context.issue.number;
+            github.rest.issues.createComment({
+              issue_number: issue_number,
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              body: 'Performance benchmarks:\n\n' + output
+            });
diff --git a/.gitignore b/.gitignore
index 1b04c123..814f6a5d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,5 @@
 # Benchmarking
-benchmarking/**/*.pickle
+benchmarks/**/*.pickle
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -69,6 +69,7 @@ target/
 
 # PyCharm environment files
 .idea/
+*.pclprof
 
 # VSCode environment files
 .vscode/
@@ -88,3 +89,4 @@ pythonenv*/
 # JS dependencies
 mesa/visualization/templates/external/
 mesa/visualization/templates/js/external/
+
diff --git a/benchmarks/Flocking/Flocking.py b/benchmarks/Flocking/Flocking.py
deleted file mode 100644
index 4a5e8529..00000000
--- a/benchmarks/Flocking/Flocking.py
+++ /dev/null
@@ -1,80 +0,0 @@
-"""
-Flockers
-=============================================================
-A Mesa implementation of Craig Reynolds's Boids flocker model.
-Uses numpy arrays to represent vectors.
-"""
-
-import numpy as np
-
-from mesa import Model
-from mesa.space import ContinuousSpace
-from mesa.time import RandomActivation
-
-from .boid import Boid
-
-
-class BoidFlockers(Model):
-    """
-    Flocker model class. Handles agent creation, placement and scheduling.
-    """
-
-    def __init__(
-        self,
-        seed,
-        population,
-        width,
-        height,
-        vision,
-        speed=1,
-        separation=1,
-        cohere=0.03,
-        separate=0.015,
-        match=0.05,
-    ):
-        """
-        Create a new Flockers model.
-
-        Args:
-            population: Number of Boids
-            width, height: Size of the space.
-            speed: How fast should the Boids move.
-            vision: How far around should each Boid look for its neighbors
-            separation: What's the minimum distance each Boid will attempt to
-                    keep from any other
-            cohere, separate, match: factors for the relative importance of
-                    the three drives."""
-        super().__init__(seed=seed)
-        self.population = population
-        self.vision = vision
-        self.speed = speed
-        self.separation = separation
-        self.schedule = RandomActivation(self)
-        self.space = ContinuousSpace(width, height, True)
-        self.factors = {"cohere": cohere, "separate": separate, "match": match}
-        self.make_agents()
-
-    def make_agents(self):
-        """
-        Create self.population agents, with random positions and starting headings.
-        """
-        for i in range(self.population):
-            x = self.random.random() * self.space.x_max
-            y = self.random.random() * self.space.y_max
-            pos = np.array((x, y))
-            velocity = np.random.random(2) * 2 - 1
-            boid = Boid(
-                i,
-                self,
-                pos,
-                self.speed,
-                velocity,
-                self.vision,
-                self.separation,
-                **self.factors,
-            )
-            self.space.place_agent(boid, pos)
-            self.schedule.add(boid)
-
-    def step(self):
-        self.schedule.step()
diff --git a/benchmarks/Flocking/boid.py b/benchmarks/Flocking/flocking.py
similarity index 52%
rename from benchmarks/Flocking/boid.py
rename to benchmarks/Flocking/flocking.py
index 598ea32b..0331d356 100644
--- a/benchmarks/Flocking/boid.py
+++ b/benchmarks/Flocking/flocking.py
@@ -1,6 +1,15 @@
+"""
+Flockers
+=============================================================
+A Mesa implementation of Craig Reynolds's Boids flocker model.
+Uses numpy arrays to represent vectors.
+"""
+
 import numpy as np
 
-from mesa import Agent
+from mesa import Agent, Model
+from mesa.space import ContinuousSpace
+from mesa.time import RandomActivation
 
 
 class Boid(Agent):
@@ -79,3 +88,82 @@ class Boid(Agent):
         self.velocity /= np.linalg.norm(self.velocity)
         new_pos = self.pos + self.velocity * self.speed
         self.model.space.move_agent(self, new_pos)
+
+
+class BoidFlockers(Model):
+    """
+    Flocker model class. Handles agent creation, placement and scheduling.
+    """
+
+    def __init__(
+        self,
+        seed,
+        population,
+        width,
+        height,
+        vision,
+        speed=1,
+        separation=1,
+        cohere=0.03,
+        separate=0.015,
+        match=0.05,
+    ):
+        """
+        Create a new Flockers model.
+
+        Args:
+            population: Number of Boids
+            width, height: Size of the space.
+            speed: How fast should the Boids move.
+            vision: How far around should each Boid look for its neighbors
+            separation: What's the minimum distance each Boid will attempt to
+                    keep from any other
+            cohere, separate, match: factors for the relative importance of
+                    the three drives."""
+        super().__init__(seed=seed)
+        self.population = population
+        self.vision = vision
+        self.speed = speed
+        self.separation = separation
+        self.schedule = RandomActivation(self)
+        self.space = ContinuousSpace(width, height, True)
+        self.factors = {"cohere": cohere, "separate": separate, "match": match}
+        self.make_agents()
+
+    def make_agents(self):
+        """
+        Create self.population agents, with random positions and starting headings.
+        """
+        for i in range(self.population):
+            x = self.random.random() * self.space.x_max
+            y = self.random.random() * self.space.y_max
+            pos = np.array((x, y))
+            velocity = np.random.random(2) * 2 - 1
+            boid = Boid(
+                i,
+                self,
+                pos,
+                self.speed,
+                velocity,
+                self.vision,
+                self.separation,
+                **self.factors,
+            )
+            self.space.place_agent(boid, pos)
+            self.schedule.add(boid)
+
+    def step(self):
+        self.schedule.step()
+
+
+if __name__ == "__main__":
+    import time
+
+    # model = BoidFlockers(15, 200, 100, 100, 5)
+    model = BoidFlockers(15, 400, 100, 100, 15)
+
+    start_time = time.perf_counter()
+    for _ in range(100):
+        model.step()
+
+    print(time.perf_counter() - start_time)
diff --git a/benchmarks/Schelling/Schelling.py b/benchmarks/Schelling/schelling.py
similarity index 74%
rename from benchmarks/Schelling/Schelling.py
rename to benchmarks/Schelling/schelling.py
index eb1379ed..624dd0b3 100644
--- a/benchmarks/Schelling/Schelling.py
+++ b/benchmarks/Schelling/schelling.py
@@ -1,4 +1,3 @@
-import random
 
 from mesa import Agent, Model
 from mesa.space import SingleGrid
@@ -10,7 +9,7 @@ class SchellingAgent(Agent):
     Schelling segregation agent
     """
 
-    def __init__(self, pos, model, agent_type):
+    def __init__(self, unique_id, model, agent_type):
         """
         Create a new Schelling agent.
         Args:
@@ -18,14 +17,12 @@ class SchellingAgent(Agent):
            x, y: Agent initial location.
            agent_type: Indicator for the agent's type (minority=1, majority=0)
         """
-        super().__init__(pos, model)
-        self.pos = pos
+        super().__init__(unique_id, model)
         self.type = agent_type
 
     def step(self):
         similar = 0
-        r = self.model.radius
-        for neighbor in self.model.grid.iter_neighbors(self.pos, moore=True, radius=r):
+        for neighbor in self.model.grid.iter_neighbors(self.pos, moore=True, radius=self.model.radius):
             if neighbor.type == self.type:
                 similar += 1
 
@@ -36,7 +33,7 @@ class SchellingAgent(Agent):
             self.model.happy += 1
 
 
-class SchellingModel(Model):
+class Schelling(Model):
     """
     Model class for the Schelling segregation model.
     """
@@ -63,9 +60,9 @@ class SchellingModel(Model):
         # the coordinates of a cell as well as
         # its contents. (coord_iter)
         for _cont, pos in self.grid.coord_iter():
-            if random.random() < self.density:  # noqa: S311
-                agent_type = 1 if random.random() < self.minority_pc else 0  # noqa: S311
-                agent = SchellingAgent(pos, self, agent_type)
+            if self.random.random() < self.density:
+                agent_type = 1 if self.random.random() < self.minority_pc else 0
+                agent = SchellingAgent(self.next_id(), self, agent_type)
                 self.grid.place_agent(agent, pos)
                 self.schedule.add(agent)
 
@@ -75,3 +72,16 @@ class SchellingModel(Model):
         """
         self.happy = 0  # Reset counter of happy agents
         self.schedule.step()
+
+
+if __name__ == "__main__":
+    import time
+
+    # model = Schelling(15, 40, 40, 3, 1, 0.625)
+    model = Schelling(15, 100, 100, 8, 2, 0.8)
+
+    start_time = time.perf_counter()
+    for _ in range(100):
+        model.step()
+
+    print(time.perf_counter() - start_time)
diff --git a/benchmarks/WolfSheep/__init__.py b/benchmarks/WolfSheep/__init__.py
index e69de29b..98b1e9fd 100644
--- a/benchmarks/WolfSheep/__init__.py
+++ b/benchmarks/WolfSheep/__init__.py
@@ -0,0 +1,14 @@
+from .wolf_sheep import WolfSheep
+
+if __name__ == "__main__":
+    # for profiling this benchmark model
+    import time
+
+    # model = WolfSheep(15, 25, 25, 60, 40, 0.2, 0.1, 20)
+    model = WolfSheep(15, 100, 100, 1000, 500, 0.4, 0.2, 20)
+
+    start_time = time.perf_counter()
+    for _ in range(100):
+        model.step()
+
+    print(time.perf_counter() - start_time)
diff --git a/benchmarks/WolfSheep/WolfSheep.py b/benchmarks/WolfSheep/wolf_sheep.py
similarity index 100%
rename from benchmarks/WolfSheep/WolfSheep.py
rename to benchmarks/WolfSheep/wolf_sheep.py
diff --git a/benchmarks/configurations.py b/benchmarks/configurations.py
index ab896040..c42f69ad 100644
--- a/benchmarks/configurations.py
+++ b/benchmarks/configurations.py
@@ -1,10 +1,10 @@
-from Flocking.Flocking import BoidFlockers
-from Schelling.Schelling import SchellingModel
-from WolfSheep.WolfSheep import WolfSheep
+from Flocking.flocking import BoidFlockers
+from Schelling.schelling import Schelling
+from WolfSheep.wolf_sheep import WolfSheep
 
 configurations = {
     # Schelling Model Configurations
-    SchellingModel: {
+    Schelling: {
         "small": {
             "seeds": 50,
             "replications": 5,
diff --git a/benchmarks/global_benchmark.py b/benchmarks/global_benchmark.py
index 9547d94a..0f8d5379 100644
--- a/benchmarks/global_benchmark.py
+++ b/benchmarks/global_benchmark.py
@@ -5,6 +5,10 @@ import sys
 import time
 import timeit
 
+# making sure we use this version of mesa and not one
+# also installed in site_packages or so.
+sys.path.insert(0, os.path.abspath(".."))
+
 from configurations import configurations
 
 
@@ -56,7 +60,7 @@ for model, model_config in configurations.items():
         mean_run = sum(results[1]) / len(results[1])
 
         print(
-            f"{time.strftime("%H:%M:%S", time.localtime())} {model.__name__:<14} ({size}) timings: Init {mean_init:.5f} s; Run {mean_run:.4f} s"
+            f"{time.strftime('%H:%M:%S', time.localtime())} {model.__name__:<14} ({size}) timings: Init {mean_init:.5f} s; Run {mean_run:.4f} s"
         )
 
         results_dict[model, size] = results
diff --git a/mesa/space.py b/mesa/space.py
index b730c2da..9be3b563 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -1539,9 +1539,11 @@ class NetworkGrid:
             neighborhood = sorted(neighbors_with_distance.keys())
         return neighborhood
 
-    def get_neighbors(self, node_id: int, include_center: bool = False) -> list[Agent]:
-        """Get all agents in adjacent nodes."""
-        neighborhood = self.get_neighborhood(node_id, include_center)
+    def get_neighbors(
+        self, node_id: int, include_center: bool = False, radius: int = 1
+    ) -> list[Agent]:
+        """Get all agents in adjacent nodes (within a certain radius)."""
+        neighborhood = self.get_neighborhood(node_id, include_center, radius)
         return self.get_cell_list_contents(neighborhood)
 
     def move_agent(self, agent: Agent, node_id: int) -> None:


Test patch:
diff --git a/tests/test_space.py b/tests/test_space.py
index b7524b6d..539d2c0e 100644
--- a/tests/test_space.py
+++ b/tests/test_space.py
@@ -867,12 +867,54 @@ class TestSingleNetworkGrid(unittest.TestCase):
             a = self.agents[i]
             assert a.pos == pos
 
-    def test_get_neighbors(self):
+    def test_get_neighborhood(self):
         assert len(self.space.get_neighborhood(0, include_center=True)) == 3
         assert len(self.space.get_neighborhood(0, include_center=False)) == 2
         assert len(self.space.get_neighborhood(2, include_center=True, radius=3)) == 7
         assert len(self.space.get_neighborhood(2, include_center=False, radius=3)) == 6
 
+    def test_get_neighbors(self):
+        """
+        Test the get_neighbors method with varying radius and include_center values. Note there are agents on node 0, 1 and 5.
+        """
+        # Test with default radius (1) and include_center = False
+        neighbors_default = self.space.get_neighbors(0, include_center=False)
+        self.assertEqual(
+            len(neighbors_default),
+            1,
+            "Should have 1 neighbors with default radius and exclude center",
+        )
+
+        # Test with default radius (1) and include_center = True
+        neighbors_include_center = self.space.get_neighbors(0, include_center=True)
+        self.assertEqual(
+            len(neighbors_include_center),
+            2,
+            "Should have 2 neighbors (including center) with default radius",
+        )
+
+        # Test with radius = 2 and include_center = False
+        neighbors_radius_2 = self.space.get_neighbors(0, include_center=False, radius=5)
+        expected_count_radius_2 = 2
+        self.assertEqual(
+            len(neighbors_radius_2),
+            expected_count_radius_2,
+            f"Should have {expected_count_radius_2} neighbors with radius 2 and exclude center",
+        )
+
+        # Test with radius = 2 and include_center = True
+        neighbors_radius_2_include_center = self.space.get_neighbors(
+            0, include_center=True, radius=5
+        )
+        expected_count_radius_2_include_center = (
+            3  # Adjust this based on your network structure
+        )
+        self.assertEqual(
+            len(neighbors_radius_2_include_center),
+            expected_count_radius_2_include_center,
+            f"Should have {expected_count_radius_2_include_center} neighbors (including center) with radius 2",
+        )
+
     def test_move_agent(self):
         initial_pos = 1
         agent_number = 1