diff --git a/mesa_frames/abstract/space.py b/mesa_frames/abstract/space.py index f20da2e..53238fd 100644 --- a/mesa_frames/abstract/space.py +++ b/mesa_frames/abstract/space.py @@ -1,14 +1,19 @@ from abc import abstractmethod -from collections.abc import Collection, Sequence -from typing import TYPE_CHECKING +from collections.abc import Callable, Collection, Sequence +from typing import TYPE_CHECKING, Literal +import polars as pl from numpy.random import Generator from typing_extensions import Self from mesa_frames.abstract.agents import AgentContainer from mesa_frames.abstract.mixin import CopyMixin, DataFrameMixin from mesa_frames.types_ import ( + BoolSeries, DataFrame, + DiscreteCoordinate, + DiscreteCoordinates, + DiscreteSpaceCapacity, GeoDataFrame, IdsLike, SpaceCoordinate, @@ -22,8 +27,63 @@ class SpaceDF(CopyMixin, DataFrameMixin): + """The SpaceDF class is an abstract class that defines the interface for all space classes in mesa_frames. + + Methods + ------- + __init__(model: 'ModelDF') + Create a new SpaceDF object. + random_agents(n: int, seed: int | None = None) -> DataFrame + Return a random sample of agents from the space. + get_directions( + pos0: SpaceCoordinate | SpaceCoordinates | None = None, + pos1: SpaceCoordinate | SpaceCoordinates | None = None, + agents0: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + normalize: bool = False, + ) -> DataFrame + Returns the directions from pos0 to pos1 or agents0 and agents1. + get_distances( + pos0: SpaceCoordinate | SpaceCoordinates | None = None, + pos1: SpaceCoordinate | SpaceCoordinates | None = None, + agents0: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + ) -> DataFrame + Returns the distances from pos0 to pos1 or agents0 and agents1. + get_neighbors( + radius: int | float | Sequence[int] | Sequence[float], + pos: Space + ) -> DataFrame + Get the neighboring agents from given positions or agents according to the specified radiuses. + move_agents( + agents: IdsLike | AgentContainer | Collection[AgentContainer], + pos + ) -> Self + Place agents in the space according to the specified coordinates. + move_to_empty( + agents: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ) -> Self + Move agents to empty cells/positions in the space. + random_pos( + n: int, + seed: int | None = None, + ) -> DataFrame + Return a random sample of positions from the space. + remove_agents( + agents: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ) + Remove agents from the space. + swap_agents( + agents0: IdsLike | AgentContainer | Collection[AgentContainer], + agents1: IdsLike | AgentContainer | Collection[AgentContainer], + ) -> Self + Swap the positions of the agents in the space. + """ + _model: "ModelDF" - _agents: DataFrame | GeoDataFrame + _agents: DataFrame | GeoDataFrame # Stores the agents placed in the space def __init__(self, model: "ModelDF") -> None: """Create a new SpaceDF object. @@ -321,3 +381,355 @@ def random(self) -> Generator: Generator """ return self.model.random + + +class DiscreteSpaceDF(SpaceDF): + """The DiscreteSpaceDF class is an abstract class that defines the interface for all discrete space classes (Grids and Networks) in mesa_frames. + + Methods + ------- + __init__(model: 'ModelDF', capacity: int | None = None) + Create a new DiscreteSpaceDF object. + is_available(pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame + Check whether the input positions are available (there exists at least one remaining spot in the cells). + is_empty(pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame + Check whether the input positions are empty (there isn't any single agent in the cells). + is_full(pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame + Check whether the input positions are full (there isn't any spot available in the cells). + move_to_empty(agents: IdsLike | AgentContainer | Collection[AgentContainer], inplace: bool = True) -> Self + Move agents to empty cells in the space (cells where there isn't any single agent). + move_to_available(agents: IdsLike | AgentContainer | Collection[AgentContainer], inplace: bool = True) -> Self + Move agents to available cells in the space (cells where there is at least one spot available). + sample_cells(n: int, cell_type: Literal["any", "empty", "available", "full"] = "any", with_replacement: bool = True) -> DataFrame + Sample cells from the grid according to the specified cell_type. + get_neighborhood(radius: int | float | Sequence[int] | Sequence[float], pos: DiscreteCoordinate | DiscreteCoordinates | None = None, agents: IdsLike | AgentContainer | Collection[AgentContainer] = None, include_center: bool = False) -> DataFrame + Get the neighborhood cells from the given positions (pos) or agents according to the specified radiuses. + get_cells(cells: DiscreteCoordinates | None = None) -> DataFrame + Retrieve a dataframe of specified cells with their properties and agents. + set_cells(properties: DataFrame, cells: DiscreteCoordinates | None = None, inplace: bool = True) -> Self + Set the properties of the specified cells. + """ + + _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) + + def __init__( + self, + model: "ModelDF", + capacity: int | None = None, + ): + """Create a DiscreteSpaceDF object. + NOTE: The capacity specified here is the default capacity, + it can be set also per cell through the set_cells method. + + Parameters + ---------- + model : ModelDF + The model to which the space belongs + capacity : int | None, optional + The maximum capacity for cells, by default None (infinite) + """ + super().__init__(model) + self._capacity = capacity + + def is_available(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: + """Check whether the input positions are available (there exists at least one remaining spot in the cells) + + Parameters + ---------- + pos : GridCoordinate | GridCoordinates + The positions to check for + + Returns + ------- + 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), + ) + + def is_empty(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: + """Check whether the input positions are empty (there isn't any single agent in the cells) + + Parameters + ---------- + pos : GridCoordinate | GridCoordinates + The positions to check for + + Returns + ------- + 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) + ) + + def is_full(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: + """Check whether the input positions are full (there isn't any spot available in the cells) + + Parameters + ---------- + pos : GridCoordinate | GridCoordinates + The positions to check for + + Returns + ------- + 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) + ) + + 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) + + def move_to_available( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ) -> Self: + """Move agents to available cells/positions in the space (cells/positions where there is at least one spot available). + + Parameters + ---------- + agents : IdsLike | AgentContainer | Collection[AgentContainer] + The agents to move to available cells/positions + inplace : bool, optional + Whether to perform the operation inplace, by default True + + Returns + ------- + Self + """ + return self._move_agents_to_cells( + agents, cell_type="available", inplace=inplace + ) + + def sample_cells( + self, + n: int, + cell_type: Literal["any", "empty", "available", "full"] = "any", + with_replacement: bool = True, + ) -> DataFrame: + """Sample cells from the grid according to the specified cell_type. + + Parameters + ---------- + n : int + The number of cells to sample + cell_type : Literal["any", "empty", "available", "full"], optional + The type of cells to sample, by default "any" + with_replacement : bool, optional + If the sampling should be with replacement, by default True + + Returns + ------- + DataFrame + A DataFrame with the sampled cells + """ + match cell_type: + case "any": + condition = self._any_cell_condition + case "empty": + condition = self._empty_cell_condition + case "available": + condition = self._available_cell_condition + case "full": + condition = self._full_cell_condition + return self._sample_cells(n, with_replacement, condition=condition) + + @abstractmethod + def get_neighborhood( + self, + radius: int | float | Sequence[int] | Sequence[float], + pos: DiscreteCoordinate | DiscreteCoordinates | None = None, + agents: IdsLike | AgentContainer | Collection[AgentContainer] = None, + include_center: bool = False, + ) -> DataFrame: + """Get the neighborhood cells from the given positions (pos) or agents according to the specified radiuses. + Either positions (pos) or agents must be specified, not both. + + Parameters + ---------- + radius : int | float | Sequence[int] | Sequence[float] + The radius(es) of the neighborhoods + pos : DiscreteCoordinate | DiscreteCoordinates | None, optional + The coordinates of the cell(s) to get the neighborhood from + agents : IdsLike | AgentContainer | Collection[AgentContainer], optional + The agent(s) to get the neighborhood from + include_center : bool, optional + If the cell in the center of the neighborhood should be included in the result, by default False + + Returns + ------- + DataFrame + A dataframe where + - Columns are called according to the coordinates of the space(['dim_0', 'dim_1', ...] in Grids, ['node_id', 'edge_id'] in Networks) + - Rows represent the coordinates of a neighboring cells + """ + ... + + @abstractmethod + def get_cells(self, cells: DiscreteCoordinates | None = None) -> DataFrame: + """Retrieve a dataframe of specified cells with their properties and agents. + + Parameters + ---------- + cells : CellCoordinates, default is optional (all cells retrieved) + + Returns + ------- + DataFrame + A DataFrame with the properties of the cells and the agents placed in them. + """ + ... + + @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. + + + 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 + + Returns + ------- + Self + """ + ... + + def _move_agents_to_cells( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + cell_type: Literal["any", "empty", "available"], + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + + # Get Ids of agents + # TODO: fix this + if isinstance(agents, AgentContainer | Collection[AgentContainer]): + agents = agents.index + + # 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") + + # Get cells of specified type + cells = obj.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 + + def _any_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return True + + def _empty_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return cap == self._capacity + + def _available_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return cap > 0 + + def _full_cell_condition(self, cap: DiscreteSpaceCapacity) -> BoolSeries: + return cap == 0 + + @abstractmethod + def _sample_cells( + self, + n: int | None, + with_replacement: bool, + condition: Callable[[DiscreteSpaceCapacity], BoolSeries], + ) -> DataFrame: + """Sample cells from the grid according to a condition on the capacity. + + Parameters + ---------- + n : int | None + The number of cells to sample + with_replacement : bool + If the sampling should be with replacement + condition : Callable[[DiscreteSpaceCapacity], BoolSeries] + The condition to apply on the capacity + + Returns + ------- + DataFrame + """ + ... + + def __getitem__(self, cells: DiscreteCoordinates): + return self.get_cells(cells) + + def __setitem__(self, cells: DiscreteCoordinates, properties: DataFrame): + self.set_cells(properties=properties, 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] + + @property + def cells(self) -> DataFrame: + return self.get_cells() + + @cells.setter + def cells(self, df: DataFrame): + return self.set_cells(df, inplace=True) + + @property + def empty_cells(self) -> DataFrame: + return self._sample_cells( + None, with_replacement=False, condition=self._empty_cell_condition + ) + + @property + def available_cells(self) -> DataFrame: + return self._sample_cells( + None, with_replacement=False, condition=self._available_cell_condition + ) + + @property + def full_cells(self) -> DataFrame: + return self._sample_cells( + None, with_replacement=False, condition=self._full_cell_condition + )