Adding SugarScape IG (polars with loops) (#71)

* adding polars implementation

* benchmarking actual execution and comparison with flame2

* performance comparison with polars

* wrong file

* adding matplotlib to docs requirements

* fix: df_constructor when data contains DF/SRS

* fix: substituting directly cells_df if there are properties

* performance: setting the indexes speeds up the merging operation

* performance: inplace operation speed up

examples/sugarscape_ig/performance_comparison.py

@@ -3,15 +3,19 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import perfplot
-import seaborn as sns
 from ss_mesa.model import SugarscapeMesa
 from ss_pandas.model import SugarscapePandas
+from ss_polars.model import SugarscapePolars
 
 
 class SugarScapeSetup:
     def __init__(self, n: int):
+        if n >= 10**6:
+            density = 0.17  # FLAME2-GPU
+        else:
+            density = 0.04  # mesa
         self.n = n
-        dimension = math.ceil(5 * math.sqrt(n))
+        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)
@@ -21,13 +25,19 @@ def __init__(self, n: int):
 def mesa_implementation(setup: SugarScapeSetup):
     return SugarscapeMesa(
         setup.n, setup.sugar_grid, setup.initial_sugar, setup.metabolism, setup.vision
-    )
+    ).run_model(100)
 
 
 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):
@@ -40,10 +50,8 @@ def plot_and_print_benchmark(labels, kernels, n_range, title, image_path):
         equality_check=None,
         title=title,
     )
-
     plt.ylabel("Execution time (s)")
     out.save(image_path)
-
     print("\nExecution times:")
     for i, label in enumerate(labels):
         print(f"---------------\n{label}:")
@@ -53,21 +61,36 @@ def plot_and_print_benchmark(labels, kernels, n_range, title, image_path):
 
 
 def main():
+    """# Mesa comparison
     sns.set_theme(style="whitegrid")
-
     labels_0 = [
         "mesa",
-        "mesa-frames (pd concise)",
+        # "mesa-frames (pd concise)",
+        "mesa-frames (pl concise)",
     ]
     kernels_0 = [
         mesa_implementation,
-        mesa_frames_pandas_concise,
+        # mesa_frames_pandas_concise,
+        mesa_frames_polars_concise,
     ]
-    n_range_0 = [k for k in range(0, 100000, 10000)]
+    n_range_0 = [k for k in range(1, 100002, 10000)]
     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)"""
 
