diff --git a/examples/sugarscape_ig/performance_comparison.py b/examples/sugarscape_ig/performance_comparison.py index c6b2ad4..1a27bcb 100644 --- a/examples/sugarscape_ig/performance_comparison.py +++ b/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,10 +183,16 @@ 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 = [ + """labels_0 = [ "mesa", # "mesa-frames (pd concise)", "mesa-frames (pl concise)", @@ -81,16 +210,32 @@ def main(): # FLAME2-GPU 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__": diff --git a/examples/sugarscape_ig/polars_comparison.png b/examples/sugarscape_ig/polars_comparison.png new file mode 100644 index 0000000..235a52e Binary files /dev/null and b/examples/sugarscape_ig/polars_comparison.png differ diff --git a/examples/sugarscape_ig/ss_mesa/agents.py b/examples/sugarscape_ig/ss_mesa/agents.py index 797824f..7577203 100644 --- a/examples/sugarscape_ig/ss_mesa/agents.py +++ b/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 diff --git a/examples/sugarscape_ig/ss_mesa/model.py b/examples/sugarscape_ig/ss_mesa/model.py index 67336ed..076af8e 100644 --- a/examples/sugarscape_ig/ss_mesa/model.py +++ b/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 diff --git a/examples/sugarscape_ig/ss_polars/agents.py b/examples/sugarscape_ig/ss_polars/agents.py index 6b3cba4..78f2558 100644 --- a/examples/sugarscape_ig/ss_polars/agents.py +++ b/examples/sugarscape_ig/ss_polars/agents.py @@ -1,10 +1,13 @@ +from abc import abstractmethod + import numpy as np import polars as pl +from numba import b1, guvectorize, int32 from mesa_frames import AgentSetPolars, ModelDF -class AntPolars(AgentSetPolars): +class AntPolarsBase(AgentSetPolars): def __init__( self, model: ModelDF, @@ -25,23 +28,52 @@ def __init__( 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), + "sugar": initial_sugar, + "metabolism": metabolism, + "vision": vision, } ) self.add(agents) + 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)) + def move(self): + neighborhood = self._get_neighborhood() + agent_order = self._get_agent_order(neighborhood) + neighborhood = self._prepare_neighborhood(neighborhood, agent_order) + best_moves = self.get_best_moves(neighborhood) + self.space.move_agents(agent_order["agent_id_center"], best_moves) + + def _get_neighborhood(self) -> pl.DataFrame: + """Get the neighborhood of each agent, completed with the sugar of the cell and the agent_id of the center cell + + NOTE: This method should be unnecessary if get_neighborhood/get_neighbors return the agent_id of the center cell and the properties of the cells + + Returns + ------- + pl.DataFrame + Neighborhood DataFrame + """ 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, @@ -49,53 +81,129 @@ def move(self): right_on=["dim_0", "dim_1"], )["unique_id"] ) + return neighborhood + + def _get_agent_order(self, neighborhood: pl.DataFrame) -> pl.DataFrame: + """Get the order of agents based on the original order of agents + + Parameters + ---------- + neighborhood : pl.DataFrame + Neighborhood DataFrame + Returns + ------- + pl.DataFrame + DataFrame with 'agent_id_center' and 'agent_order' columns + """ # 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") + return ( + neighborhood.unique( + subset=["agent_id_center"], keep="first", maintain_order=True + ) + .with_row_count("agent_order") + .select(["agent_id_center", "agent_order"]) + ) + + def _prepare_neighborhood( + self, neighborhood: pl.DataFrame, agent_order: pl.DataFrame + ) -> pl.DataFrame: + """Prepare the neighborhood DataFrame to find the best moves + + Parameters + ---------- + neighborhood : pl.DataFrame + Neighborhood DataFrame + agent_order : pl.DataFrame + DataFrame with 'agent_id_center' and 'agent_order' columns + + Returns + ------- + pl.DataFrame + Prepared neighborhood DataFrame + """ neighborhood = neighborhood.join(agent_order, on="agent_id_center") + # Add blocking agent order 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", - ) + how="left", + ).rename({"agent_id": "blocking_agent_id"}) - # Filter impossible moves + # Filter only possible moves (agent is in his cell, blocking agent has moved before him or there is no blocking agent) neighborhood = neighborhood.filter( - pl.col("agent_order") >= pl.col("blocking_agent_order") + (pl.col("agent_order") >= pl.col("blocking_agent_order")) + | pl.col("blocking_agent_order").is_null() ) - # Sort cells by sugar and radius (nearest first) - neighborhood = neighborhood.sort(["sugar", "radius"], descending=[True, False]) + # Sort neighborhood by agent_order & max_sugar (max_sugar because we will check anyway if the cell is empty) + # However, we need to make sure that the current agent cell is ordered by current sugar (since it's 0 until agent hasn't moved) + neighborhood = neighborhood.with_columns( + max_sugar=pl.when(pl.col("blocking_agent_id") == pl.col("agent_id_center")) + .then(pl.lit(0)) + .otherwise(pl.col("max_sugar")) + ).sort( + ["agent_order", "max_sugar", "radius", "dim_0"], + descending=[False, True, False, False], + ) + return neighborhood + + def get_best_moves(self, neighborhood: pl.DataFrame) -> pl.DataFrame: + """Get the best moves for each agent + + Parameters + ---------- + neighborhood : pl.DataFrame + Neighborhood DataFrame + Returns + ------- + pl.DataFrame + DataFrame with the best moves for each agent + """ + raise NotImplementedError("Subclasses must implement this method") + + +class AntPolarsLoopDF(AntPolarsBase): + def get_best_moves(self, neighborhood: pl.DataFrame): 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 + # Check if there are previous agents that might make the same move (priority for the given move is > 1) + neighborhood = neighborhood.with_columns( + priority=pl.col("agent_order").cum_count().over(["dim_0", "dim_1"]) + ) + + # Get the best moves for each agent: + # 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") + .unique(subset=["dim_0", "dim_1"], keep="first", maintain_order=True) ) - # 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") + # - There isn't a higher priority agent that might make the same move + + condition = pl.col("blocking_agent_id").is_null() | ( + pl.col("blocking_agent_id") == pl.col("agent_id_center") ) if len(best_moves) > 0: - condition = condition | pl.col("agent_id").is_in( + condition = condition | pl.col("blocking_agent_id").is_in( best_moves["agent_id_center"] ) + + condition = condition & (pl.col("priority") == 1) + new_best_moves = new_best_moves.filter(condition) best_moves = pl.concat([best_moves, new_best_moves]) @@ -110,16 +218,184 @@ def move(self): 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"])) + # Check if there are previous agents that might make the same move (priority for the given move is > 1) + neighborhood = neighborhood.with_columns( + priority=pl.col("agent_order").cum_count().over(["dim_0", "dim_1"]) + ) + return best_moves.sort("agent_order").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"] + +class AntPolarsLoop(AntPolarsBase): + numba_target = None + + def get_best_moves(self, neighborhood: pl.DataFrame): + occupied_cells, free_cells, target_cells = self._prepare_cells(neighborhood) + best_moves_func = self._get_best_moves() + + processed_agents = np.zeros(len(self.agents), dtype=np.bool_) + + if self.numba_target is None: + # Non-vectorized case: we need to create and pass the best_moves array + map_batches_func = lambda df: best_moves_func( + occupied_cells, + free_cells, + target_cells, + df.struct.field("agent_order"), + df.struct.field("blocking_agent_order"), + processed_agents, + best_moves=np.full(len(self.agents), -1, dtype=np.int32), + ) + else: + # Vectorized case: Polars will create the output array (best_moves) automatically + map_batches_func = lambda df: best_moves_func( + occupied_cells.astype(np.int32), + free_cells.astype(np.bool_), + target_cells.astype(np.int32), + df.struct.field("agent_order"), + df.struct.field("blocking_agent_order"), + processed_agents.astype(np.bool_), + ) + + best_moves = ( + neighborhood.fill_null(-1) + .cast({"agent_order": pl.Int32, "blocking_agent_order": pl.Int32}) + .select( + pl.struct(["agent_order", "blocking_agent_order"]).map_batches( + map_batches_func + ) + ) + .with_columns( + dim_0=pl.col("agent_order") // self.space.dimensions[1], + dim_1=pl.col("agent_order") % self.space.dimensions[1], + ) + .drop("agent_order") ) + return best_moves - def step(self): - self.shuffle().do("move").do("eat") - self.discard(self.agents.filter(pl.col("sugar") <= 0)) + def _prepare_cells( + self, neighborhood: pl.DataFrame + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Get the occupied and free cells and the target cells for each agent, + based on the neighborhood DataFrame such that the arrays refer to a flattened version of the grid + + Parameters + ---------- + neighborhood : pl.DataFrame + Neighborhood DataFrame + + Returns + ------- + Tuple[np.ndarray, np.ndarray, np.ndarray] + occupied_cells, free_cells, target_cells + """ + occupied_cells = ( + neighborhood[["agent_id_center", "agent_order"]] + .unique() + .join(self.pos, left_on="agent_id_center", right_on="unique_id") + .with_columns( + flattened=(pl.col("dim_0") * self.space.dimensions[1] + pl.col("dim_1")) + ) + .sort("agent_order")["flattened"] + .to_numpy() + ) + free_cells = np.ones( + self.space.dimensions[0] * self.space.dimensions[1], dtype=np.bool_ + ) + free_cells[occupied_cells] = False + + target_cells = ( + neighborhood["dim_0"] * self.space.dimensions[1] + neighborhood["dim_1"] + ).to_numpy() + return occupied_cells, free_cells, target_cells + + def _get_best_moves(self): + raise NotImplementedError("Subclasses must implement this method") + + +class AntPolarsLoopNoVec(AntPolarsLoop): + # Non-vectorized case + def _get_best_moves(self): + def inner_get_best_moves( + occupied_cells: np.ndarray, + free_cells: np.ndarray, + target_cells: np.ndarray, + agent_id_center: np.ndarray, + blocking_agent: np.ndarray, + processed_agents: np.ndarray, + best_moves: np.ndarray, + ) -> np.ndarray: + for i, agent in enumerate(agent_id_center): + # If the agent has not moved yet + if not processed_agents[agent]: + # If the target cell is free + if free_cells[target_cells[i]] or blocking_agent[i] == agent: + best_moves[agent] = target_cells[i] + # Free current cell + free_cells[occupied_cells[agent]] = True + # Occupy target cell + free_cells[target_cells[i]] = False + processed_agents[agent] = True + return best_moves + + return inner_get_best_moves + + +class AntPolarsNumba(AntPolarsLoop): + # Vectorized case + def _get_best_moves(self): + @guvectorize( + [ + ( + int32[:], + b1[:], + int32[:], + int32[:], + int32[:], + b1[:], + int32[:], + ) + ], + "(n), (m), (p), (p), (p), (n)->(n)", + nopython=True, + target=self.numba_target, + # Writable inputs should be declared according to https://numba.pydata.org/numba-doc/dev/user/vectorize.html#overwriting-input-values + # In this case, there doesn't seem to be a difference. I will leave it commented for reference so that we can use CUDA target (which doesn't support writable_args) + # writable_args=( + # "free_cells", + # "processed_agents", + # ), + ) + def vectorized_get_best_moves( + occupied_cells, + free_cells, + target_cells, + agent_id_center, + blocking_agent, + processed_agents, + best_moves, + ): + for i, agent in enumerate(agent_id_center): + # If the agent has not moved yet + if not processed_agents[agent]: + # If the target cell is free + if free_cells[target_cells[i]] or blocking_agent[i] == agent: + best_moves[agent] = target_cells[i] + # Free current cell + free_cells[occupied_cells[agent]] = True + # Occupy target cell + free_cells[target_cells[i]] = False + processed_agents[agent] = True + + return vectorized_get_best_moves + + +class AntPolarsNumbaCPU(AntPolarsNumba): + numba_target = "cpu" + + +class AntPolarsNumbaParallel(AntPolarsNumba): + numba_target = "parallel" + + +class AntPolarsNumbaGPU(AntPolarsNumba): + numba_target = "cuda" diff --git a/examples/sugarscape_ig/ss_polars/model.py b/examples/sugarscape_ig/ss_polars/model.py index 70830e0..be9768c 100644 --- a/examples/sugarscape_ig/ss_polars/model.py +++ b/examples/sugarscape_ig/ss_polars/model.py @@ -3,21 +3,24 @@ from mesa_frames import GridPolars, ModelDF -from .agents import AntPolars +from .agents import AntPolarsBase class SugarscapePolars(ModelDF): def __init__( self, + agent_type: type[AntPolarsBase], 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, + initial_positions: pl.DataFrame | None = None, + seed: int | None = None, width: int | None = None, height: int | None = None, ): - super().__init__() + super().__init__(seed) if sugar_grid is None: sugar_grid = self.random.integers(0, 4, (width, height)) grid_dimensions = sugar_grid.shape @@ -30,20 +33,21 @@ def __init__( 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) + 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]: 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}) + self.step() diff --git a/pyproject.toml b/pyproject.toml index d7b94db..ae8a1c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,6 +88,7 @@ test = [ dev = [ "mesa_frames[test, docs]", "mesa", + "numba", ] [tool.hatch.envs.test]