examples: Adding Sugarscape IG - polars with numba vectorized loop (#91)

* adding numba to project.toml

* splitting agent_type

* splitting antpolars and antpolarsloop

* adding py-spy to dev tools

* splitting between numba, loop with DF and loop non vectorized implementations

* performance comparison between loop and numba

* fix: cuda target, not gpu

* adding polars comparison

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

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

* correcting dtypes

* adding writable_args (see https://numba.readthedocs.io/en/stable/user/vectorize.html#overwriting-input-values)

* processed_agents is of type bool

* fix: occupied cells sorted by agent_order, assertion to verify best moves aren't duplicated

* adding fixed initial positions (to assert equality between simulations)

* adding complete reproducibility via seed

* fix: taking into account potential (max) sugar when creating the neighborhood. For both numba and completely vectorized it's easier to reason this way then update the current sugar and "best moves" ranking when agents move

* fix: considering priority (if there are previous order agents that might make the same move and haven't found the optimal move yet). This avoids race conditions.

* fix: formatting

* clarifying with comments

* updating actual sugar before executing the step (NOTE: might be unncessary since we prepare the neighborhood looking at potential/max sugar anyway)

* whitespace fixes

* adding initial_positions and seed to mesa model

* renaming grid to space for mesa_models

* fix: best_moves only uses neighborhood and not agent_order

* adding documentation and type hints

* fix: logic for the priority condition (right order of parentheses)

* removing assertion (testing purposes only)

* adding equality_check on model state

* changing n_range to reflect million of agents

* fix: changing callable to typing.Callable

* removing extra requirement

* removing flame_graph (it was just a one-off, we can add memory-profiling in the future)

* fix: comparing DFs

* removing outdated picture

* adding mesa_comparison.png

* adding polars_cocmparison.png

* updating kernels for comparison

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

docs/general/user-guide/4_benchmarks.md

@@ -18,4 +18,4 @@ mesa-frames offers significant performance improvements over the original mesa f
 
 [View the benchmark script](https://github.com/projectmesa/mesa-frames/blob/main/examples/sugarscape_ig/performance_comparison.py)
 
-![Performance Graph SS IG](https://github.com/projectmesa/mesa-frames/raw/main/examples/sugarscape_ig/benchmark_plot_0.png)
+![Performance Graph SS IG](https://github.com/projectmesa/mesa-frames/raw/main/examples/sugarscape_ig/mesa_comparison.png)

examples/sugarscape_ig/performance_comparison.py

@@ -3,10 +3,20 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import perfplot
+import polars as pl
+import seaborn as sns
+from polars.testing import assert_frame_equal
 from ss_mesa.model import SugarscapeMesa
 from ss_pandas.model import SugarscapePandas
+from ss_polars.agents import (
+    AntPolarsLoopDF,
+    AntPolarsLoopNoVec,
+    AntPolarsNumbaCPU,
+    AntPolarsNumbaGPU,
+    AntPolarsNumbaParallel,
+)
 from ss_polars.model import SugarscapePolars
-
+from typing_extensions import Callable
 
 class SugarScapeSetup:
     def __init__(self, n: int):
@@ -15,39 +25,152 @@ def __init__(self, n: int):
         else:
             density = 0.04  # mesa
         self.n = n
+        self.seed = 42
         dimension = math.ceil(math.sqrt(n / density))
-        self.sugar_grid = np.random.randint(0, 4, (dimension, dimension))
-        self.initial_sugar = np.random.randint(6, 25, n)
-        self.metabolism = np.random.randint(2, 4, n)
-        self.vision = np.random.randint(1, 6, n)
+        random_gen = np.random.default_rng(self.seed)
+        self.sugar_grid = random_gen.integers(0, 4, (dimension, dimension))
+        self.initial_sugar = random_gen.integers(6, 25, n)
+        self.metabolism = random_gen.integers(2, 4, n)
+        self.vision = random_gen.integers(1, 6, n)
+        self.initial_positions = pl.DataFrame(
+            schema={"dim_0": pl.Int64, "dim_1": pl.Int64}
+        )
+        while self.initial_positions.shape[0] < n:
+            initial_pos_0 = random_gen.integers(
+                0, dimension, n - self.initial_positions.shape[0]
+            )
+            initial_pos_1 = random_gen.integers(
+                0, dimension, n - self.initial_positions.shape[0]
+            )
+            self.initial_positions = self.initial_positions.vstack(
+                pl.DataFrame(
+                    {
+                        "dim_0": initial_pos_0,
+                        "dim_1": initial_pos_1,
+                    }
+                )
+            ).unique(maintain_order=True)
+        return
 
 
 def mesa_implementation(setup: SugarScapeSetup):
-    return SugarscapeMesa(
-        setup.n, setup.sugar_grid, setup.initial_sugar, setup.metabolism, setup.vision
-    ).run_model(100)
+    model = SugarscapeMesa(
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
 
 
 def mesa_frames_pandas_concise(setup: SugarScapeSetup):
-    return SugarscapePandas(
-        setup.n, setup.sugar_grid, setup.initial_sugar, setup.metabolism, setup.vision
-    ).run_model(100)
-
-
-def mesa_frames_polars_concise(setup: SugarScapeSetup):
-    return SugarscapePolars(
-        setup.n, setup.sugar_grid, setup.initial_sugar, setup.metabolism, setup.vision
-    ).run_model(100)
-
-
-def plot_and_print_benchmark(labels, kernels, n_range, title, image_path):
+    model = SugarscapePandas(
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
+
+
+def mesa_frames_polars_loop_DF(setup: SugarScapeSetup):
+    model = SugarscapePolars(
+        AntPolarsLoopDF,
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
+
+
+def mesa_frames_polars_loop_no_vec(setup: SugarScapeSetup):
+    model = SugarscapePolars(
+        AntPolarsLoopNoVec,
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
+
+
+def mesa_frames_polars_numba_cpu(setup: SugarScapeSetup):
+    model = SugarscapePolars(
+        AntPolarsNumbaCPU,
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
+
+
+def mesa_frames_polars_numba_gpu(setup: SugarScapeSetup):
+    model = SugarscapePolars(
+        AntPolarsNumbaGPU,
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
+
+
+def mesa_frames_polars_numba_parallel(setup: SugarScapeSetup):
+    model = SugarscapePolars(
+        AntPolarsNumbaParallel,
+        setup.n,
+        setup.sugar_grid,
+        setup.initial_sugar,
+        setup.metabolism,
+        setup.vision,
+        setup.initial_positions,
+        setup.seed,
+    )
+    model.run_model(100)
+    return model
+
+
+def plot_and_print_benchmark(
+    labels: list[str],
+    kernels: list[Callable],
+    n_range: list[int],
+    title: str,
+    image_path: str,
+    equality_check: Callable | None = None,
+):
     out = perfplot.bench(
         setup=SugarScapeSetup,
         kernels=kernels,
         labels=labels,
         n_range=n_range,
         xlabel="Number of agents",
-        equality_check=None,
+        equality_check=equality_check,
         title=title,
     )
     plt.ylabel("Execution time (s)")
@@ -60,37 +183,58 @@ def plot_and_print_benchmark(labels, kernels, n_range, title, image_path):
         print("---------------")
 
 
+def polars_equality_check(a: SugarscapePolars, b: SugarscapePolars):
+    assert_frame_equal(a.space.agents, b.space.agents, check_row_order=False)
+    assert_frame_equal(a.space.cells, b.space.cells, check_row_order=False)
+    return True
+
+
 def main():
-    """# Mesa comparison
+    # Mesa comparison
     sns.set_theme(style="whitegrid")
     labels_0 = [
+        # "mesa-frames (pd concise)", # Pandas to be removed because of performance
+        "mesa-frames (pl numba parallel)",
         "mesa",
-        # "mesa-frames (pd concise)",
-        "mesa-frames (pl concise)",
     ]
     kernels_0 = [
-        mesa_implementation,
         # mesa_frames_pandas_concise,
-        mesa_frames_polars_concise,
+        mesa_frames_polars_numba_parallel,
+        mesa_implementation,
     ]
-    n_range_0 = [k for k in range(1, 100002, 10000)]
+    n_range_0 = [k for k in range(10**5, 5*10**5 + 2, 10**5)]
     title_0 = "100 steps of the SugarScape IG model:\n" + " vs ".join(labels_0)
-    image_path_0 = "benchmark_plot_0.png"
-    plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0)"""
+    image_path_0 = "mesa_comparison.png"
+    plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0)
 
-    # FLAME2-GPU comparison
+    # mesa-frames comparison
     labels_1 = [
         # "mesa-frames (pd concise)",
-        "mesa-frames (pl concise)",
+        "mesa-frames (pl loop DF)",
+        "mesa-frames (pl loop no vec)",
+        "mesa-frames (pl numba CPU)",
+        "mesa-frames (pl numba parallel)",
+        "mesa-frames (pl numba GPU)",
     ]
+    # Polars best_moves (non-vectorized loop vs DF loop vs numba loop)
     kernels_1 = [
-        # mesa_frames_pandas_concise,
-        mesa_frames_polars_concise,
+        mesa_frames_polars_loop_DF,
+        mesa_frames_polars_loop_no_vec,
+        mesa_frames_polars_numba_cpu,
+        mesa_frames_polars_numba_parallel,
+        mesa_frames_polars_numba_gpu,
     ]
-    n_range_1 = [k for k in range(1, 3 * 10**6 + 2, 10**6)]
+    n_range_1 = [k for k in range(10**6, 3 * 10**6 + 2, 10**6)]
     title_1 = "100 steps of the SugarScape IG model:\n" + " vs ".join(labels_1)
-    image_path_1 = "benchmark_plot_1.png"
-    plot_and_print_benchmark(labels_1, kernels_1, n_range_1, title_1, image_path_1)
+    image_path_1 = "polars_comparison.png"
+    plot_and_print_benchmark(
+        labels_1,
+        kernels_1,
+        n_range_1,
+        title_1,
+        image_path_1,
+        equality_check=polars_equality_check,
+    )
 
 
 if __name__ == "__main__":

examples/sugarscape_ig/ss_mesa/agents.py

@@ -25,20 +25,20 @@ def __init__(self, unique_id, model, moore=False, sugar=0, metabolism=0, vision=
         self.vision = vision
 
     def get_sugar(self, pos):
-        this_cell = self.model.grid.get_cell_list_contents([pos])
+        this_cell = self.model.space.get_cell_list_contents([pos])
         for agent in this_cell:
             if type(agent) is Sugar:
                 return agent
 
     def is_occupied(self, pos):
-        this_cell = self.model.grid.get_cell_list_contents([pos])
+        this_cell = self.model.space.get_cell_list_contents([pos])
         return any(isinstance(agent, AntMesa) for agent in this_cell)
 
     def move(self):
         # Get neighborhood within vision
         neighbors = [
             i
-            for i in self.model.grid.get_neighborhood(
+            for i in self.model.space.get_neighborhood(
                 self.pos, self.moore, False, radius=self.vision
             )
             if not self.is_occupied(i)
@@ -55,7 +55,7 @@ def move(self):
             pos for pos in candidates if get_distance(self.pos, pos) == min_dist
         ]
         self.random.shuffle(final_candidates)
-        self.model.grid.move_agent(self, final_candidates[0])
+        self.model.space.move_agent(self, final_candidates[0])
 
     def eat(self):
         sugar_patch = self.get_sugar(self.pos)
@@ -66,7 +66,7 @@ def step(self):
         self.move()
         self.eat()
         if self.sugar <= 0:
-            self.model.grid.remove_agent(self)
+            self.model.space.remove_agent(self)
             self.model.agents.remove(self)
 
 
@@ -77,7 +77,7 @@ def __init__(self, unique_id, model, max_sugar):
         self.max_sugar = max_sugar
 
     def step(self):
-        if self.model.grid.is_cell_empty(self.pos):
+        if self.model.space.is_cell_empty(self.pos):
             self.amount = self.max_sugar
         else:
             self.amount = 0

examples/sugarscape_ig/ss_mesa/model.py

@@ -1,5 +1,6 @@
 import mesa
 import numpy as np
+import polars as pl
 
 from .agents import AntMesa, Sugar
 
@@ -16,6 +17,8 @@ def __init__(
         initial_sugar: np.ndarray | None = None,
         metabolism: np.ndarray | None = None,
         vision: np.ndarray | None = None,
+        initial_positions: pl.DataFrame | None = None,
+        seed: int | None = None,
         width: int | None = None,
         height: int | None = None,
     ):
@@ -34,30 +37,37 @@ def __init__(
             metabolism = np.random.randint(2, 4, n_agents)
         if vision is None:
             vision = np.random.randint(1, 6, n_agents)
+        if seed is not None:
+            self.reset_randomizer(seed)
 
         self.width, self.height = sugar_grid.shape
         self.n_agents = n_agents
-        self.grid = mesa.space.MultiGrid(self.width, self.height, torus=False)
+        self.space = mesa.space.MultiGrid(self.width, self.height, torus=False)
         self.agents: list = []
 
         agent_id = 0
         self.sugars = []
-        for _, (x, y) in self.grid.coord_iter():
+
+        for _, (x, y) in self.space.coord_iter():
             max_sugar = sugar_grid[x, y]
             sugar = Sugar(agent_id, self, max_sugar)
             agent_id += 1
-            self.grid.place_agent(sugar, (x, y))
+            self.space.place_agent(sugar, (x, y))
             self.sugars.append(sugar)
 
         # Create agent:
         for i in range(self.n_agents):
-            x = self.random.randrange(self.width)
-            y = self.random.randrange(self.height)
+            if initial_positions is not None:
+                x = initial_positions["dim_0"][i]
+                y = initial_positions["dim_1"][i]
+            else:
+                x = self.random.randrange(self.width)
+                y = self.random.randrange(self.height)
             ssa = AntMesa(
                 agent_id, self, False, initial_sugar[i], metabolism[i], vision[i]
             )
             agent_id += 1
-            self.grid.place_agent(ssa, (x, y))
+            self.space.place_agent(ssa, (x, y))
             self.agents.append(ssa)
 
         self.running = True