addind conversion to geobjects

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
 

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.
 

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,

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

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)

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_
+        """

mesa_frames/types.py → 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

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"]