diff --git a/examples/benchmark_plot_0.png b/examples/benchmark_plot_0.png deleted file mode 100644 index 3132939..0000000 Binary files a/examples/benchmark_plot_0.png and /dev/null differ diff --git a/examples/sugarscape_ig/benchmark_plot_0.png b/examples/sugarscape_ig/benchmark_plot_0.png new file mode 100644 index 0000000..92408c4 Binary files /dev/null and b/examples/sugarscape_ig/benchmark_plot_0.png differ diff --git a/examples/sugarscape_ig/benchmark_plot_1.png b/examples/sugarscape_ig/benchmark_plot_1.png new file mode 100644 index 0000000..198d50a Binary files /dev/null and b/examples/sugarscape_ig/benchmark_plot_1.png differ diff --git a/examples/sugarscape_ig/performance_comparison.py b/examples/sugarscape_ig/performance_comparison.py index c46e3c3..c6b2ad4 100644 --- a/examples/sugarscape_ig/performance_comparison.py +++ b/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__": diff --git a/examples/sugarscape_ig/ss_polars/__init__.py b/examples/sugarscape_ig/ss_polars/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/sugarscape_ig/ss_polars/agents.py b/examples/sugarscape_ig/ss_polars/agents.py new file mode 100644 index 0000000..6b3cba4 --- /dev/null +++ b/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)) diff --git a/examples/sugarscape_ig/ss_polars/model.py b/examples/sugarscape_ig/ss_polars/model.py new file mode 100644 index 0000000..70830e0 --- /dev/null +++ b/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})