diff --git a/mesa_frames/__init__.py b/mesa_frames/__init__.py index 4288c36..3cc21e0 100644 --- a/mesa_frames/__init__.py +++ b/mesa_frames/__init__.py @@ -1,11 +1,13 @@ from mesa_frames.concrete.agents import AgentsDF +from mesa_frames.concrete.model import ModelDF from mesa_frames.concrete.pandas.agentset import AgentSetPandas +from mesa_frames.concrete.pandas.space import GridPandas from mesa_frames.concrete.polars.agentset import AgentSetPolars -from mesa_frames.concrete.model import ModelDF __all__ = [ "AgentsDF", "AgentSetPandas", "AgentSetPolars", "ModelDF", + "GridPandas", ] diff --git a/mesa_frames/abstract/space.py b/mesa_frames/abstract/space.py index 9237617..6fe6ef0 100644 --- a/mesa_frames/abstract/space.py +++ b/mesa_frames/abstract/space.py @@ -4,11 +4,13 @@ from typing import TYPE_CHECKING, Literal from warnings import warn +import numpy as np import polars as pl from numpy.random import Generator -from typing_extensions import Self +from typing_extensions import Any, Self -from mesa_frames.abstract.agents import AgentContainer +from mesa_frames import AgentsDF +from mesa_frames.abstract.agents import AgentContainer, AgentSetDF from mesa_frames.abstract.mixin import CopyMixin, DataFrameMixin from mesa_frames.types_ import ( BoolSeries, @@ -21,6 +23,7 @@ GridCoordinate, GridCoordinates, IdsLike, + Series, SpaceCoordinate, SpaceCoordinates, ) @@ -89,6 +92,12 @@ class SpaceDF(CopyMixin, DataFrameMixin): _model: "ModelDF" _agents: DataFrame | GeoDataFrame # Stores the agents placed in the space + _center_col_names: list[ + str + ] # The column names of the center pos/agents in the neighbors/neighborhood method (eg. ['dim_0_center', 'dim_1_center', ...] in Grids, ['node_id_center', 'edge_id_center'] in Networks) + _pos_col_names: list[ + str + ] # The column names of the positions in the _agents dataframe (eg. ['dim_0', 'dim_1', ...] in Grids, ['node_id', 'edge_id'] in Networks) def __init__(self, model: "ModelDF") -> None: """Create a new SpaceDF object. @@ -103,6 +112,71 @@ def __init__(self, model: "ModelDF") -> None: """ self._model = model + def move_agents( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + pos: SpaceCoordinate | SpaceCoordinates, + inplace: bool = True, + ) -> Self: + """Move agents in the Space to the specified coordinates. If some agents are not placed, + raises a RuntimeWarning. + + Parameters + ---------- + agents : IdsLike | AgentContainer | Collection[AgentContainer] + The agents to move + pos : SpaceCoordinate | SpaceCoordinates + The coordinates for each agents. The length of the coordinates must match the number of agents. + inplace : bool, optional + Whether to perform the operation inplace, by default True + + Raises + ------ + RuntimeWarning + If some agents are not placed in the space. + ValueError + - If some agents are not part of the model. + - If agents is IdsLike and some agents are present multiple times. + + Returns + ------- + Self + """ + obj = self._get_obj(inplace=inplace) + return obj._place_or_move_agents(agents=agents, pos=pos, is_move=True) + + def place_agents( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + pos: SpaceCoordinate | SpaceCoordinates, + inplace: bool = True, + ) -> Self: + """Place agents in the space according to the specified coordinates. If some agents are already placed, raises a RuntimeWarning. + + Parameters + ---------- + agents : IdsLike | AgentContainer | Collection[AgentContainer] + The agents to place in the space + pos : SpaceCoordinate | SpaceCoordinates + The coordinates for each agents. The length of the coordinates must match the number of agents. + inplace : bool, optional + Whether to perform the operation inplace, by default True + + Returns + ------- + Self + + Raises + ------ + RuntimeWarning + If some agents are already placed in the space. + ValueError + - If some agents are not part of the model. + - If agents is IdsLike and some agents are present multiple times. + """ + obj = self._get_obj(inplace=inplace) + return obj._place_or_move_agents(agents=agents, pos=pos, is_move=False) + def random_agents( self, n: int, @@ -124,9 +198,58 @@ def random_agents( A DataFrame with the sampled agents """ if seed is None: - seed = self.random.integers(0) + seed = self.random.integers(np.iinfo(np.int32).max) return self._df_sample(self._agents, n=n, seed=seed) + def swap_agents( + self, + agents0: IdsLike | AgentContainer | Collection[AgentContainer], + agents1: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ) -> Self: + """Swap the positions of the agents in the space. + agents0 and agents1 must have the same length and all agents must be placed in the space. + + Parameters + ---------- + agents0 : IdsLike | AgentContainer | Collection[AgentContainer] + The first set of agents to swap + agents1 : IdsLike | AgentContainer | Collection[AgentContainer] + The second set of agents to swap + + Returns + ------- + Self + """ + agents0 = self._get_ids_srs(agents0) + agents1 = self._get_ids_srs(agents1) + if __debug__: + if len(agents0) != len(agents1): + raise ValueError("The two sets of agents must have the same length") + if not self._df_contains(self._agents, "agent_id", agents0).all(): + raise ValueError("Some agents in agents0 are not in the space") + if not self._df_contains(self._agents, "agent_id", agents1).all(): + raise ValueError("Some agents in agents1 are not in the space") + if self._srs_contains(agents0, agents1).any(): + raise ValueError("Some agents are present in both agents0 and agents1") + obj = self._get_obj(inplace) + agents0_df = obj._df_get_masked_df( + obj._agents, index_cols="agent_id", mask=agents0 + ) + agents1_df = obj._df_get_masked_df( + obj._agents, index_cols="agent_id", mask=agents1 + ) + agents0_df = obj._df_set_index(agents0_df, "agent_id", agents1) + agents1_df = obj._df_set_index(agents1_df, "agent_id", agents0) + obj._agents = obj._df_combine_first( + agents0_df, obj._agents, index_cols="agent_id" + ) + obj._agents = obj._df_combine_first( + agents1_df, obj._agents, index_cols="agent_id" + ) + + return obj + @abstractmethod def get_directions( self, @@ -168,7 +291,7 @@ def get_distances( pos1: SpaceCoordinate | SpaceCoordinates | None = None, agents0: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, - ) -> DataFrame: + ) -> Series: """Returns the distances from pos0 to pos1 or agents0 and agents1. If the space is a Network, the distance is the number of nodes of the shortest path between the two nodes. In all other cases, the distance is Euclidean/l2/Frobenius norm. @@ -228,32 +351,20 @@ def get_neighbors( ... @abstractmethod - def move_agents( + def move_to_empty( self, agents: IdsLike | AgentContainer | Collection[AgentContainer], - pos: SpaceCoordinate | SpaceCoordinates, inplace: bool = True, ) -> Self: - """Place agents in the space according to the specified coordinates. If some agents are already placed, - raises a RuntimeWarning. + """Move agents to empty cells/positions in the space (cells/positions where there isn't any single agent). Parameters ---------- agents : IdsLike | AgentContainer | Collection[AgentContainer] - The agents to place in the space - pos : SpaceCoordinate | SpaceCoordinates - The coordinates for each agents. The length of the coordinates must match the number of agents. + The agents to move to empty cells/positions inplace : bool, optional Whether to perform the operation inplace, by default True - Raises - ------ - RuntimeWarning - If some agents are already placed in the space. - ValueError - - If some agents are not part of the model. - - If agents is IdsLike and some agents are present multiple times. - Returns ------- Self @@ -261,17 +372,17 @@ def move_agents( ... @abstractmethod - def move_to_empty( + def place_to_empty( self, agents: IdsLike | AgentContainer | Collection[AgentContainer], inplace: bool = True, ) -> Self: - """Move agents to empty cells/positions in the space (cells/positions where there isn't any single agent). + """Place agents in empty cells/positions in the space (cells/positions where there isn't any single agent). Parameters ---------- agents : IdsLike | AgentContainer | Collection[AgentContainer] - The agents to move to empty cells/positions + The agents to place in empty cells/positions inplace : bool, optional Whether to perform the operation inplace, by default True @@ -330,21 +441,43 @@ def remove_agents( """ ... + def _get_ids_srs( + self, agents: IdsLike | AgentContainer | Collection[AgentContainer] + ) -> Series: + if isinstance(agents, AgentSetDF): + return self._srs_constructor(agents.index, name="agent_id") + elif isinstance(agents, AgentsDF): + return self._srs_constructor(agents._ids, name="agent_id") + elif isinstance(agents, Collection) and (isinstance(agents[0], AgentContainer)): + ids = [] + for a in agents: + if isinstance(a, AgentSetDF): + ids.append(self._srs_constructor(a.index, name="agent_id")) + elif isinstance(a, AgentsDF): + ids.append(self._srs_constructor(a._ids, name="agent_id")) + return self._df_concat(ids, ignore_index=True) + elif isinstance(agents, int): + return self._srs_constructor([agents], name="agent_id") + else: # IDsLike + return self._srs_constructor(agents, name="agent_id") + @abstractmethod - def swap_agents( + def _place_or_move_agents( self, - agents0: IdsLike | AgentContainer | Collection[AgentContainer], - agents1: IdsLike | AgentContainer | Collection[AgentContainer], + agents: IdsLike | AgentContainer | Collection[AgentContainer], + pos: SpaceCoordinate | SpaceCoordinates, + is_move: bool, ) -> Self: - """Swap the positions of the agents in the space. - agents0 and agents1 must have the same length and all agents must be placed in the space. + """A unique method for moving or placing agents (only the RuntimeWarning changes). Parameters ---------- - agents0 : IdsLike | AgentContainer | Collection[AgentContainer] - The first set of agents to swap - agents1 : IdsLike | AgentContainer | Collection[AgentContainer] - The second set of agents to swap + agents : IdsLike | AgentContainer | Collection[AgentContainer] + The agents to move/place + pos : SpaceCoordinate | SpaceCoordinates + The position to move/place agents to + is_move : bool + Whether the operation is "move" or "place" Returns ------- @@ -415,14 +548,12 @@ class DiscreteSpaceDF(SpaceDF): Set the properties of the specified cells. """ + _agents: DataFrame _capacity: int | None # The maximum capacity for cells (default is infinite) _cells: DataFrame # Stores the properties of the cells - _cells_col_names: list[ - str - ] # The column names of the _cells dataframe (eg. ['dim_0', 'dim_1', ...] in Grids, ['node_id', 'edge_id'] in Networks) - _center_col_names: list[ - str - ] # The column names of the center cells/agents in the get_neighbors method (eg. ['dim_0_center', 'dim_1_center', ...] in Grids, ['node_id_center', 'edge_id_center'] in Networks) + _cells_capacity: ( + DiscreteSpaceCapacity # Storing the remaining capacity of the cells in the grid + ) def __init__( self, @@ -456,12 +587,7 @@ def is_available(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFra DataFrame A dataframe with positions and a boolean column "available" """ - df = self._df_constructor(data=pos, columns=self._cells_col_names) - return self._df_add_columns( - df, - ["available"], - self._df_get_bool_mask(df, mask=self.full_cells, negate=True), - ) + return self._check_cells(pos, "available") def is_empty(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: """Check whether the input positions are empty (there isn't any single agent in the cells) @@ -476,10 +602,7 @@ def is_empty(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: DataFrame A dataframe with positions and a boolean column "empty" """ - df = self._df_constructor(data=pos, columns=self._cells_col_names) - return self._df_add_columns( - df, ["empty"], self._df_get_bool_mask(df, mask=self._cells, negate=True) - ) + return self._check_cells(pos, "empty") def is_full(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: """Check whether the input positions are full (there isn't any spot available in the cells) @@ -494,17 +617,17 @@ def is_full(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: DataFrame A dataframe with positions and a boolean column "full" """ - df = self._df_constructor(data=pos, columns=self._cells_col_names) - return self._df_add_columns( - df, ["full"], self._df_get_bool_mask(df, mask=self.full_cells, negate=True) - ) + return self._check_cells(pos, "full") def move_to_empty( self, agents: IdsLike | AgentContainer | Collection[AgentContainer], inplace: bool = True, ) -> Self: - return self._move_agents_to_cells(agents, cell_type="empty", inplace=inplace) + obj = self._get_obj(inplace) + return obj._place_or_move_agents_to_cells( + agents, cell_type="empty", is_move=True + ) def move_to_available( self, @@ -524,15 +647,41 @@ def move_to_available( ------- Self """ - return self._move_agents_to_cells( - agents, cell_type="available", inplace=inplace + obj = self._get_obj(inplace) + return obj._place_or_move_agents_to_cells( + agents, cell_type="available", is_move=True ) + def place_to_empty( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + return obj._place_or_move_agents_to_cells( + agents, cell_type="empty", is_move=False + ) + + def place_to_available( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + return obj._place_or_move_agents_to_cells( + agents, cell_type="available", is_move=False + ) + + def random_pos(self, n: int, seed: int | None = None) -> DataFrame | pl.DataFrame: + return self.sample_cells(n, cell_type="any", with_replacement=True, seed=seed) + def sample_cells( self, n: int, cell_type: Literal["any", "empty", "available", "full"] = "any", with_replacement: bool = True, + seed: int | None = None, + respect_capacity: bool = True, ) -> DataFrame: """Sample cells from the grid according to the specified cell_type. @@ -544,6 +693,12 @@ def sample_cells( The type of cells to sample, by default "any" with_replacement : bool, optional If the sampling should be with replacement, by default True + seed : int | None, optional + The seed for the sampling, by default None + If None, an integer from the model's random number generator is used. + respect_capacity : bool, optional + If the capacity of the cells should be respected in the sampling. + This is only relevant if cell_type is "empty" or "available", by default True Returns ------- @@ -559,7 +714,66 @@ def sample_cells( condition = self._available_cell_condition case "full": condition = self._full_cell_condition - return self._sample_cells(n, with_replacement, condition=condition) + return self._sample_cells( + n, + with_replacement, + condition=condition, + seed=seed, + respect_capacity=respect_capacity, + ) + + def set_cells( + self, + cells: DataFrame | DiscreteCoordinate | DiscreteCoordinates, + properties: DataFrame | dict[str, Any] | None = None, + inplace: bool = True, + ) -> Self: + """Set the properties of the specified cells. + This method mirrors the functionality of mesa's PropertyLayer, but allows also to set properties only of specific cells. + Either the cells DF must contain both the cells' coordinates and the properties + or the cells' coordinates can be specified separately with the cells argument. + If the Space is a Grid, the cell coordinates must be GridCoordinates. + If the Space is a Network, the cell coordinates must be NetworkCoordinates. + + + Parameters + ---------- + cells : DataFrame | DiscreteCoordinate | DiscreteCoordinates + The cells to set the properties for. It can contain the coordinates of the cells or both the coordinates and the properties. + properties : DataFrame | dict[str, Any] | None, optional + The properties of the cells, by default None if the cells argument contains the properties + inplace : bool + Whether to perform the operation inplace + + Returns + ------- + Self + """ + obj = self._get_obj(inplace) + cells_col_names = obj._df_column_names(obj._cells) + if __debug__: + if isinstance(cells, DataFrame) and any( + k not in cells_col_names for k in obj._pos_col_names + ): + raise ValueError( + f"The cells DataFrame must have the columns {obj._pos_col_names}" + ) + if properties: + pos_df = obj._get_df_coords(cells) + properties = obj._df_constructor(data=properties, index=pos_df.index) + cells = obj._df_concat( + [pos_df, properties], how="horizontal", index_cols=obj._pos_col_names + ) + else: + cells = obj._df_constructor(data=cells, index_cols=obj._pos_col_names) + + if "capacity" in cells_col_names: + obj._cells_capacity = obj._update_capacity_cells(cells) + + obj._cells = obj._df_combine_first( + cells, obj._cells, index_cols=obj._pos_col_names + ) + return obj @abstractmethod def get_neighborhood( @@ -610,76 +824,110 @@ def get_cells( """ ... + # We define the cell conditions here, because ruff does not allow lambda functions + + def _any_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return self._cells_capacity + @abstractmethod - def set_cells( - self, - properties: DataFrame, - cells: DiscreteCoordinates | None = None, - inplace: bool = True, - ) -> Self: - """Set the properties of the specified cells. - This method mirrors the functionality of mesa's PropertyLayer, but allows also to set properties only of specific cells. - Either the properties DF must contain both the cell coordinates and the properties - or the cell coordinates must be specified separately with the cells argument. - If the Space is a Grid, the cell coordinates must be GridCoordinates. - If the Space is a Network, the cell coordinates must be NetworkCoordinates. + def _empty_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: ... + def _available_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return cap > 0 + + def _full_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return cap == 0 + + def _check_cells( + self, + pos: DiscreteCoordinate | DiscreteCoordinates, + state: Literal["empty", "full", "available"], + ) -> DataFrame: + """ + Check the state of cells at given positions. Parameters ---------- - properties : DataFrame - The properties of the cells - cells : DiscreteCoordinates | None, optional - The coordinates of the cells to set the properties for, by default None (all cells) - inplace : bool - Whether to perform the operation inplace + pos : DiscreteCoordinate | DiscreteCoordinates + The positions to check + state : Literal["empty", "full", "available"] + The state to check for ("empty", "full", or "available") Returns ------- - Self + DataFrame + A dataframe with positions and a boolean column indicating the state """ - ... + pos_df = self._get_df_coords(pos) + + if state == "empty": + mask = self.empty_cells + elif state == "full": + mask = self.full_cells + elif state == "available": + mask = self.available_cells + + return self._df_with_columns( + original_df=pos_df, + data=self._df_get_bool_mask( + pos_df, + index_cols=self._pos_col_names, + mask=mask, + ), + new_columns=state, + ) - def _move_agents_to_cells( + def _place_or_move_agents_to_cells( self, agents: IdsLike | AgentContainer | Collection[AgentContainer], cell_type: Literal["any", "empty", "available"], - inplace: bool = True, + is_move: bool, ) -> Self: - obj = self._get_obj(inplace) - # Get Ids of agents - # TODO: fix this - if isinstance(agents, AgentContainer | Collection[AgentContainer]): - agents = agents.index + agents = self._get_ids_srs(agents) - # Check ids presence in model - b_contained = obj.model.agents.contains(agents) - if (isinstance(b_contained, pl.Series) and not b_contained.all()) or ( - isinstance(b_contained, bool) and not b_contained - ): - raise ValueError("Some agents are not in the model") + if __debug__: + # Check ids presence in model + b_contained = self.model.agents.contains(agents) + if (isinstance(b_contained, pl.Series) and not b_contained.all()) or ( + isinstance(b_contained, bool) and not b_contained + ): + raise ValueError("Some agents are not in the model") # Get cells of specified type - cells = obj.sample_cells(len(agents), cell_type=cell_type) + cells = self.sample_cells(len(agents), cell_type=cell_type) # Place agents - obj._agents = obj.move_agents(agents, cells) - return obj - - # We define the cell conditions here, because ruff does not allow lambda functions + if is_move: + self.move_agents(agents, cells) + else: + self.place_agents(agents, cells) + return self - def _any_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: - return True + @abstractmethod + def _get_df_coords( + self, + pos: DiscreteCoordinate | DiscreteCoordinates | None = None, + agents: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + ) -> DataFrame: + """Get the DataFrame of coordinates from the specified positions or agents. - def _empty_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: - return cap == self._capacity + Parameters + ---------- + pos : DiscreteCoordinate | DiscreteCoordinates | None, optional + agents : IdsLike | AgentContainer | Collection[AgentContainer], optional - def _available_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: - return cap > 0 + Returns + ------- + DataFrame + A dataframe where the columns are the coordinates col_names and the rows are the positions - def _full_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: - return cap == 0 + Raises + ------ + ValueError + If neither pos or agents are specified + """ + ... @abstractmethod def _sample_cells( @@ -687,17 +935,25 @@ def _sample_cells( n: int | None, with_replacement: bool, condition: Callable[[DiscreteSpaceCapacity], BoolSeries], + seed: int | None = None, + respect_capacity: bool = True, ) -> DataFrame: """Sample cells from the grid according to a condition on the capacity. Parameters ---------- n : int | None - The number of cells to sample + The number of cells to sample. If None, samples the maximum available. with_replacement : bool If the sampling should be with replacement condition : Callable[[DiscreteSpaceCapacity], BoolSeries] The condition to apply on the capacity + seed : int | None, optional + The seed for the sampling, by default None + If None, an integer from the model's random number generator is used. + respect_capacity : bool, optional + If the capacity should be respected in the sampling. + This is only relevant if cell_type is "empty" or "available", by default True Returns ------- @@ -705,19 +961,67 @@ def _sample_cells( """ ... - def __getitem__(self, cells: DiscreteCoordinates): - return self.get_cells(cells) + @abstractmethod + def _update_capacity_cells(self, cells: DataFrame) -> DiscreteSpaceCapacity: + """Update the cells' capacity after setting new properties. - def __setitem__(self, cells: DiscreteCoordinates, properties: DataFrame): - self.set_cells(properties=properties, cells=cells) + Parameters + ---------- + cells : DataFrame + A DF with the cells to update the capacity and the 'capacity' column + + Returns + ------- + DiscreteSpaceCapacity + The updated cells' capacity + """ + ... + + @abstractmethod + def _update_capacity_agents( + self, agents: DataFrame, operation: Literal["movement", "removal"] + ) -> DiscreteSpaceCapacity: + """Update the cells' capacity after moving agents. + + Parameters + ---------- + agents : DataFrame + The moved agents with their new positions + + Returns + ------- + DiscreteSpaceCapacity + The updated cells' capacity + """ + ... + + def __getitem__(self, cells: DiscreteCoordinate | DiscreteCoordinates): + return self.get_cells(cells) def __getattr__(self, key: str) -> DataFrame: # Fallback, if key (property) is not found in the object, # then it must mean that it's in the _cells dataframe return self._cells[key] + def __setitem__(self, cells: DiscreteCoordinates, properties: DataFrame): + self.set_cells(properties=properties, cells=cells) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}\nCells:\n{self._cells.__repr__()}\nAgents:\n{self._agents.__repr__()}" + + def __str__(self) -> str: + return ( + f"{self.__class__.__name__}\nCells:\n{self._cells}\nAgents:\n{self._agents}" + ) + @property def cells(self) -> DataFrame: + """ + Returns + ------- + DataFrame + A Dataframe with all cells, their properties and their agents + """ return self.get_cells() @cells.setter @@ -742,6 +1046,18 @@ def full_cells(self) -> DataFrame: None, with_replacement=False, condition=self._full_cell_condition ) + @property + @abstractmethod + def remaining_capacity(self) -> int | None: + """The remaining capacity of the cells in the grid. + + Returns + ------- + int | None + None if the capacity is infinite, otherwise the remaining capacity + """ + ... + class GridDF(DiscreteSpaceDF): """The GridDF class is an abstract class that defines the interface for all grid classes in mesa-frames. @@ -749,6 +1065,7 @@ class GridDF(DiscreteSpaceDF): Warning ------- + For rectangular grids: In this implementation, [0, ..., 0] is the bottom-left corner and [dimensions[0]-1, ..., dimensions[n-1]-1] is the top-right corner, consistent with Cartesian coordinates and Matplotlib/Seaborn plot outputs. @@ -756,6 +1073,12 @@ class GridDF(DiscreteSpaceDF): `mesa-examples Sugarscape model`_, where [0, ..., 0] is the top-left corner and [dimensions[0]-1, ..., dimensions[n-1]-1] is the bottom-right corner. + For hexagonal grids: + The coordinates are ordered according to the axial coordinate system. + In this system, the hexagonal grid uses two axes (q and r) at 60 degrees to each other. + The q-axis points to the right, and the r-axis points up and to the right. + The [0, 0] coordinate is at the bottom-left corner of the grid. + .. _np.genfromtxt: https://numpy.org/doc/stable/reference/generated/numpy.genfromtxt.html .. _mesa-examples Sugarscape model: https://github.com/projectmesa/mesa-examples/blob/e137a60e4e2f2546901bec497e79c4a7b0cc69bb/examples/sugarscape_g1mt/sugarscape_g1mt/model.py#L93-L94 @@ -777,7 +1100,7 @@ class GridDF(DiscreteSpaceDF): If the grid is a torus """ - _grid_capacity: ( + _cells_capacity: ( GridCapacity # Storing the remaining capacity of the cells in the grid ) _neighborhood_type: Literal[ @@ -798,6 +1121,7 @@ def __init__( Warning ------- + For rectangular grids: In this implementation, [0, ..., 0] is the bottom-left corner and [dimensions[0]-1, ..., dimensions[n-1]-1] is the top-right corner, consistent with Cartesian coordinates and Matplotlib/Seaborn plot outputs. @@ -805,15 +1129,21 @@ def __init__( `mesa-examples Sugarscape model`_, where [0, ..., 0] is the top-left corner and [dimensions[0]-1, ..., dimensions[n-1]-1] is the bottom-right corner. + For hexagonal grids: + The coordinates are ordered according to the axial coordinate system. + In this system, the hexagonal grid uses two axes (q and r) at 60 degrees to each other. + The q-axis points to the right, and the r-axis points up and to the right. + The [0, 0] coordinate is at the bottom-left corner of the grid. + .. _np.genfromtxt: https://numpy.org/doc/stable/reference/generated/numpy.genfromtxt.html .. _mesa-examples Sugarscape model: https://github.com/projectmesa/mesa-examples/blob/e137a60e4e2f2546901bec497e79c4a7b0cc69bb/examples/sugarscape_g1mt/sugarscape_g1mt/model.py#L93-L94 Parameters ---------- model : 'ModelDF' - The model selfect to which the grid belongs + The model object to which the grid belongs dimensions: Sequence[int] - The dimensions of the grid + The dimensions of the grid. For hexagonal grids, this should be [q_max, r_max]. torus : bool, optional If the grid should be a torus, by default False capacity : int | None, optional @@ -827,17 +1157,19 @@ def __init__( super().__init__(model, capacity) self._dimensions = dimensions self._torus = torus - self._cells_col_names = [f"dim_{k}" for k in range(len(dimensions))] - self._center_col_names = [x + "_center" for x in self._cells_col_names] + self._pos_col_names = [f"dim_{k}" for k in range(len(dimensions))] + self._center_col_names = [x + "_center" for x in self._pos_col_names] self._agents = self._df_constructor( - columns=["agent_id"] + self._cells_col_names, index_col="agent_id" + columns=["agent_id"] + self._pos_col_names, + index_cols="agent_id", + dtypes={col: int for col in self._pos_col_names}, ) self._cells = self._df_constructor( - columns=self._cells_col_names + ["capacity"], - index_cols=self._cells_col_names, + columns=self._pos_col_names + ["capacity"], + index_cols=self._pos_col_names, ) self._offsets = self._compute_offsets(neighborhood_type) - self._grid_capacity = self._generate_empty_grid(dimensions, capacity) + self._cells_capacity = self._generate_empty_grid(dimensions, capacity) self._neighborhood_type = neighborhood_type def get_directions( @@ -850,7 +1182,7 @@ def get_directions( ) -> DataFrame: result = self._calculate_differences(pos0, pos1, agents0, agents1) if normalize: - result = result / self._df_norm(result) + result = self._df_div(result, other=self._df_norm(result)) return result def get_distances( @@ -861,7 +1193,7 @@ def get_distances( agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, ) -> DataFrame: result = self._calculate_differences(pos0, pos1, agents0, agents1) - return self._df_norm(result) + return self._df_norm(result, "distance", True) def get_neighbors( self, @@ -875,55 +1207,202 @@ def get_neighbors( ) return self._df_get_masked_df( df=self._agents, - index_col="agent_id", + index_cols=self._pos_col_names, mask=neighborhood_df, - columns=self._agents.columns, ) - def get_cells( - self, coords: GridCoordinate | GridCoordinates | None = None + def get_neighborhood( + self, + radius: int | Sequence[int], + pos: GridCoordinate | GridCoordinates | None = None, + agents: IdsLike | AgentContainer | Collection[AgentContainer] = None, + include_center: bool = False, ) -> DataFrame: - coords_df = self._get_df_coords(pos=coords) - return self._df_get_masked_df( - df=self._cells, - index_cols=self._cells_col_names, - mask=coords_df, - columns=self._cells.columns, + pos_df = self._get_df_coords(pos, agents) + + if __debug__: + if isinstance(radius, Sequence): + if len(radius) != len(pos_df): + raise ValueError( + "The length of the radius sequence must be equal to the number of positions/agents" + ) + + ## Create all possible neighbors by multiplying offsets by the radius and adding original pos + + # If radius is a sequence, get the maximum radius (we will drop unnecessary neighbors later, time-efficient but memory-inefficient) + if isinstance(radius, Sequence): + radius_srs = self._srs_constructor(radius, name="radius") + max_radius = radius_srs.max() + else: + max_radius = radius + + range_srs = self._srs_range(name="radius", start=1, end=max_radius + 1) + + neighbors_df = self._df_join( + self._offsets, + range_srs, + how="cross", ) - def move_agents( - self, - agents: IdsLike | AgentContainer | Collection[AgentContainer], - pos: GridCoordinate | GridCoordinates, - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) + neighbors_df = self._df_with_columns( + neighbors_df, + data=self._df_mul( + neighbors_df[self._pos_col_names], neighbors_df["radius"] + ), + new_columns=self._pos_col_names, + ) - # Get Ids of agents - if isinstance(agents, AgentContainer | Collection[AgentContainer]): - agents = agents.index + if self.neighborhood_type == "hexagonal": + # We need to add in-between cells for hexagonal grids + # In-between offsets (for every radius k>=2, we need k-1 in-between cells) + in_between_cols = ["in_between_dim_0", "in_between_dim_1"] + radius_srs = self._srs_constructor( + np.repeat(np.arange(1, max_radius + 1), np.arange(0, max_radius)), + name="radius", + ) + radius_df = self._srs_to_df(radius_srs) + radius_df = self._df_with_columns( + radius_df, + self._df_groupby_cumcount(radius_df, "radius") + 1, + new_columns="offset", + ) - if __debug__: - # Check ids presence in model - b_contained = obj.model.agents.contains(agents) - if (isinstance(b_contained, pl.Series) and not b_contained.all()) or ( - isinstance(b_contained, bool) and not b_contained - ): - raise ValueError("Some agents are not in the model") + in_between_df = self._df_join( + self._in_between_offsets, + radius_df, + how="cross", + ) + # We multiply the radius to get the directional cells + in_between_df = self._df_with_columns( + in_between_df, + data=self._df_mul( + in_between_df[self._pos_col_names], in_between_df["radius"] + ), + new_columns=self._pos_col_names, + ) + # We multiply the offset (from the directional cells) to get the in-between offset for each radius + in_between_df = self._df_with_columns( + in_between_df, + data=self._df_mul( + in_between_df[in_between_cols], in_between_df["offset"] + ), + new_columns=in_between_cols, + ) + # We add the in-between offset to the directional cells to obtain the in-between cells + in_between_df = self._df_with_columns( + in_between_df, + data=self._df_add( + in_between_df[self._pos_col_names], + self._df_rename_columns( + in_between_df[in_between_cols], + in_between_cols, + self._pos_col_names, + ), + ), + new_columns=self._pos_col_names, + ) - # Check ids are unique - agents = pl.Series(agents) - if agents.unique_counts() != len(agents): - raise ValueError("Some agents are present multiple times") + in_between_df = self._df_drop_columns( + in_between_df, in_between_cols + ["offset"] + ) - # Warn if agents are already placed - if agents.is_in(obj._agents["agent_id"]): - warn("Some agents are already placed in the grid", RuntimeWarning) + neighbors_df = self._df_concat( + [neighbors_df, in_between_df], how="vertical" + ) - # Place agents (checking that capacity is not) - coords = obj._get_df_coords(pos) - obj._agents = obj._place_agents_df(agents, coords) - return obj + neighbors_df = self._df_join( + neighbors_df, pos_df, how="cross", suffix="_center" + ) + + center_df = self._df_rename_columns( + neighbors_df[self._center_col_names], + self._center_col_names, + self._pos_col_names, + ) # We rename the columns to the original names for the addition + + neighbors_df = self._df_with_columns( + original_df=neighbors_df, + new_columns=self._pos_col_names, + data=self._df_add( + neighbors_df[self._pos_col_names], + center_df, + ), + ) + + # If radius is a sequence, filter unnecessary neighbors + if isinstance(radius, Sequence): + radius_df = self._df_rename_columns( + self._df_concat([pos_df, radius_srs], how="horizontal"), + self._pos_col_names + ["radius"], + self._center_col_names + ["max_radius"], + ) + neighbors_df = self._df_join( + neighbors_df, + radius_df, + on=self._center_col_names, + ) + neighbors_df = self._df_filter( + neighbors_df, neighbors_df["radius"] <= neighbors_df["max_radius"] + ) + neighbors_df = self._df_drop_columns(neighbors_df, "max_radius") + + # If torus, "normalize" (take modulo) for out-of-bounds cells + if self._torus: + neighbors_df = self._df_with_columns( + neighbors_df, + data=self.torus_adj(neighbors_df[self._pos_col_names]), + new_columns=self._pos_col_names, + ) + # Remove duplicates + neighbors_df = self._df_drop_duplicates(neighbors_df, self._pos_col_names) + + # Filter out-of-bound neighbors + neighbors_df = self._df_filter( + neighbors_df, + ( + (neighbors_df[self._pos_col_names] < self._dimensions) + & (neighbors_df >= 0) + ), + all=True, + ) + + if include_center: + center_df = self._df_rename_columns( + pos_df, self._pos_col_names, self._center_col_names + ) + pos_df = self._df_with_columns( + pos_df, + data=0, + new_columns=["radius"], + ) + pos_df = self._df_concat([pos_df, center_df], how="horizontal") + + neighbors_df = self._df_concat( + [pos_df, neighbors_df], how="vertical", ignore_index=True + ) + + neighbors_df = self._df_reset_index(neighbors_df, drop=True) + return neighbors_df + + def get_cells( + self, coords: GridCoordinate | GridCoordinates | None = None + ) -> DataFrame: + # TODO : Consider whether not outputting the agents at all (fastest), + # outputting a single agent per cell (current) + # or outputting all agents per cell in a imploded list (slowest, https://stackoverflow.com/a/66018377) + if not coords: + cells_df = self._cells + else: + coords_df = self._get_df_coords(pos=coords) + cells_df = self._df_get_masked_df( + df=self._cells, index_cols=self._pos_col_names, mask=coords_df + ) + return self._df_join( + left=cells_df, + right=self._agents, + index_cols=self._pos_col_names, + on=self._pos_col_names, + ) def out_of_bounds(self, pos: GridCoordinate | GridCoordinates) -> DataFrame: """Check if a position is out of bounds in a non-toroidal grid. @@ -943,13 +1422,15 @@ def out_of_bounds(self, pos: GridCoordinate | GridCoordinates) -> DataFrame: ValueError If the grid is a torus """ - if self._torus: + if self.torus: raise ValueError("This method is only valid for non-torus grids") - pos_df = self._get_df_coords(pos) - out_of_bounds = pos_df < 0 | pos_df >= self._dimensions - return self._df_constructor( - data=[pos_df, out_of_bounds], + pos_df = self._get_df_coords(pos, check_bounds=False) + out_of_bounds = self._df_all( + (pos_df < 0) | (pos_df >= self._dimensions), + name="out_of_bounds", + index_cols=self._pos_col_names, ) + return self._df_concat(objs=[pos_df, out_of_bounds], how="horizontal") def remove_agents( self, @@ -958,9 +1439,7 @@ def remove_agents( ) -> Self: obj = self._get_obj(inplace) - # Get Ids of agents - if isinstance(agents, AgentContainer | Collection[AgentContainer]): - agents = agents.index + agents = obj._get_ids_srs(agents) if __debug__: # Check ids presence in model @@ -971,16 +1450,18 @@ def remove_agents( raise ValueError("Some agents are not in the model") # Remove agents - obj._agents = obj._df_remove(obj._agents, ids=agents, index_col="agent_id") + obj._cells_capacity = obj._update_capacity_agents(agents, operation="removal") + + obj._agents = obj._df_remove(obj._agents, mask=agents, index_cols="agent_id") return obj - def torus_adj(self, pos: GridCoordinates) -> DataFrame: + def torus_adj(self, pos: GridCoordinate | GridCoordinates) -> DataFrame: """Get the toroidal adjusted coordinates of a position. Parameters ---------- - pos : GridCoordinates + pos : GridCoordinate | GridCoordinates The coordinates to adjust Returns @@ -1060,24 +1541,38 @@ def _compute_offsets(self, neighborhood_type: str) -> DataFrame: raise ValueError( "Hexagonal neighborhood is only valid for 2-dimensional grids" ) - even_offsets = [(-1, -1), (-1, 0), (0, -1), (0, 1), (1, -1), (1, 0)] - odd_offsets = [(-1, 0), (-1, 1), (0, -1), (0, 1), (1, 0), (1, 1)] - - # Create a DataFrame with three columns: dim_0, dim_1, and is_even - offsets_data = [(d[0], d[1], True) for d in even_offsets] + [ - (d[0], d[1], False) for d in odd_offsets + directions = [ + (1, 0), # East + (1, -1), # South-West + (0, -1), # South-East + (-1, 0), # West + (-1, 1), # North-West + (0, 1), # North-East + ] + in_between = [ + (-1, -1), # East -> South-East + (0, 1), # South-West -> West + (-1, 0), # South-East -> South-West + (1, 1), # West -> North-West + (1, 0), # North-West -> North-East + (0, -1), # North-East -> East ] - return self._df_constructor( - data=offsets_data, columns=self._cells_col_names + ["is_even"] + df = self._df_constructor(data=directions, columns=self._pos_col_names) + self._in_between_offsets = self._df_with_columns( + df, + data=in_between, + new_columns=["in_between_dim_0", "in_between_dim_1"], ) + return df else: raise ValueError("Invalid neighborhood type specified") - return self._df_constructor(data=directions, columns=self._cells_col_names) + return self._df_constructor(data=directions, columns=self._pos_col_names) def _get_df_coords( self, pos: GridCoordinate | GridCoordinates | None = None, agents: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + check_bounds: bool = True, ) -> DataFrame: """Get the DataFrame of coordinates from the specified positions or agents. @@ -1101,12 +1596,47 @@ def _get_df_coords( raise ValueError("Neither pos or agents are specified") elif pos is not None and agents is not None: raise ValueError("Both pos and agents are specified") - if agents: - return self._df_get_masked_df( - self._agents, index_col="agent_id", mask=agents + # If the grid is non-toroidal, we have to check whether any position is out of bounds + if not self.torus and pos is not None and check_bounds: + pos = self.out_of_bounds(pos) + if pos["out_of_bounds"].any(): + raise ValueError( + "If the grid is non-toroidal, every position must be in-bound" + ) + if agents is not None: + agents = self._get_ids_srs(agents) + # Check ids presence in model + b_contained = self.model.agents.contains(agents) + if (isinstance(b_contained, pl.Series) and not b_contained.all()) or ( + isinstance(b_contained, bool) and not b_contained + ): + raise ValueError("Some agents are not present in the model") + + # Check ids presence in the grid + b_contained = self._df_contains(self._agents, "agent_id", agents) + if not b_contained.all(): + raise ValueError("Some agents are not placed in the grid") + # Check ids are unique + agents = pl.Series(agents) + if agents.n_unique() != len(agents): + raise ValueError("Some agents are present multiple times") + if agents is not None: + return self._df_reset_index( + self._df_get_masked_df( + self._agents, index_cols="agent_id", mask=agents + ), + index_cols="agent_id", + drop=True, ) if isinstance(pos, DataFrame): - return pos[self._cells_col_names] + return pos[self._pos_col_names] + elif ( + isinstance(pos, Collection) + and isinstance(pos[0], Collection) + and (len(pos[0]) == len(self._dimensions)) + ): # We only test the first coordinate for performance + # This means that we have a collection of coordinates + return self._df_constructor(data=pos, columns=self._pos_col_names) elif isinstance(pos, Sequence) and len(pos) == len(self._dimensions): # This means that the sequence is already a sequence where each element is the # sequence of coordinates for dimension i @@ -1115,23 +1645,66 @@ def _get_df_coords( start = c.start if c.start is not None else 0 step = c.step if c.step is not None else 1 stop = c.stop if c.stop is not None else self._dimensions[i] - pos[i] = pl.arange(start=start, end=stop, step=step) - elif isinstance(c, int): - pos[i] = [c] - return self._df_constructor(data=pos, columns=self._cells_col_names) - elif isinstance(pos, Collection) and all( - len(c) == len(self._dimensions) for c in pos - ): - # This means that we have a collection of coordinates - sequences = [] - for i in range(len(self._dimensions)): - sequences.append([c[i] for c in pos]) - return self._df_constructor(data=sequences, columns=self._cells_col_names) + pos[i] = self._srs_range(start=start, stop=stop, step=step) + return self._df_constructor(data=[pos], columns=self._pos_col_names) elif isinstance(pos, int) and len(self._dimensions) == 1: - return self._df_constructor(data=[pos], columns=self._cells_col_names) + return self._df_constructor(data=[pos], columns=self._pos_col_names) else: raise ValueError("Invalid coordinates") + def _place_or_move_agents( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + pos: GridCoordinate | GridCoordinates, + is_move: bool, + ) -> Self: + agents = self._get_ids_srs(agents) + + if __debug__: + # Warn if agents are already placed + if is_move: + if not self._df_contains(self._agents, "agent_id", agents).all(): + warn("Some agents are not present in the grid", RuntimeWarning) + else: # is "place" + if self._df_contains(self._agents, "agent_id", agents).any(): + warn("Some agents are already present in the grid", RuntimeWarning) + + # Check if agents are present in the model + b_contained = self.model.agents.contains(agents) + if not b_contained.all(): + raise ValueError("Some agents are not present in the model") + + # Check if there is enough capacity + if self._capacity: + # If len(agents) > remaining_capacity + len(agents that will move) + if len(agents) > self.remaining_capacity + len( + self._df_get_masked_df( + self._agents, + index_cols="agent_id", + mask=agents, + ) + ): + raise ValueError("Not enough capacity in the space for all agents") + + # Place or move agents (checking that capacity is respected) + pos_df = self._get_df_coords(pos) + agents_df = self._srs_to_df(agents) + + if __debug__: + if len(agents_df) != len(pos_df): + raise ValueError("The number of agents and positions must be equal") + + new_df = self._df_concat( + [agents_df, pos_df], how="horizontal", index_cols="agent_id" + ) + self._cells_capacity = self._update_capacity_agents( + new_df, operation="movement" + ) + self._agents = self._df_combine_first( + new_df, self._agents, index_cols="agent_id" + ) + return self + @abstractmethod def _generate_empty_grid( self, dimensions: Sequence[int], capacity: int @@ -1148,24 +1721,6 @@ def _generate_empty_grid( """ ... - @abstractmethod - def _place_agents_df(self, agents: IdsLike, coords: DataFrame) -> DataFrame: - """Place agents in the grid according to the specified coordinates. - - Parameters - ---------- - agents : IDsLike - The agents to place in the grid - coords : DataFrame - The coordinates for each agent - - Returns - ------- - DataFrame - A DataFrame with the agents placed in the grid - """ - ... - @property def dimensions(self) -> Sequence[int]: return self._dimensions diff --git a/mesa_frames/concrete/pandas/space.py b/mesa_frames/concrete/pandas/space.py new file mode 100644 index 0000000..1ac9898 --- /dev/null +++ b/mesa_frames/concrete/pandas/space.py @@ -0,0 +1,171 @@ +from collections.abc import Callable, Sequence + +import numpy as np +import pandas as pd +from typing import Literal + +from mesa_frames.abstract.space import GridDF +from mesa_frames.concrete.pandas.mixin import PandasMixin + + +class GridPandas(GridDF, PandasMixin): + _agents: pd.DataFrame + _copy_with_method: dict[str, tuple[str, list[str]]] = { + "_agents": ("copy", ["deep"]), + "_cells": ("copy", ["deep"]), + "_cells_capacity": ("copy", []), + "_offsets": ("copy", ["deep"]), + } + _cells: pd.DataFrame + _cells_capacity: np.ndarray + _offsets: pd.DataFrame + + def _empty_cell_condition(self, cap: np.ndarray) -> np.ndarray: + # Create a boolean mask of the same shape as cap + empty_mask = np.ones_like(cap, dtype=bool) + + if not self._agents.empty: + # Get the coordinates of all agents + agent_coords = self._agents[self._pos_col_names].to_numpy(int) + + # Mark cells containing agents as not empty + empty_mask[tuple(agent_coords.T)] = False + + return empty_mask + + def _generate_empty_grid( + self, dimensions: Sequence[int], capacity: int + ) -> np.ndarray: + if not capacity: + capacity = np.inf + return np.full(dimensions, capacity) + + def _sample_cells( + self, + n: int | None, + with_replacement: bool, + condition: Callable[[np.ndarray], np.ndarray], + seed: int | None = None, + respect_capacity: bool = True, + ) -> pd.DataFrame: + # Set up the random number generator + if seed is None: + rng = self.model.random + else: + rng = np.random.default_rng(seed) + + # Get the coordinates of cells that meet the condition + coords = np.array(np.where(condition(self._cells_capacity))).T + + if respect_capacity and condition != self._full_cell_condition: + capacities = self._cells_capacity[tuple(coords.T)] + else: + # If not respecting capacity or for full cells, set capacities to 1 + capacities = np.ones(len(coords), dtype=int) + + if n is not None: + if with_replacement: + if respect_capacity and condition != self._full_cell_condition: + assert ( + n <= capacities.sum() + ), "Requested sample size exceeds the total available capacity." + + sampled_coords = np.empty((0, coords.shape[1]), dtype=coords.dtype) + while len(sampled_coords) < n: + remaining_samples = n - len(sampled_coords) + sampled_indices = rng.choice( + len(coords), + size=remaining_samples, + replace=True, + ) + unique_indices, counts = np.unique( + sampled_indices, return_counts=True + ) + + if respect_capacity and condition != self._full_cell_condition: + # Calculate valid counts for each unique index + valid_counts = np.minimum(counts, capacities[unique_indices]) + # Update capacities + capacities[unique_indices] -= valid_counts + else: + valid_counts = counts + + # Create array of repeated coordinates + new_coords = np.repeat(coords[unique_indices], valid_counts, axis=0) + # Extend sampled_coords + sampled_coords = np.vstack((sampled_coords, new_coords)) + + if respect_capacity and condition != self._full_cell_condition: + # Update coords and capacities + mask = capacities > 0 + coords = coords[mask] + capacities = capacities[mask] + + sampled_coords = sampled_coords[:n] + rng.shuffle(sampled_coords) + else: + assert n <= len( + coords + ), "Requested sample size exceeds the number of available cells." + sampled_indices = rng.choice(len(coords), size=n, replace=False) + sampled_coords = coords[sampled_indices] + else: + sampled_coords = coords + + # Convert the coordinates to a DataFrame + sampled_cells = pd.DataFrame(sampled_coords, columns=self._pos_col_names) + return sampled_cells + + def _update_capacity_agents( + self, + agents: pd.DataFrame, + operation: Literal["movement", "removal"], + ) -> np.ndarray: + # Update capacity for agents that were already on the grid + masked_df = self._df_get_masked_df( + self._agents, index_cols="agent_id", mask=agents + ) + + if operation == "movement": + # Increase capacity at old positions + old_positions = tuple(masked_df[self._pos_col_names].to_numpy(int).T) + np.add.at(self._cells_capacity, old_positions, 1) + + # Decrease capacity at new positions + new_positions = tuple(agents[self._pos_col_names].to_numpy(int).T) + np.add.at(self._cells_capacity, new_positions, -1) + elif operation == "removal": + # Increase capacity at the positions of removed agents + positions = tuple(masked_df[self._pos_col_names].to_numpy(int).T) + np.add.at(self._cells_capacity, positions, 1) + return self._cells_capacity + + def _update_capacity_cells(self, cells: pd.DataFrame) -> np.ndarray: + # Get the coordinates of the cells to update + coords = cells.index + + # Get the current capacity of updatable cells + current_capacity = self._cells.reindex(coords, fill_value=self._capacity)[ + "capacity" + ].to_numpy() + + # Calculate the number of agents currently in each cell + agents_in_cells = current_capacity - self._cells_capacity[tuple(zip(*coords))] + + # Update the capacity in self._cells_capacity + new_capacity = cells["capacity"].to_numpy() - agents_in_cells + + # Assert that no new capacity is negative + assert np.all( + new_capacity >= 0 + ), "New capacity of a cell cannot be less than the number of agents in it." + + self._cells_capacity[tuple(zip(*coords))] = new_capacity + + return self._cells_capacity + + @property + def remaining_capacity(self) -> int: + if not self._capacity: + return np.inf + return self._cells_capacity.sum() diff --git a/tests/pandas/__init__.py b/tests/pandas/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/pandas/test_grid_pandas.py b/tests/pandas/test_grid_pandas.py new file mode 100644 index 0000000..6df0841 --- /dev/null +++ b/tests/pandas/test_grid_pandas.py @@ -0,0 +1,1279 @@ +import numpy as np +import pandas as pd +import pytest +import typeguard as tg + +from mesa_frames import GridPandas, ModelDF +from tests.pandas.test_agentset_pandas import ( + ExampleAgentSetPandas, + fix1_AgentSetPandas, +) +from tests.polars.test_agentset_polars import ( + ExampleAgentSetPolars, + fix2_AgentSetPolars, +) + + +# This serves otherwise ruff complains about the two fixtures not being used +def not_called(): + fix1_AgentSetPandas() + fix2_AgentSetPolars() + + +@tg.typechecked +class TestGridPandas: + @pytest.fixture + def model( + self, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ) -> ModelDF: + model = ModelDF() + model.agents.add([fix1_AgentSetPandas, fix2_AgentSetPolars]) + return model + + @pytest.fixture + def grid_moore(self, model: ModelDF) -> GridPandas: + space = GridPandas(model, dimensions=[3, 3], capacity=2) + space.place_agents(agents=[0, 1], pos=[[0, 0], [1, 1]]) + space.set_cells( + [[0, 0], [1, 1]], properties={"capacity": [1, 3], "property_0": "value_0"} + ) + return space + + @pytest.fixture + def grid_moore_torus(self, model: ModelDF) -> GridPandas: + space = GridPandas(model, dimensions=[3, 3], capacity=2, torus=True) + space.place_agents(agents=[0, 1], pos=[[0, 0], [1, 1]]) + space.set_cells( + [[0, 0], [1, 1]], properties={"capacity": [1, 3], "property_0": "value_0"} + ) + return space + + @pytest.fixture + def grid_von_neumann(self, model: ModelDF) -> GridPandas: + space = GridPandas(model, dimensions=[3, 3], neighborhood_type="von_neumann") + space.place_agents(agents=[0, 1], pos=[[0, 0], [1, 1]]) + return space + + @pytest.fixture + def grid_hexagonal(self, model: ModelDF) -> GridPandas: + space = GridPandas(model, dimensions=[10, 10], neighborhood_type="hexagonal") + space.place_agents(agents=[0, 1], pos=[[5, 4], [5, 5]]) + return space + + def test___init__(self, model: ModelDF): + # Test with default parameters + grid1 = GridPandas(model, dimensions=[3, 3]) + assert isinstance(grid1, GridPandas) + assert isinstance(grid1.agents, pd.DataFrame) + assert grid1.agents.empty + assert isinstance(grid1.cells, pd.DataFrame) + assert grid1.cells.empty + assert isinstance(grid1.dimensions, list) + assert len(grid1.dimensions) == 2 + assert isinstance(grid1.neighborhood_type, str) + assert grid1.neighborhood_type == "moore" + assert grid1.remaining_capacity == float("inf") + assert grid1.model == model + + # Test with capacity = 10 + grid2 = GridPandas(model, dimensions=[3, 3], capacity=10) + assert grid2.remaining_capacity == (10 * 3 * 3) + + # Test with torus = True + grid3 = GridPandas(model, dimensions=[3, 3], torus=True) + assert grid3.torus + + # Test with neighborhood_type = "von_neumann" + grid4 = GridPandas(model, dimensions=[3, 3], neighborhood_type="von_neumann") + assert grid4.neighborhood_type == "von_neumann" + + # Test with neighborhood_type = "moore" + grid5 = GridPandas(model, dimensions=[3, 3], neighborhood_type="moore") + assert grid5.neighborhood_type == "moore" + + # Test with neighborhood_type = "hexagonal" + grid6 = GridPandas(model, dimensions=[3, 3], neighborhood_type="hexagonal") + assert grid6.neighborhood_type == "hexagonal" + + def test_get_cells(self, grid_moore: GridPandas): + # Test with None (all cells) + result = grid_moore.get_cells() + assert isinstance(result, pd.DataFrame) + assert result.reset_index()["dim_0"].tolist() == [0, 1] + assert result.reset_index()["dim_1"].tolist() == [0, 1] + assert result["capacity"].tolist() == [1, 3] + assert result["property_0"].tolist() == ["value_0", "value_0"] + + # Test with GridCoordinate + result = grid_moore.get_cells([0, 0]) + assert isinstance(result, pd.DataFrame) + assert result.reset_index()["dim_0"].tolist() == [0] + assert result.reset_index()["dim_1"].tolist() == [0] + assert result["capacity"].tolist() == [1] + assert result["property_0"].tolist() == ["value_0"] + + # Test with GridCoordinates + result = grid_moore.get_cells([[0, 0], [1, 1]]) + assert isinstance(result, pd.DataFrame) + assert result.reset_index()["dim_0"].tolist() == [0, 1] + assert result.reset_index()["dim_1"].tolist() == [0, 1] + assert result["capacity"].tolist() == [1, 3] + assert result["property_0"].tolist() == ["value_0", "value_0"] + + def test_get_directions( + self, + grid_moore: GridPandas, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ): + # Test with GridCoordinate + dir = grid_moore.get_directions(pos0=[1, 1], pos1=[2, 2]) + assert isinstance(dir, pd.DataFrame) + assert dir["dim_0"].to_list() == [1] + assert dir["dim_1"].to_list() == [1] + + # Test with GridCoordinates + dir = grid_moore.get_directions(pos0=[[0, 0], [2, 2]], pos1=[[1, 2], [1, 1]]) + assert isinstance(dir, pd.DataFrame) + assert dir["dim_0"].to_list() == [1, -1] + assert dir["dim_1"].to_list() == [2, -1] + + # Test with missing agents (raises ValueError) + with pytest.raises(ValueError): + grid_moore.get_directions( + agents0=fix1_AgentSetPandas, agents1=fix2_AgentSetPolars + ) + + # Test with IdsLike + grid_moore.place_agents(fix2_AgentSetPolars, [[0, 1], [0, 2], [1, 0], [1, 2]]) + dir = grid_moore.get_directions(agents0=[0, 1], agents1=[4, 5]) + assert isinstance(dir, pd.DataFrame) + assert dir["dim_0"].to_list() == [0, -1] + assert dir["dim_1"].to_list() == [1, 1] + + # Test with two AgentSetDFs + grid_moore.place_agents([2, 3], [[1, 1], [2, 2]]) + dir = grid_moore.get_directions( + agents0=fix1_AgentSetPandas, agents1=fix2_AgentSetPolars + ) + assert isinstance(dir, pd.DataFrame) + assert dir["dim_0"].to_list() == [0, -1, 0, -1] + assert dir["dim_1"].to_list() == [1, 1, -1, 0] + + # Test with AgentsDF + dir = grid_moore.get_directions( + agents0=grid_moore.model.agents, agents1=grid_moore.model.agents + ) + assert isinstance(dir, pd.DataFrame) + assert (dir == 0).all().all() + + # Test with normalize + dir = grid_moore.get_directions(agents0=[0, 1], agents1=[4, 5], normalize=True) + # Check if the vectors are normalized (length should be 1) + assert np.allclose(np.sqrt(dir["dim_0"] ** 2 + dir["dim_1"] ** 2), 1.0) + # Check specific normalized values + assert np.allclose(dir["dim_0"].to_list(), [0, -1 / np.sqrt(2)]) + assert np.allclose(dir["dim_1"].to_list(), [1, 1 / np.sqrt(2)]) + + def test_get_distances( + self, + grid_moore: GridPandas, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ): + # Test with GridCoordinate + dist = grid_moore.get_distances(pos0=[1, 1], pos1=[2, 2]) + assert isinstance(dist, pd.DataFrame) + assert np.allclose(dist["distance"].to_list(), [np.sqrt(2)]) + + # Test with GridCoordinates + dist = grid_moore.get_distances(pos0=[[0, 0], [2, 2]], pos1=[[1, 2], [1, 1]]) + assert isinstance(dist, pd.DataFrame) + assert np.allclose(dist["distance"].to_list(), [np.sqrt(5), np.sqrt(2)]) + + # Test with missing agents (raises ValueError) + with pytest.raises(ValueError): + grid_moore.get_distances( + agents0=fix1_AgentSetPandas, agents1=fix2_AgentSetPolars + ) + + # Test with IdsLike + grid_moore.place_agents(fix2_AgentSetPolars, [[0, 1], [0, 2], [1, 0], [1, 2]]) + dist = grid_moore.get_distances(agents0=[0, 1], agents1=[4, 5]) + assert isinstance(dist, pd.DataFrame) + assert np.allclose(dist["distance"].to_list(), [1.0, np.sqrt(2)]) + + # Test with two AgentSetDFs + grid_moore.place_agents([2, 3], [[1, 1], [2, 2]]) + dist = grid_moore.get_distances( + agents0=fix1_AgentSetPandas, agents1=fix2_AgentSetPolars + ) + assert isinstance(dist, pd.DataFrame) + assert np.allclose(dist["distance"].to_list(), [1.0, np.sqrt(2), 1.0, 1.0]) + + # Test with AgentsDF + dist = grid_moore.get_distances( + agents0=grid_moore.model.agents, agents1=grid_moore.model.agents + ) + assert (dist == 0).all().all() + + def test_get_neighborhood( + self, + grid_moore: GridPandas, + grid_hexagonal: GridPandas, + grid_von_neumann: GridPandas, + grid_moore_torus: GridPandas, + ): + # Test with radius = int, pos=GridCoordinate + neighborhood = grid_moore.get_neighborhood(radius=1, pos=[1, 1]) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.columns.to_list() == [ + "dim_0", + "dim_1", + "radius", + "dim_0_center", + "dim_1_center", + ] + assert neighborhood.shape == (8, 5) + assert neighborhood["dim_0"].to_list() == [0, 0, 0, 1, 1, 2, 2, 2] + assert neighborhood["dim_1"].to_list() == [0, 1, 2, 0, 2, 0, 1, 2] + assert neighborhood["radius"].to_list() == [1] * 8 + assert neighborhood["dim_0_center"].to_list() == [1] * 8 + assert neighborhood["dim_1_center"].to_list() == [1] * 8 + + # Test with Sequence[int], pos=Sequence[GridCoordinate] + neighborhood = grid_moore.get_neighborhood(radius=[1, 2], pos=[[1, 1], [2, 2]]) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == (8 + 6, 5) + assert neighborhood["radius"].sort_values().to_list() == [1] * 11 + [2] * 3 + assert neighborhood["dim_0_center"].sort_values().to_list() == [1] * 8 + [2] * 6 + assert neighborhood["dim_1_center"].sort_values().to_list() == [1] * 8 + [2] * 6 + neighborhood = neighborhood.sort_values(["dim_0", "dim_1"]) + assert neighborhood["dim_0"].to_list() == [0] * 5 + [1] * 4 + [2] * 5 + assert neighborhood["dim_1"].to_list() == [ + 0, + 0, + 1, + 2, + 2, + 0, + 1, + 2, + 2, + 0, + 0, + 1, + 1, + 2, + ] + + grid_moore.place_agents([0, 1], [[1, 1], [2, 2]]) + + # Test with agent=int, pos=GridCoordinate + neighborhood = grid_moore.get_neighborhood(radius=1, agents=0) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == (8, 5) + assert neighborhood["dim_0"].to_list() == [0, 0, 0, 1, 1, 2, 2, 2] + assert neighborhood["dim_1"].to_list() == [0, 1, 2, 0, 2, 0, 1, 2] + assert neighborhood["radius"].to_list() == [1] * 8 + assert neighborhood["dim_0_center"].to_list() == [1] * 8 + assert neighborhood["dim_1_center"].to_list() == [1] * 8 + + # Test with agent=Sequence[int], pos=Sequence[GridCoordinate] + neighborhood = grid_moore.get_neighborhood(radius=[1, 2], agents=[0, 1]) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == (8 + 6, 5) + assert neighborhood["radius"].sort_values().to_list() == [1] * 11 + [2] * 3 + assert neighborhood["dim_0_center"].sort_values().to_list() == [1] * 8 + [2] * 6 + assert neighborhood["dim_1_center"].sort_values().to_list() == [1] * 8 + [2] * 6 + neighborhood = neighborhood.sort_values(["dim_0", "dim_1"]) + assert neighborhood["dim_0"].to_list() == [0] * 5 + [1] * 4 + [2] * 5 + assert neighborhood["dim_1"].to_list() == [ + 0, + 0, + 1, + 2, + 2, + 0, + 1, + 2, + 2, + 0, + 0, + 1, + 1, + 2, + ] + + # Test with include_center + neighborhood = grid_moore.get_neighborhood( + radius=1, pos=[1, 1], include_center=True + ) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == (9, 5) + assert neighborhood["dim_0"].to_list() == [1, 0, 0, 0, 1, 1, 2, 2, 2] + assert neighborhood["dim_1"].to_list() == [1, 0, 1, 2, 0, 2, 0, 1, 2] + assert neighborhood["radius"].to_list() == [0] + [1] * 8 + assert neighborhood["dim_0_center"].to_list() == [1] * 9 + assert neighborhood["dim_1_center"].to_list() == [1] * 9 + + # Test with torus + neighborhood = grid_moore_torus.get_neighborhood(radius=1, pos=[0, 0]) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == (8, 5) + assert neighborhood["dim_0"].to_list() == [2, 2, 2, 0, 0, 1, 1, 1] + assert neighborhood["dim_1"].to_list() == [2, 0, 1, 2, 1, 2, 0, 1] + assert neighborhood["radius"].to_list() == [1] * 8 + assert neighborhood["dim_0_center"].to_list() == [0] * 8 + assert neighborhood["dim_1_center"].to_list() == [0] * 8 + + # Test with radius and pos of different length + with pytest.raises(ValueError): + neighborhood = grid_moore.get_neighborhood(radius=[1, 2], pos=[1, 1]) + + # Test with von_neumann neighborhood + neighborhood = grid_von_neumann.get_neighborhood(radius=1, pos=[1, 1]) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == (4, 5) + assert neighborhood["dim_0"].to_list() == [0, 1, 1, 2] + assert neighborhood["dim_1"].to_list() == [1, 0, 2, 1] + assert neighborhood["radius"].to_list() == [1] * 4 + assert neighborhood["dim_0_center"].to_list() == [1] * 4 + assert neighborhood["dim_1_center"].to_list() == [1] * 4 + + # Test with hexagonal neighborhood (odd cell [2,1] and even cell [2,2]) + neighborhood = grid_hexagonal.get_neighborhood( + radius=[2, 3], pos=[[5, 4], [5, 5]] + ) + assert isinstance(neighborhood, pd.DataFrame) + assert neighborhood.shape == ( + 6 * 2 + 12 * 2 + 18, + 5, + ) # 6 neighbors for radius 1, 12 for radius 2, 18 for radius 3 + + # Sort the neighborhood for consistent ordering + neighborhood = neighborhood.sort_values( + ["dim_0_center", "dim_1_center", "radius", "dim_0", "dim_1"] + ).reset_index(drop=True) + + # Expected neighbors for [5,4] and [5,5] + expected_neighbors = [ + # Neighbors of [5,4] + # radius 1 + (4, 4), + (4, 5), + (5, 3), + (5, 5), + (6, 3), + (6, 4), + # radius 2 + (3, 4), + (3, 6), + (4, 2), + (4, 5), + (4, 6), + (5, 2), + (5, 5), + (5, 6), + (6, 3), + (7, 2), + (7, 3), + (7, 4), + # Neighbors of [5,5] + # radius 1 + (4, 5), + (4, 6), + (5, 4), + (5, 6), + (6, 4), + (6, 5), + # radius 2 + (3, 5), + (3, 7), + (4, 3), + (4, 6), + (4, 7), + (5, 3), + (5, 6), + (5, 7), + (6, 4), + (7, 3), + (7, 4), + (7, 5), + # radius 3 + (2, 5), + (2, 8), + (3, 2), + (3, 6), + (3, 8), + (4, 2), + (4, 7), + (4, 8), + (5, 2), + (5, 6), + (5, 7), + (5, 8), + (6, 3), + (7, 4), + (8, 2), + (8, 3), + (8, 4), + (8, 5), + ] + + assert ( + list(zip(neighborhood["dim_0"], neighborhood["dim_1"])) + == expected_neighbors + ) + + def test_get_neighbors( + self, + fix2_AgentSetPolars: ExampleAgentSetPolars, + grid_moore: GridPandas, + grid_hexagonal: GridPandas, + grid_von_neumann: GridPandas, + grid_moore_torus: GridPandas, + ): + # Place agents in the grid + grid_moore.move_agents( + [0, 1, 2, 3, 4, 5, 6, 7], + [[0, 0], [0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1], [2, 2]], + ) + + # Test with radius = int, pos=GridCoordinate + neighbors = grid_moore.get_neighbors(radius=1, pos=[1, 1]) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.columns.to_list() == ["dim_0", "dim_1"] + assert neighbors.shape == (8, 2) + assert neighbors["dim_0"].to_list() == [0, 0, 0, 1, 1, 2, 2, 2] + assert neighbors["dim_1"].to_list() == [0, 1, 2, 0, 2, 0, 1, 2] + assert set(neighbors.index) == {0, 1, 2, 3, 4, 5, 6, 7} + + # Test with Sequence[int], pos=Sequence[GridCoordinate] + neighbors = grid_moore.get_neighbors(radius=[1, 2], pos=[[1, 1], [2, 2]]) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (8, 2) + neighbors = neighbors.sort_values(["dim_0", "dim_1"]) + assert neighbors["dim_0"].to_list() == [0, 0, 0, 1, 1, 2, 2, 2] + assert neighbors["dim_1"].to_list() == [0, 1, 2, 0, 2, 0, 1, 2] + assert set(neighbors.index) == {0, 1, 2, 3, 4, 5, 6, 7} + + # Test with agent=int + neighbors = grid_moore.get_neighbors(radius=1, agents=0) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (2, 2) + assert neighbors["dim_0"].to_list() == [0, 1] + assert neighbors["dim_1"].to_list() == [1, 0] + assert set(neighbors.index) == {1, 3} + + # Test with agent=Sequence[int] + neighbors = grid_moore.get_neighbors(radius=[1, 2], agents=[0, 7]) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (7, 2) + neighbors = neighbors.sort_values(["dim_0", "dim_1"]) + assert neighbors["dim_0"].to_list() == [0, 0, 0, 1, 1, 2, 2] + assert neighbors["dim_1"].to_list() == [0, 1, 2, 0, 2, 0, 1] + assert set(neighbors.index) == {0, 1, 2, 3, 4, 5, 6} + + # Test with include_center + neighbors = grid_moore.get_neighbors(radius=1, pos=[1, 1], include_center=True) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (8, 2) # No agent at [1, 1], so still 8 neighbors + assert neighbors["dim_0"].to_list() == [0, 0, 0, 1, 1, 2, 2, 2] + assert neighbors["dim_1"].to_list() == [0, 1, 2, 0, 2, 0, 1, 2] + assert set(neighbors.index) == {0, 1, 2, 3, 4, 5, 6, 7} + + # Test with torus + grid_moore_torus.move_agents( + [0, 1, 2, 3, 4, 5, 6, 7], + [[2, 2], [2, 0], [2, 1], [0, 2], [0, 1], [1, 2], [1, 0], [1, 1]], + ) + neighbors = grid_moore_torus.get_neighbors(radius=1, pos=[0, 0]) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (8, 2) + assert neighbors["dim_0"].to_list() == [2, 2, 2, 0, 0, 1, 1, 1] + assert neighbors["dim_1"].to_list() == [2, 0, 1, 2, 1, 2, 0, 1] + assert set(neighbors.index) == {0, 1, 2, 3, 4, 5, 6, 7} + + # Test with radius and pos of different length + with pytest.raises(ValueError): + neighbors = grid_moore.get_neighbors(radius=[1, 2], pos=[1, 1]) + + # Test with von_neumann neighborhood + grid_von_neumann.move_agents([0, 1, 2, 3], [[0, 1], [1, 0], [1, 2], [2, 1]]) + neighbors = grid_von_neumann.get_neighbors(radius=1, pos=[1, 1]) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (4, 2) + assert neighbors["dim_0"].to_list() == [0, 1, 1, 2] + assert neighbors["dim_1"].to_list() == [1, 0, 2, 1] + assert set(neighbors.index) == {0, 1, 2, 3} + + # Test with hexagonal neighborhood (odd cell [5,4] and even cell [5,5]) + grid_hexagonal.move_agents( + range(8), [[4, 4], [4, 5], [5, 3], [5, 5], [6, 3], [6, 4], [5, 4], [5, 6]] + ) + neighbors = grid_hexagonal.get_neighbors(radius=[2, 3], pos=[[5, 4], [5, 5]]) + assert isinstance(neighbors, pd.DataFrame) + assert neighbors.index.name == "agent_id" + assert neighbors.shape == (8, 2) # All agents are within the neighborhood + + # Sort the neighbors for consistent ordering + neighbors = neighbors.sort_values(["dim_0", "dim_1"]).reset_index(drop=True) + + assert neighbors["dim_0"].to_list() == [ + 4, + 4, + 5, + 5, + 5, + 5, + 6, + 6, + ] + assert neighbors["dim_1"].to_list() == [4, 5, 3, 4, 5, 6, 3, 4] + assert set(neighbors.index) == set(range(8)) + + def test_is_available(self, grid_moore: GridPandas): + # Test with GridCoordinate + result = grid_moore.is_available([0, 0]) + assert isinstance(result, pd.DataFrame) + assert result["available"].tolist() == [False] + result = grid_moore.is_available([1, 1]) + assert result["available"].tolist() == [True] + + # Test with GridCoordinates + result = grid_moore.is_available([[0, 0], [1, 1]]) + assert result["available"].tolist() == [False, True] + + def test_is_empty(self, grid_moore: GridPandas): + # Test with GridCoordinate + result = grid_moore.is_empty([0, 0]) + assert isinstance(result, pd.DataFrame) + assert result["empty"].tolist() == [False] + result = grid_moore.is_empty([1, 1]) + assert result["empty"].tolist() == [False] + + # Test with GridCoordinates + result = grid_moore.is_empty([[0, 0], [1, 1]]) + assert result["empty"].tolist() == [False, False] + + def test_is_full(self, grid_moore: GridPandas): + # Test with GridCoordinate + result = grid_moore.is_full([0, 0]) + assert isinstance(result, pd.DataFrame) + assert result["full"].tolist() == [True] + result = grid_moore.is_full([1, 1]) + assert result["full"].tolist() == [False] + + # Test with GridCoordinates + result = grid_moore.is_full([[0, 0], [1, 1]]) + assert result["full"].tolist() == [True, False] + + def test_move_agents( + self, + grid_moore: GridPandas, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ): + # Test with IdsLike + space = grid_moore.move_agents(agents=1, pos=[1, 1], inplace=False) + assert space.remaining_capacity == (2 * 3 * 3 - 2) + assert len(space.agents) == 2 + assert space.agents.index.to_list() == [0, 1] + assert space.agents["dim_0"].to_list() == [0, 1] + assert space.agents["dim_1"].to_list() == [0, 1] + + # Test with AgentSetDF + with pytest.warns(RuntimeWarning): + space = grid_moore.move_agents( + agents=fix2_AgentSetPolars, + pos=[[0, 0], [1, 0], [2, 0], [0, 1]], + inplace=False, + ) + assert space.remaining_capacity == (2 * 3 * 3 - 6) + assert len(space.agents) == 6 + assert space.agents.index.to_list() == [0, 1, 4, 5, 6, 7] + assert space.agents["dim_0"].to_list() == [0, 1, 0, 1, 2, 0] + assert space.agents["dim_1"].to_list() == [0, 1, 0, 0, 0, 1] + + # Test with Collection[AgentSetDF] + with pytest.warns(RuntimeWarning): + space = grid_moore.move_agents( + agents=[fix1_AgentSetPandas, fix2_AgentSetPolars], + pos=[[0, 2], [1, 2], [2, 2], [0, 1], [1, 1], [2, 1], [0, 0], [1, 0]], + inplace=False, + ) + assert space.remaining_capacity == (2 * 3 * 3 - 8) + assert len(space.agents) == 8 + assert space.agents.index.to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + assert space.agents["dim_0"].to_list() == [0, 1, 2, 0, 1, 2, 0, 1] + assert space.agents["dim_1"].to_list() == [2, 2, 2, 1, 1, 1, 0, 0] + + # Raises ValueError if len(agents) != len(pos) + with pytest.raises(ValueError): + space = grid_moore.move_agents( + agents=[0, 1], pos=[[0, 0], [1, 1], [2, 2]], inplace=False + ) + + # Test with AgentsDF, pos=DataFrame + pos = pd.DataFrame( + { + "dim_0": [0, 1, 2, 0, 1, 2, 0, 1], + "dim_1": [2, 2, 2, 1, 1, 1, 0, 0], + } + ) + + with pytest.warns(RuntimeWarning): + space = grid_moore.move_agents( + agents=grid_moore.model.agents, + pos=pos, + inplace=False, + ) + assert space.remaining_capacity == (2 * 3 * 3 - 8) + assert len(space.agents) == 8 + assert space.agents.index.to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + assert space.agents["dim_0"].to_list() == [0, 1, 2, 0, 1, 2, 0, 1] + assert space.agents["dim_1"].to_list() == [2, 2, 2, 1, 1, 1, 0, 0] + + # Test with agents=int, pos=DataFrame + pos = pd.DataFrame({"dim_0": [0], "dim_1": [2]}) + space = grid_moore.move_agents(agents=1, pos=pos, inplace=False) + assert space.remaining_capacity == (2 * 3 * 3 - 2) + assert len(space.agents) == 2 + assert space.agents.index.to_list() == [0, 1] + assert space.agents["dim_0"].to_list() == [0, 0] + assert space.agents["dim_1"].to_list() == [0, 2] + + def test_move_to_available(self, grid_moore: GridPandas): + # Test with GridCoordinate + last = None + different = False + for _ in range(10): + available_cells = grid_moore.available_cells + space = grid_moore.move_to_available(0, inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert space.agents[["dim_0", "dim_1"]].values[0] in available_cells.values + last = space.agents[["dim_0", "dim_1"]].values + assert different + + # Test with GridCoordinates + last = None + different = False + for _ in range(10): + available_cells = grid_moore.available_cells + space = grid_moore.move_to_available([0, 1], inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in available_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in available_cells.values) + last = space.agents[["dim_0", "dim_1"]].values + assert different + + # Test with AgentSetDF + last = None + different = False + for _ in range(10): + available_cells = grid_moore.available_cells + space = grid_moore.move_to_available(grid_moore.model.agents, inplace=False) + if last is not None and not different: + if (space.agents["dim_0"].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in available_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in available_cells.values) + last = space.agents["dim_0"].values + assert different + + def test_move_to_empty(self, grid_moore: GridPandas): + # Test with GridCoordinate + last = None + different = False + for _ in range(10): + empty_cells = grid_moore.empty_cells + space = grid_moore.move_to_empty(0, inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert space.agents[["dim_0", "dim_1"]].values[0] in empty_cells.values + last = space.agents[["dim_0", "dim_1"]].values + assert different + + # Test with GridCoordinates + last = None + different = False + for _ in range(10): + empty_cells = grid_moore.empty_cells + space = grid_moore.move_to_empty([0, 1], inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in empty_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in empty_cells.values) + last = space.agents[["dim_0", "dim_1"]].values + assert different + + # Test with AgentSetDF + last = None + different = False + for _ in range(10): + empty_cells = grid_moore.empty_cells + space = grid_moore.move_to_empty(grid_moore.model.agents, inplace=False) + if last is not None and not different: + if (space.agents["dim_0"].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in empty_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in empty_cells.values) + last = space.agents["dim_0"].values + assert different + + def test_out_of_bounds(self, grid_moore: GridPandas): + # Test with GridCoordinate + out_of_bounds = grid_moore.out_of_bounds([11, 11]) + assert isinstance(out_of_bounds, pd.DataFrame) + assert out_of_bounds.shape == (1, 3) + assert out_of_bounds.columns.to_list() == ["dim_0", "dim_1", "out_of_bounds"] + assert out_of_bounds.iloc[0].to_list() == [11, 11, True] + + # Test with GridCoordinates + out_of_bounds = grid_moore.out_of_bounds([[0, 0], [11, 11]]) + assert isinstance(out_of_bounds, pd.DataFrame) + assert out_of_bounds.shape == (2, 3) + assert out_of_bounds.columns.to_list() == ["dim_0", "dim_1", "out_of_bounds"] + assert out_of_bounds.iloc[0].to_list() == [0, 0, False] + assert out_of_bounds.iloc[1].to_list() == [11, 11, True] + + def test_place_agents( + self, + grid_moore: GridPandas, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ): + # Test with IdsLike + with pytest.warns(RuntimeWarning): + space = grid_moore.place_agents( + agents=[1, 2], pos=[[1, 1], [2, 2]], inplace=False + ) + assert space.remaining_capacity == (2 * 3 * 3 - 3) + assert len(space.agents) == 3 + assert space.agents.index.to_list() == [0, 1, 2] + assert space.agents["dim_0"].to_list() == [0, 1, 2] + assert space.agents["dim_1"].to_list() == [0, 1, 2] + + # Test with agents not in the model + with pytest.raises(ValueError): + space = grid_moore.place_agents( + agents=[10, 11], + pos=[[0, 0], [1, 0]], + inplace=False, + ) + + # Test with AgentSetDF + space = grid_moore.place_agents( + agents=fix2_AgentSetPolars, + pos=[[0, 0], [1, 0], [2, 0], [0, 1]], + inplace=False, + ) + assert space.remaining_capacity == (2 * 3 * 3 - 6) + assert len(space.agents) == 6 + assert space.agents.index.to_list() == [0, 1, 4, 5, 6, 7] + assert space.agents["dim_0"].to_list() == [0, 1, 0, 1, 2, 0] + assert space.agents["dim_1"].to_list() == [0, 1, 0, 0, 0, 1] + + # Test with Collection[AgentSetDF] + with pytest.warns(RuntimeWarning): + space = grid_moore.place_agents( + agents=[fix1_AgentSetPandas, fix2_AgentSetPolars], + pos=[[0, 2], [1, 2], [2, 2], [0, 1], [1, 1], [2, 1], [0, 0], [1, 0]], + inplace=False, + ) + assert space.remaining_capacity == (2 * 3 * 3 - 8) + assert len(space.agents) == 8 + assert space.agents.index.to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + assert space.agents["dim_0"].to_list() == [0, 1, 2, 0, 1, 2, 0, 1] + assert space.agents["dim_1"].to_list() == [2, 2, 2, 1, 1, 1, 0, 0] + + # Test with AgentsDF, pos=DataFrame + pos = pd.DataFrame( + { + "dim_0": [0, 1, 2, 0, 1, 2, 0, 1], + "dim_1": [2, 2, 2, 1, 1, 1, 0, 0], + } + ) + with pytest.warns(RuntimeWarning): + space = grid_moore.place_agents( + agents=grid_moore.model.agents, + pos=pos, + inplace=False, + ) + assert space.remaining_capacity == (2 * 3 * 3 - 8) + assert len(space.agents) == 8 + assert space.agents.index.to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + assert space.agents["dim_0"].to_list() == [0, 1, 2, 0, 1, 2, 0, 1] + assert space.agents["dim_1"].to_list() == [2, 2, 2, 1, 1, 1, 0, 0] + + # Test with agents=int, pos=DataFrame + pos = pd.DataFrame({"dim_0": [0], "dim_1": [2]}) + with pytest.warns(RuntimeWarning): + space = grid_moore.place_agents(agents=1, pos=pos, inplace=False) + assert space.remaining_capacity == (2 * 3 * 3 - 2) + assert len(space.agents) == 2 + assert space.agents.index.to_list() == [0, 1] + assert space.agents["dim_0"].to_list() == [0, 0] + assert space.agents["dim_1"].to_list() == [0, 2] + + def test_place_to_available(self, grid_moore: GridPandas): + # Test with GridCoordinate + last = None + different = False + for _ in range(10): + available_cells = grid_moore.available_cells + space = grid_moore.place_to_available(0, inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert space.agents[["dim_0", "dim_1"]].values[0] in available_cells.values + last = space.agents[["dim_0", "dim_1"]].values + assert different + # Test with GridCoordinates + last = None + different = False + for _ in range(10): + available_cells = grid_moore.available_cells + space = grid_moore.place_to_available([0, 1], inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in available_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in available_cells.values) + last = space.agents[["dim_0", "dim_1"]].values + assert different + # Test with AgentSetDF + last = None + different = False + for _ in range(10): + available_cells = grid_moore.available_cells + space = grid_moore.place_to_available( + grid_moore.model.agents, inplace=False + ) + if last is not None and not different: + if (space.agents["dim_0"].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in available_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in available_cells.values) + last = space.agents["dim_0"].values + assert different + + def test_place_to_empty(self, grid_moore: GridPandas): + # Test with GridCoordinate + last = None + different = False + for _ in range(10): + empty_cells = grid_moore.empty_cells + space = grid_moore.place_to_empty(0, inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert space.agents[["dim_0", "dim_1"]].values[0] in empty_cells.values + last = space.agents[["dim_0", "dim_1"]].values + assert different + # Test with GridCoordinates + last = None + different = False + for _ in range(10): + empty_cells = grid_moore.empty_cells + space = grid_moore.place_to_empty([0, 1], inplace=False) + if last is not None and not different: + if (space.agents[["dim_0", "dim_1"]].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in empty_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in empty_cells.values) + last = space.agents[["dim_0", "dim_1"]].values + assert different + # Test with AgentSetDF + last = None + different = False + for _ in range(10): + empty_cells = grid_moore.empty_cells + space = grid_moore.place_to_empty(grid_moore.model.agents, inplace=False) + if last is not None and not different: + if (space.agents["dim_0"].values != last).any(): + different = True + assert ( + space.agents[["dim_0", "dim_1"]].values[0] in empty_cells.values + ) and (space.agents[["dim_0", "dim_1"]].values[1] in empty_cells.values) + last = space.agents["dim_0"].values + assert different + + def test_random_agents(self, grid_moore: GridPandas): + # Test without seed + different = False + for _ in range(10): + agents0 = grid_moore.random_agents(1) + agents1 = grid_moore.random_agents(1) + if (agents0.values != agents1.values).all().all(): + different = True + break + assert different + + # Test with seed + agents0 = grid_moore.random_agents(1, seed=42) + agents1 = grid_moore.random_agents(1, seed=42) + assert (agents0 == agents1).all().all() + + def test_random_pos(self, grid_moore: GridPandas): + # Test without seed + different = False + last = None + for _ in range(10): + random_pos = grid_moore.random_pos(5) + assert isinstance(random_pos, pd.DataFrame) + assert len(random_pos) == 5 + assert random_pos.columns.to_list() == ["dim_0", "dim_1"] + assert not grid_moore.out_of_bounds(random_pos)["out_of_bounds"].any() + if last is not None and not different: + if (last != random_pos).any().any(): + different = True + break + last = random_pos + assert different + + # Test with seed + random_pos0 = grid_moore.random_pos(5, seed=42) + random_pos1 = grid_moore.random_pos(5, seed=42) + assert (random_pos0 == random_pos1).all().all() + + def test_remove_agents( + self, + grid_moore: GridPandas, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ): + grid_moore.move_agents( + [0, 1, 2, 3, 4, 5, 6, 7], + [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]], + ) + capacity = grid_moore.remaining_capacity + # Test with IdsLike + space = grid_moore.remove_agents([1, 2], inplace=False) + assert space.agents.shape == (6, 2) + assert space.remaining_capacity == capacity + 2 + assert space.agents.index.to_list() == [0, 3, 4, 5, 6, 7] + assert [ + x for id in space.model.agents.index.values() for x in id.to_list() + ] == [x for x in range(8)] + + # Test with AgentSetDF + space = grid_moore.remove_agents(fix1_AgentSetPandas, inplace=False) + assert space.agents.shape == (4, 2) + assert space.remaining_capacity == capacity + 4 + assert space.agents.index.to_list() == [4, 5, 6, 7] + assert [ + x for id in space.model.agents.index.values() for x in id.to_list() + ] == [x for x in range(8)] + + # Test with Collection[AgentSetDF] + space = grid_moore.remove_agents( + [fix1_AgentSetPandas, fix2_AgentSetPolars], inplace=False + ) + assert [ + x for id in space.model.agents.index.values() for x in id.to_list() + ] == [x for x in range(8)] + assert space.agents.empty + assert space.remaining_capacity == capacity + 8 + # Test with AgentsDF + space = grid_moore.remove_agents(grid_moore.model.agents, inplace=False) + assert space.remaining_capacity == capacity + 8 + assert space.agents.empty + assert [ + x for id in space.model.agents.index.values() for x in id.to_list() + ] == [x for x in range(8)] + + def test_sample_cells(self, grid_moore: GridPandas): + # Test with default parameters + replacement = False + same = True + last = None + for _ in range(10): + result = grid_moore.sample_cells(10) + assert len(result) == 10 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + counts = result.groupby(result.columns.to_list()).size() + assert (counts <= 2).all() + if not replacement and (counts > 1).any(): + replacement = True + if same and last is not None: + same = (result == last).all().all() + if not same and replacement: + break + last = result + assert replacement and not same + + # Test with too many samples + with pytest.raises(AssertionError): + grid_moore.sample_cells(100) + + # Test with 'empty' cell_type + + result = grid_moore.sample_cells(14, cell_type="empty") + assert len(result) == 14 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + counts = result.groupby(result.columns.to_list()).size() + + ## (0, 1) and (1, 1) are not in the result + assert not ( + (result["dim_0"] == 0) & (result["dim_1"] == 0) + ).any(), "Found (0, 1) in the result" + assert not ( + (result["dim_0"] == 1) & (result["dim_1"] == 1) + ).any(), "Found (1, 1) in the result" + + # 14 should be the max number of empty cells + with pytest.raises(AssertionError): + grid_moore.sample_cells(15, cell_type="empty") + + # Test with 'available' cell_type + result = grid_moore.sample_cells(16, cell_type="available") + assert len(result) == 16 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + counts = result.groupby(result.columns.to_list()).size() + + # 16 should be the max number of available cells + with pytest.raises(AssertionError): + grid_moore.sample_cells(17, cell_type="available") + + # Test with 'full' cell_type and no replacement + grid_moore.set_cells([[0, 0], [1, 1]], properties={"capacity": 1}) + result = grid_moore.sample_cells(2, cell_type="full", with_replacement=False) + assert len(result) == 2 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + assert ( + ((result["dim_0"] == 0) & (result["dim_1"] == 0)) + | ((result["dim_0"] == 1) & (result["dim_1"] == 1)) + ).all() + # 2 should be the max number of full cells + with pytest.raises(AssertionError): + grid_moore.sample_cells(3, cell_type="full", with_replacement=False) + + # Test with 'seed' + result = grid_moore.sample_cells(10, seed=42) + result2 = grid_moore.sample_cells(10, seed=42) + assert (result == result2).all().all() + + def test_set_cells(self, model: ModelDF): + grid_moore = GridPandas(model, dimensions=[3, 3], capacity=2) + + # Test with GridCoordinate + grid_moore.set_cells( + [0, 0], properties={"capacity": 1, "property_0": "value_0"} + ) + assert grid_moore.remaining_capacity == (2 * 3 * 3 - 1) + cell_df = grid_moore.get_cells([0, 0]) + assert cell_df.iloc[0]["capacity"] == 1 + assert cell_df.iloc[0]["property_0"] == "value_0" + + # Test with GridCoordinates + grid_moore.set_cells( + [[1, 1], [2, 2]], properties={"capacity": 3, "property_1": "value_1"} + ) + assert grid_moore.remaining_capacity == (2 * 3 * 3 - 1 + 2) + cell_df = grid_moore.get_cells([[1, 1], [2, 2]]) + assert cell_df.iloc[0]["capacity"] == 3 + assert cell_df.iloc[0]["property_1"] == "value_1" + assert cell_df.iloc[1]["capacity"] == 3 + assert cell_df.iloc[1]["property_1"] == "value_1" + cell_df = grid_moore.get_cells([0, 0]) + assert cell_df.iloc[0]["capacity"] == 1 + assert cell_df.iloc[0]["property_0"] == "value_0" + + # Test with DataFrame + df = pd.DataFrame( + {"dim_0": [0, 1, 2], "dim_1": [0, 1, 2], "capacity": [2, 2, 2]} + ) + grid_moore.set_cells(df) + assert grid_moore.remaining_capacity == (2 * 3 * 3) + + cells_df = grid_moore.get_cells([[0, 0], [1, 1], [2, 2]]) + + assert cells_df.iloc[0]["capacity"] == 2 + assert cells_df.iloc[1]["capacity"] == 2 + assert cells_df.iloc[2]["capacity"] == 2 + assert cells_df.iloc[0]["property_0"] == "value_0" + assert cells_df.iloc[1]["property_1"] == "value_1" + assert cells_df.iloc[2]["property_1"] == "value_1" + + # Add 2 agents to a cell, then set the cell capacity to 1 + grid_moore.place_agents([1, 2], [[0, 0], [0, 0]]) + with pytest.raises(AssertionError): + grid_moore.set_cells([0, 0], properties={"capacity": 1}) + + def test_swap_agents( + self, + grid_moore: GridPandas, + fix1_AgentSetPandas: ExampleAgentSetPandas, + fix2_AgentSetPolars: ExampleAgentSetPolars, + ): + grid_moore.move_agents( + [0, 1, 2, 3, 4, 5, 6, 7], + [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]], + ) + # Test with IdsLike + space = grid_moore.swap_agents([0, 1], [2, 3], inplace=False) + assert space.agents.loc[0].tolist() == grid_moore.agents.loc[2].tolist() + assert space.agents.loc[1].tolist() == grid_moore.agents.loc[3].tolist() + assert space.agents.loc[2].tolist() == grid_moore.agents.loc[0].tolist() + assert space.agents.loc[3].tolist() == grid_moore.agents.loc[1].tolist() + # Test with AgentSetDFs + space = grid_moore.swap_agents( + fix1_AgentSetPandas, fix2_AgentSetPolars, inplace=False + ) + assert space.agents.loc[0].to_list() == grid_moore.agents.loc[4].to_list() + assert space.agents.loc[1].to_list() == grid_moore.agents.loc[5].to_list() + assert space.agents.loc[2].to_list() == grid_moore.agents.loc[6].to_list() + assert space.agents.loc[3].tolist() == grid_moore.agents.loc[7].tolist() + + def test_torus_adj(self, grid_moore: GridPandas, grid_moore_torus: GridPandas): + # Test with non-toroidal grid + with pytest.raises(ValueError): + grid_moore.torus_adj([10, 10]) + + # Test with toroidal grid (GridCoordinate) + adj_df = grid_moore_torus.torus_adj([10, 8]) + assert isinstance(adj_df, pd.DataFrame) + assert adj_df.shape == (1, 2) + assert adj_df.columns.to_list() == ["dim_0", "dim_1"] + assert adj_df.iloc[0].to_list() == [1, 2] + + # Test with toroidal grid (GridCoordinates) + adj_df = grid_moore_torus.torus_adj([[10, 8], [15, 11]]) + assert isinstance(adj_df, pd.DataFrame) + assert adj_df.shape == (2, 2) + assert adj_df.columns.to_list() == ["dim_0", "dim_1"] + assert adj_df.iloc[0].to_list() == [1, 2] + assert adj_df.iloc[1].to_list() == [0, 2] + + def test___getitem__(self, grid_moore: GridPandas): + # Test out of bounds + with pytest.raises(ValueError): + grid_moore[[5, 5]] + + # Test with GridCoordinate + df = grid_moore[[0, 0]] + assert isinstance(df, pd.DataFrame) + assert df.index.names == ["dim_0", "dim_1"] + assert df.index.to_list() == [(0, 0)] + assert df.columns.to_list() == ["capacity", "property_0", "agent_id"] + assert df.iloc[0].to_list() == [1, "value_0", 0] + + # Test with GridCoordinates + df = grid_moore[[[0, 0], [1, 1]]] + assert isinstance(df, pd.DataFrame) + assert df.index.names == ["dim_0", "dim_1"] + assert df.index.to_list() == [(0, 0), (1, 1)] + assert df.columns.to_list() == ["capacity", "property_0", "agent_id"] + assert df.iloc[0].to_list() == [1, "value_0", 0] + assert df.iloc[1].to_list() == [3, "value_0", 1] + + def test___setitem__(self, grid_moore: GridPandas): + # Test with out-of-bounds + with pytest.raises(ValueError): + grid_moore[[5, 5]] = {"capacity": 10} + + # Test with GridCoordinate + grid_moore[[0, 0]] = {"capacity": 10} + assert grid_moore.get_cells([[0, 0]]).iloc[0]["capacity"] == 10 + # Test with GridCoordinates + grid_moore[[[0, 0], [1, 1]]] = {"capacity": 20} + assert grid_moore.get_cells([[0, 0], [1, 1]])["capacity"].tolist() == [20, 20] + + # Property tests + def test_agents(self, grid_moore: GridPandas): + assert isinstance(grid_moore.agents, pd.DataFrame) + assert grid_moore.agents.index.name == "agent_id" + assert grid_moore.agents.index.to_list() == [0, 1] + assert grid_moore.agents.columns.to_list() == ["dim_0", "dim_1"] + assert grid_moore.agents["dim_0"].to_list() == [0, 1] + assert grid_moore.agents["dim_1"].to_list() == [0, 1] + + def test_available_cells(self, grid_moore: GridPandas): + result = grid_moore.available_cells + assert len(result) == 8 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + + def test_cells(self, grid_moore: GridPandas): + result = grid_moore.cells + assert isinstance(result, pd.DataFrame) + assert result.index.names == ["dim_0", "dim_1"] + assert result.columns.to_list() == ["capacity", "property_0", "agent_id"] + assert result.index.to_list() == [(0, 0), (1, 1)] + assert result["capacity"].to_list() == [1, 3] + assert result["property_0"].to_list() == ["value_0", "value_0"] + assert result["agent_id"].to_list() == [0, 1] + + def test_dimensions(self, grid_moore: GridPandas): + assert isinstance(grid_moore.dimensions, list) + assert len(grid_moore.dimensions) == 2 + + def test_empty_cells(self, grid_moore: GridPandas): + result = grid_moore.empty_cells + assert len(result) == 7 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + + def test_full_cells(self, grid_moore: GridPandas): + grid_moore.set_cells([[0, 0], [1, 1]], {"capacity": 1}) + result = grid_moore.full_cells + assert len(result) == 2 + assert isinstance(result, pd.DataFrame) + assert result.columns.to_list() == ["dim_0", "dim_1"] + assert ( + ((result["dim_0"] == 0) & (result["dim_1"] == 0)) + | ((result["dim_0"] == 1) & (result["dim_1"] == 1)) + ).all() + + def test_model(self, grid_moore: GridPandas, model: ModelDF): + assert grid_moore.model == model + + def test_neighborhood_type( + self, + grid_moore: GridPandas, + grid_von_neumann: GridPandas, + grid_hexagonal: GridPandas, + ): + assert grid_moore.neighborhood_type == "moore" + assert grid_von_neumann.neighborhood_type == "von_neumann" + assert grid_hexagonal.neighborhood_type == "hexagonal" + + def test_random(self, grid_moore: GridPandas): + assert grid_moore.random == grid_moore.model.random + + def test_remaining_capacity(self, grid_moore: GridPandas): + assert grid_moore.remaining_capacity == (3 * 3 * 2 - 2) + + def test_torus(self, model: ModelDF, grid_moore: GridPandas): + assert not grid_moore.torus + + grid_2 = GridPandas(model, [3, 3], torus=True) + assert grid_2.torus diff --git a/tests/polars/__init__.py b/tests/polars/__init__.py new file mode 100644 index 0000000..e69de29