-    plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0)
+    # FLAME2-GPU comparison
+    labels_1 = [
+        # "mesa-frames (pd concise)",
+        "mesa-frames (pl concise)",
+    ]
+    kernels_1 = [
+        # mesa_frames_pandas_concise,
+        mesa_frames_polars_concise,
+    ]
+    n_range_1 = [k for k in range(1, 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)
 
 
 if __name__ == "__main__":

examples/sugarscape_ig/ss_polars/agents.py

@@ -0,0 +1,125 @@
+import numpy as np
+import polars as pl
+
+from mesa_frames import AgentSetPolars, ModelDF
+
+
+class AntPolars(AgentSetPolars):
+    def __init__(
+        self,
+        model: ModelDF,
+        n_agents: int,
+        initial_sugar: np.ndarray | None = None,
+        metabolism: np.ndarray | None = None,
+        vision: np.ndarray | None = None,
+    ):
+        super().__init__(model)
+
+        if initial_sugar is None:
+            initial_sugar = model.random.integers(6, 25, n_agents)
+        if metabolism is None:
+            metabolism = model.random.integers(2, 4, n_agents)
+        if vision is None:
+            vision = model.random.integers(1, 6, n_agents)
+
+        agents = pl.DataFrame(
+            {
+                "unique_id": pl.arange(n_agents, eager=True),
+                "sugar": model.random.integers(6, 25, n_agents),
+                "metabolism": model.random.integers(2, 4, n_agents),
+                "vision": model.random.integers(1, 6, n_agents),
+            }
+        )
+        self.add(agents)
+
+    def move(self):
+        neighborhood: pl.DataFrame = self.space.get_neighborhood(
+            radius=self["vision"], agents=self, include_center=True
+        )
+
+        # Join self.space.cells to obtain properties ('sugar') per cell
+        neighborhood = neighborhood.join(self.space.cells, on=["dim_0", "dim_1"])
+
+        # Join self.pos to obtain the agent_id of the center cell
+        # TODO: get_neighborhood/get_neighbors should return 'agent_id_center' instead of center position when input is AgentLike
+        neighborhood = neighborhood.with_columns(
+            agent_id_center=neighborhood.join(
+                self.pos,
+                left_on=["dim_0_center", "dim_1_center"],
+                right_on=["dim_0", "dim_1"],
+            )["unique_id"]
+        )
+
+        # Order of agents moves based on the original order of agents.
+        # The agent in his cell has order 0 (highest)
+        agent_order = neighborhood.unique(
+            subset=["agent_id_center"], keep="first", maintain_order=True
+        ).with_row_count("agent_order")
+
+        neighborhood = neighborhood.join(agent_order, on="agent_id_center")
+
+        neighborhood = neighborhood.join(
+            agent_order.select(
+                pl.col("agent_id_center").alias("agent_id"),
+                pl.col("agent_order").alias("blocking_agent_order"),
+            ),
+            on="agent_id",
+        )
+
+        # Filter impossible moves
+        neighborhood = neighborhood.filter(
+            pl.col("agent_order") >= pl.col("blocking_agent_order")
+        )
+
+        # Sort cells by sugar and radius (nearest first)
+        neighborhood = neighborhood.sort(["sugar", "radius"], descending=[True, False])
+
+        best_moves = pl.DataFrame()
+        # While there are agents that do not have a best move, keep looking for one
+        while len(best_moves) < len(self.agents):
+            # Get the best moves for each agent and if duplicates are found, select the one with the highest order
+            new_best_moves = (
+                neighborhood.group_by("agent_id_center", maintain_order=True)
+                .first()
+                .sort("agent_order")
+                .unique(subset=["dim_0", "dim_1"], keep="first")
+            )
+
+            # Agents can make the move if:
+            # - There is no blocking agent
+            # - The agent is in its own cell
+            # - The blocking agent has moved before him
+            condition = pl.col("agent_id").is_null() | (
+                pl.col("agent_id") == pl.col("agent_id_center")
+            )
+            if len(best_moves) > 0:
+                condition = condition | pl.col("agent_id").is_in(
+                    best_moves["agent_id_center"]
+                )
+            new_best_moves = new_best_moves.filter(condition)
+
+            best_moves = pl.concat([best_moves, new_best_moves])
+
+            # Remove agents that have already moved
+            neighborhood = neighborhood.filter(
+                ~pl.col("agent_id_center").is_in(best_moves["agent_id_center"])
+            )
+
+            # Remove cells that have been already selected
+            neighborhood = neighborhood.join(
+                best_moves.select(["dim_0", "dim_1"]), on=["dim_0", "dim_1"], how="anti"
+            )
+
+        self.space.move_agents(self, best_moves.select(["dim_0", "dim_1"]))
+
+    def eat(self):
+        cells = self.space.cells.filter(pl.col("agent_id").is_not_null())
+        self[cells["agent_id"], "sugar"] = (
+            self[cells["agent_id"], "sugar"]
+            + cells["sugar"]
+            - self[cells["agent_id"], "metabolism"]
+        )
+
+    def step(self):
+        self.shuffle().do("move").do("eat")
+        self.discard(self.agents.filter(pl.col("sugar") <= 0))

examples/sugarscape_ig/ss_polars/model.py

@@ -0,0 +1,49 @@
+import numpy as np
+import polars as pl
+
+from mesa_frames import GridPolars, ModelDF
+
+from .agents import AntPolars
+
+
+class SugarscapePolars(ModelDF):
+    def __init__(
+        self,
+        n_agents: int,
+        sugar_grid: np.ndarray | None = None,
+        initial_sugar: np.ndarray | None = None,
+        metabolism: np.ndarray | None = None,
+        vision: np.ndarray | None = None,
+        width: int | None = None,
+        height: int | None = None,
+    ):
+        super().__init__()
+        if sugar_grid is None:
+            sugar_grid = self.random.integers(0, 4, (width, height))
+        grid_dimensions = sugar_grid.shape
+        self.space = GridPolars(
+            self, grid_dimensions, neighborhood_type="von_neumann", capacity=1
+        )
+        dim_0 = pl.Series("dim_0", pl.arange(grid_dimensions[0], eager=True)).to_frame()
+        dim_1 = pl.Series("dim_1", pl.arange(grid_dimensions[1], eager=True)).to_frame()
+        sugar_grid = dim_0.join(dim_1, how="cross").with_columns(
+            sugar=sugar_grid.flatten(), max_sugar=sugar_grid.flatten()
+        )
+        self.space.set_cells(sugar_grid)
+        self.agents += AntPolars(self, n_agents, initial_sugar, metabolism, vision)
+        self.space.place_to_empty(self.agents)
+
+    def run_model(self, steps: int) -> list[int]:
+        for _ in range(steps):
+            if len(self.agents) == 0:
+                return
+            self.step()
+            empty_cells = self.space.empty_cells
+            full_cells = self.space.full_cells
+
+            max_sugar = self.space.cells.join(
+                empty_cells, on=["dim_0", "dim_1"]
+            ).select(pl.col("max_sugar"))
+
+            self.space.set_cells(full_cells, {"sugar": 0})
+            self.space.set_cells(empty_cells, {"sugar": max_sugar})