diff --git a/mesa_frames/abstract/agents.py b/mesa_frames/abstract/agents.py index 4217648..299aa2a 100644 --- a/mesa_frames/abstract/agents.py +++ b/mesa_frames/abstract/agents.py @@ -9,11 +9,12 @@ from typing_extensions import Any, Self, overload from mesa_frames.abstract.mixin import CopyMixin -from mesa_frames.types import BoolSeries, DataFrame, IdsLike, Index, MaskLike, Series +from mesa_frames.types_ import BoolSeries, DataFrame, IdsLike, Index, MaskLike, Series if TYPE_CHECKING: from mesa_frames.concrete.agents import AgentSetDF from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.space import SpaceDF class AgentContainer(CopyMixin): @@ -359,6 +360,23 @@ def sort( A new or updated AgentContainer. """ + @abstractmethod + def _convert_to_geobject(self, space: SpaceDF, inplace: bool = True) -> Self: + """Converts the DataFrame(s) of AgentContainer to GeoDataFrame(s). + + Parameters + ---------- + space : SpaceDF + The space to add to the AgentContainer. Determines the geometry type. + inplace : bool + Whether to add the space column in place. + + Returns + ------- + Self + """ + ... + def __add__(self, other) -> Self: return self.add(agents=other, inplace=False) @@ -735,7 +753,9 @@ def __init__(self, model: ModelDF) -> None: @abstractmethod def add( - self, agents: DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True + self, + agents: DataFrame | Sequence[Any] | dict[str, Any], + inplace: bool = True, ) -> Self: """Add agents to the AgentSetDF diff --git a/mesa_frames/concrete/agents.py b/mesa_frames/concrete/agents.py index 14891b8..fec000d 100644 --- a/mesa_frames/concrete/agents.py +++ b/mesa_frames/concrete/agents.py @@ -1,14 +1,13 @@ from collections import defaultdict from collections.abc import Callable, Collection, Iterable, Iterator, Sequence -from typing import Literal, cast +from typing import TYPE_CHECKING, Literal, cast import polars as pl from typing_extensions import Any, Self, overload -from typing import TYPE_CHECKING - from mesa_frames.abstract.agents import AgentContainer, AgentSetDF -from mesa_frames.types import ( +from mesa_frames.concrete.space import SpaceDF +from mesa_frames.types_ import ( AgnosticMask, BoolSeries, DataFrame, @@ -319,6 +318,12 @@ def sort( ] return obj + def _convert_to_geobject(self, space: SpaceDF, inplace: bool = True) -> Self: + obj = self._get_obj(inplace) + for agentset in obj._agentsets: + agentset._convert_to_geobject(space, inplace=True) + return obj + def _check_ids_presence(self, other: list[AgentSetDF]) -> pl.DataFrame: """Check if the IDs of the agents to be added are unique. diff --git a/mesa_frames/concrete/agentset_pandas.py b/mesa_frames/concrete/agentset_pandas.py index 8f4ce4d..848e247 100644 --- a/mesa_frames/concrete/agentset_pandas.py +++ b/mesa_frames/concrete/agentset_pandas.py @@ -1,17 +1,20 @@ from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import TYPE_CHECKING +import geopandas as gpd import pandas as pd import polars as pl from typing_extensions import Any, Self, overload from mesa_frames.abstract.agents import AgentSetDF from mesa_frames.concrete.agentset_polars import AgentSetPolars -from mesa_frames.types import PandasIdsLike, PandasMaskLike +from mesa_frames.types_ import PandasIdsLike, PandasMaskLike if TYPE_CHECKING: from mesa_frames.concrete.model import ModelDF +from mesa_frames.concrete.space import SpaceDF + class AgentSetPandas(AgentSetDF): _agents: pd.DataFrame @@ -108,11 +111,18 @@ def __init__(self, model: "ModelDF") -> None: def add( self, - agents: pd.DataFrame | Sequence[Any] | dict[str, Any], + agents: pd.DataFrame | gpd.GeoDataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True, ) -> Self: obj = self._get_obj(inplace) - if isinstance(agents, pd.DataFrame): + if isinstance(agents, gpd.GeoDataFrame): + try: + self.model.space + except ValueError: + raise ValueError( + "You are adding agents with a GeoDataFrame but haven't set model.space. Set it before adding agents with a GeoDataFrame or add agents with a standard DataFrame" + ) + if isinstance(agents, (pd.DataFrame, gpd.GeoDataFrame)): new_agents = agents if "unique_id" != agents.index.name: try: @@ -312,6 +322,12 @@ def _concatenate_agentsets( self.remove(ids_to_remove, inplace=True) return self + def _convert_to_geobject(self, space: SpaceDF, inplace: bool = True) -> Self: + obj = self._get_obj(inplace) + + obj._agents = gpd.GeoDataFrame(obj._agents, geometry=None) + return obj + def _get_bool_mask( self, mask: PandasMaskLike = None, diff --git a/mesa_frames/concrete/agentset_polars.py b/mesa_frames/concrete/agentset_polars.py index 358a310..f7f6dff 100644 --- a/mesa_frames/concrete/agentset_polars.py +++ b/mesa_frames/concrete/agentset_polars.py @@ -1,12 +1,13 @@ from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import TYPE_CHECKING +import geopolars as gpl import polars as pl from polars._typing import IntoExpr from typing_extensions import Any, Self, overload from mesa_frames.concrete.agents import AgentSetDF -from mesa_frames.types import PolarsIdsLike, PolarsMaskLike +from mesa_frames.types_ import PolarsIdsLike, PolarsMaskLike if TYPE_CHECKING: from mesa_frames.concrete.agentset_pandas import AgentSetPandas @@ -136,7 +137,14 @@ def add( The updated AgentSetPolars. """ obj = self._get_obj(inplace) - if isinstance(agents, pl.DataFrame): + if isinstance(agents, gpl.GeoDataFrame): + try: + self.model.space + except ValueError: + raise ValueError( + "You are adding agents with a GeoDataFrame but haven't set model.space. Set it before adding agents with a GeoDataFrame or add agents with a standard DataFrame" + ) + if isinstance(agents, gpl.GeoDataFrame, pl.DataFrame): if "unique_id" not in agents.columns: raise KeyError("DataFrame must have a unique_id column.") new_agents = agents diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py index be4ef43..3bea3ce 100644 --- a/mesa_frames/concrete/model.py +++ b/mesa_frames/concrete/model.py @@ -5,6 +5,7 @@ from typing_extensions import Any from mesa_frames.concrete.agents import AgentsDF +from mesa_frames.concrete.space import SpaceDF if TYPE_CHECKING: from mesa_frames.abstract.agents import AgentSetDF @@ -59,6 +60,7 @@ class ModelDF: _seed: int | Sequence[int] running: bool _agents: AgentsDF + _space: SpaceDF | None def __new__( cls, seed: int | Sequence[int] | None = None, *args: Any, **kwargs: Any @@ -147,3 +149,14 @@ def agents(self, agents: AgentsDF) -> None: @property def agent_types(self) -> list[type]: return [agent.__class__ for agent in self._agents._agentsets] + + @property + def space(self) -> SpaceDF: + if not self._space: + raise ValueError("You haven't set the space for the model. Use model.space = your_space") + return self._space + + @space.setter + def space(self, space: SpaceDF) -> None: + self._space = space + self._agents._convert_to_geobject(self._space, inplace=True) diff --git a/mesa_frames/concrete/space.py b/mesa_frames/concrete/space.py new file mode 100644 index 0000000..ff6138b --- /dev/null +++ b/mesa_frames/concrete/space.py @@ -0,0 +1,76 @@ +""" +Mesa Frames Space Module +================= + +Objects used to add a spatial component to a model. + +Grid: base grid, which creates a rectangular grid. +SingleGrid: extension to Grid which strictly enforces one agent per cell. +MultiGrid: extension to Grid where each cell can contain a set of agents. +HexGrid: extension to Grid to handle hexagonal neighbors. +ContinuousSpace: a two-dimensional space where each agent has an arbitrary + position of `float`'s. +NetworkGrid: a network where each node contains zero or more agents. +""" + +from mesa_frames.abstract.agents import AgentContainer +from mesa_frames.types_ import IdsLike, PositionsLike + + +class SpaceDF: + def _check_empty_pos(pos: PositionsLike) -> bool: + """Check if the given positions are empty. + + Parameters + ---------- + pos : DataFrame | tuple[Series, Series] | Series + Input positions to check. + + Returns + ------- + Series[bool] + Whether + """ + + +class SingleGrid(SpaceDF): + """Rectangular grid where each cell contains exactly at most one agent. + + Grid cells are indexed by [x, y], where [0, 0] is assumed to be the + bottom-left and [width-1, height-1] is the top-right. If a grid is + toroidal, the top and bottom, and left and right, edges wrap to each other. + + This class provides a property `empties` that returns a set of coordinates + for all empty cells in the grid. It is automatically updated whenever + agents are added or removed from the grid. The `empties` property should be + used for efficient access to current empty cells rather than manually + iterating over the grid to check for emptiness. + + """ + + def place_agents(self, agents: IdsLike | AgentContainer, pos: PositionsLike): + """Place agents on the grid at the coordinates specified in pos. + NOTE: The cells must be empty. + + + Parameters + ---------- + agents : IdsLike | AgentContainer + + pos : DataFrame | tuple[Series, Series] + _description_ + """ + + def _check_empty_pos(pos: PositionsLike) -> bool: + """Check if the given positions are empty. + + Parameters + ---------- + pos : DataFrame | tuple[Series, Series] + Input positions to check. + + Returns + ------- + bool + _description_ + """ diff --git a/mesa_frames/types.py b/mesa_frames/types_.py similarity index 50% rename from mesa_frames/types.py rename to mesa_frames/types_.py index 232aa0f..a2c273c 100644 --- a/mesa_frames/types.py +++ b/mesa_frames/types_.py @@ -1,9 +1,11 @@ from collections.abc import Collection -from typing import Literal +import geopandas as gpd +import geopolars as gpl import pandas as pd import polars as pl from numpy import ndarray +from typing_extensions import Literal, Sequence ####----- Agnostic Types -----#### AgnosticMask = Literal["all", "active"] | None @@ -23,10 +25,38 @@ ###----- Generic -----### -DataFrame = pd.DataFrame | pl.DataFrame -Series = pd.Series | pl.Series +GeoDataFame = gpd.GeoDataFrame | gpl.GeoDataFrame +GeoSeries = gpd.GeoSeries | gpl.GeoSeries +DataFrame = pd.DataFrame | pl.DataFrame | GeoDataFame +Series = pd.Series | pl.Series | GeoSeries Index = pd.Index | pl.Series BoolSeries = pd.Series | pl.Series MaskLike = AgnosticMask | PandasMaskLike | PolarsMaskLike IdsLike = AgnosticIds | PandasIdsLike | PolarsIdsLike + + +###----- Time ------### TimeT = float | int + + +###----- Space -----### +Coordinates = tuple[int, int] | tuple[float, float] +Node_ID = int +AgnosticPositionsLike = ( + Sequence[Coordinates] | Sequence[Node_ID] | Coordinates | Node_ID +) +PolarsPositionsLike = ( + AgnosticPositionsLike + | pl.DataFrame + | tuple[pl.Series, pl.Series] + | gpl.GeoSeries + | pl.Series +) +PandasPositionsLike = ( + AgnosticPositionsLike + | pd.DataFrame + | tuple[pd.Series, pd.Series] + | gpd.GeoSeries + | pd.Series +) +PositionsLike = AgnosticPositionsLike | PolarsPositionsLike | PandasPositionsLike diff --git a/pyproject.toml b/pyproject.toml index b9c5d27..ece6ac9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,10 +19,13 @@ dependencies = [ pandas = [ "pandas~=2.2", "pyarrow", + "geopandas" ] polars = [ "polars>=1.0.0", #polars._typing (see mesa_frames.types) added in 1.0.0 + "geopolars" ] + dev = [ "mesa_frames[pandas,polars]", "pytest", @@ -31,5 +34,8 @@ dev = [ "mesa", ] +[tool.hatch.envs.dev] +features = ["dev"] + [tool.hatch.build.targets.wheel] packages = ["mesa_frames"]