adding fixed initial positions (to assert equality between simulations)

projectmesa · Sep 20, 2024 · ced1d16 · ced1d16
1 parent 845aa61
commit ced1d16
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 10 deletions.
diff --git a/examples/sugarscape_ig/performance_comparison.py b/examples/sugarscape_ig/performance_comparison.py
@@ -3,6 +3,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import perfplot
+import polars as pl
 import seaborn as sns
 from ss_mesa.model import SugarscapeMesa
 from ss_pandas.model import SugarscapePandas
@@ -28,6 +29,24 @@ def __init__(self, n: int):
         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)
+self.initial_positions = pl.DataFrame(
+            schema={"dim_0": pl.Int64, "dim_1": pl.Int64}
+        )
+        while self.initial_positions.shape[0] < n:
+            initial_pos_0 = np.random.randint(
+                0, dimension, n - self.initial_positions.shape[0]
+            )
+            initial_pos_1 = np.random.randint(
+                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()
 
 
 def mesa_implementation(setup: SugarScapeSetup):
@@ -43,58 +62,68 @@ def mesa_frames_pandas_concise(setup: SugarScapeSetup):
 
 
 def mesa_frames_polars_loop_DF(setup: SugarScapeSetup):
-    return SugarscapePolars(
+    model = SugarscapePolars(
         AntPolarsLoopDF,
         setup.n,
         setup.sugar_grid,
         setup.initial_sugar,
         setup.metabolism,
         setup.vision,
-    ).run_model(100)
+    setup.initial_positions,
+    )
+    model.run_model(100)
 
 
 def mesa_frames_polars_loop_no_vec(setup: SugarScapeSetup):
-    return SugarscapePolars(
+    model = SugarscapePolars(
         AntPolarsLoopNoVec,
         setup.n,
         setup.sugar_grid,
         setup.initial_sugar,
         setup.metabolism,
         setup.vision,
-    ).run_model(100)
+    setup.initial_positions,
+    )
+    model.run_model(100)
 
 
 def mesa_frames_polars_numba_cpu(setup: SugarScapeSetup):
-    return SugarscapePolars(
+    model = SugarscapePolars(
         AntPolarsNumbaCPU,
         setup.n,
         setup.sugar_grid,
         setup.initial_sugar,
         setup.metabolism,
         setup.vision,
-    ).run_model(100)
+    setup.initial_positions,
+    )
+    model.run_model(100)
 
 
 def mesa_frames_polars_numba_gpu(setup: SugarScapeSetup):
-    return SugarscapePolars(
+    model = SugarscapePolars(
         AntPolarsNumbaGPU,
         setup.n,
         setup.sugar_grid,
         setup.initial_sugar,
         setup.metabolism,
         setup.vision,
-    ).run_model(100)
+    setup.initial_positions,
+    )
+    model.run_model(100)
 
 
 def mesa_frames_polars_numba_parallel(setup: SugarScapeSetup):
-    return SugarscapePolars(
+    model = SugarscapePolars(
         AntPolarsNumbaParallel,
         setup.n,
         setup.sugar_grid,
         setup.initial_sugar,
         setup.metabolism,
         setup.vision,
-    ).run_model(100)
+    setup.initial_positions,
+    )
+    model.run_model(100)
 
 
 def plot_and_print_benchmark(labels, kernels, n_range, title, image_path):

diff --git a/examples/sugarscape_ig/ss_polars/model.py b/examples/sugarscape_ig/ss_polars/model.py
@@ -15,6 +15,7 @@ def __init__(
         initial_sugar: np.ndarray | None = None,
         metabolism: np.ndarray | None = None,
         vision: np.ndarray | None = None,
+initial_positions: pl.DataFrame | None = None,
         width: int | None = None,
         height: int | None = None,
     ):
@@ -32,6 +33,9 @@ def __init__(
         )
         self.space.set_cells(sugar_grid)
         self.agents += agent_type(self, n_agents, initial_sugar, metabolism, vision)
+if initial_positions is not None:
+            self.space.place_agents(self.agents, initial_positions)
+        else:
         self.space.place_to_empty(self.agents)
 
     def run_model(self, steps: int) -> list[int]: