diff --git a/mesa_frames/abstract/agents.py b/mesa_frames/abstract/agents.py index 93baaaa..ce412c3 100644 --- a/mesa_frames/abstract/agents.py +++ b/mesa_frames/abstract/agents.py @@ -735,7 +735,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/abstract/space.py b/mesa_frames/abstract/space.py new file mode 100644 index 0000000..f20da2e --- /dev/null +++ b/mesa_frames/abstract/space.py @@ -0,0 +1,323 @@ +from abc import abstractmethod +from collections.abc import Collection, Sequence +from typing import TYPE_CHECKING + +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 ( + DataFrame, + GeoDataFrame, + IdsLike, + SpaceCoordinate, + SpaceCoordinates, +) + +ESPG = int + +if TYPE_CHECKING: + from mesa_frames.concrete.model import ModelDF + + +class SpaceDF(CopyMixin, DataFrameMixin): + _model: "ModelDF" + _agents: DataFrame | GeoDataFrame + + def __init__(self, model: "ModelDF") -> None: + """Create a new SpaceDF object. + + Parameters + ---------- + model : 'ModelDF' + + Returns + ------- + None + """ + self._model = model + + def random_agents( + self, + n: int, + seed: int | None = None, + ) -> DataFrame: + """Return a random sample of agents from the space. + + Parameters + ---------- + n : int + The number of agents to sample + 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. + + Returns + ------- + DataFrame + A DataFrame with the sampled agents + """ + if seed is None: + seed = self.random.integers(0) + return self._df_sample(self._agents, n=n, seed=seed) + + @abstractmethod + def get_directions( + self, + 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. + If the space is a Network, the direction is the shortest path between the two nodes. + In all other cases, the direction is the direction vector between the two positions. + Either positions (pos0, pos1) or agents (agents0, agents1) must be specified, not both and they must have the same length. + + Parameters + ---------- + pos0 : SpaceCoordinate | SpaceCoordinates | None, optional + The starting positions + pos1 : SpaceCoordinate | SpaceCoordinates | None, optional + The ending positions + agents0 : IdsLike | AgentContainer | Collection[AgentContainer] | None, optional + The starting agents + agents1 : IdsLike | AgentContainer | Collection[AgentContainer] | None, optional + The ending agents + normalize : bool, optional + Whether to normalize the vectors to unit norm. By default False + + Returns + ------- + DataFrame + A DataFrame where each row represents the direction from pos0 to pos1 or agents0 to agents1 + """ + ... + + @abstractmethod + def get_distances( + self, + 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. + 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. + You should specify either positions (pos0, pos1) or agents (agents0, agents1), not both and they must have the same length. + + Parameters + ---------- + pos0 : SpaceCoordinate | SpaceCoordinates | None, optional + The starting positions + pos1 : SpaceCoordinate | SpaceCoordinates | None, optional + The ending positions + agents0 : IdsLike | AgentContainer | Collection[AgentContainer], optional + The starting agents + agents1 : IdsLike | AgentContainer | Collection[AgentContainer], optional + The ending agents + + Returns + ------- + DataFrame + A DataFrame where each row represents the distance from pos0 to pos1 or agents0 to agents1 + """ + ... + + @abstractmethod + def get_neighbors( + self, + radius: int | float | Sequence[int] | Sequence[float], + pos: SpaceCoordinate | SpaceCoordinates | None = None, + agents: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, + include_center: bool = False, + ) -> DataFrame: + """Get the neighboring agents from given positions or agents according to the specified radiuses. + Either positions (pos0, pos1) or agents (agents0, agents1) must be specified, not both and they must have the same length. + + Parameters + ---------- + radius : int | float | Sequence[int] | Sequence[float] + The radius(es) of the neighborhood + pos : SpaceCoordinate | SpaceCoordinates | None, optional + The coordinates of the cell to get the neighborhood from, by default None + agents : IdsLike | AgentContainer | Collection[AgentContainer] | None, optional + The id of the agents to get the neighborhood from, by default None + include_center : bool, optional + If the center cells or agents should be included in the result, by default False + + Returns + ------- + DataFrame + A dataframe with neighboring agents. + The columns with '_center' suffix represent the center agent/position. + + Raises + ------ + ValueError + If both pos and agent are None or if both pos and agent are not None. + """ + ... + + @abstractmethod + def move_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 + + 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 + """ + ... + + @abstractmethod + def move_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). + + Parameters + ---------- + agents : IdsLike | AgentContainer | Collection[AgentContainer] + The agents to move to empty cells/positions + inplace : bool, optional + Whether to perform the operation inplace, by default True + + Returns + ------- + Self + """ + ... + + @abstractmethod + def random_pos( + self, + n: int, + seed: int | None = None, + ) -> DataFrame: + """Return a random sample of positions from the space. + + Parameters + ---------- + n : int + The number of positions to sample + 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. + + Returns + ------- + DataFrame + A DataFrame with the sampled positions + """ + ... + + @abstractmethod + def remove_agents( + self, + agents: IdsLike | AgentContainer | Collection[AgentContainer], + inplace: bool = True, + ): + """Remove agents from the space. + + Parameters + ---------- + agents : IdsLike | AgentContainer | Collection[AgentContainer] + The agents to remove from the space + inplace : bool, optional + Whether to perform the operation inplace, by default True + + Raises + ------ + ValueError + If some agents are not part of the model. + + Returns + ------- + Self + """ + ... + + @abstractmethod + def swap_agents( + self, + agents0: IdsLike | AgentContainer | Collection[AgentContainer], + agents1: IdsLike | AgentContainer | Collection[AgentContainer], + ) -> 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 + """ + + @abstractmethod + def __repr__(self) -> str: ... + + @abstractmethod + def __str__(self) -> str: ... + + @property + def agents(self) -> DataFrame | GeoDataFrame: + """Get the ids of the agents placed in the cell set, along with their coordinates or geometries + + Returns + ------- + AgentsDF + """ + return self._agents + + @property + def model(self) -> "ModelDF": + """The model to which the space belongs. + + Returns + ------- + 'ModelDF' + """ + return self._model + + @property + def random(self) -> Generator: + """The model's random number generator. + + Returns + ------- + Generator + """ + return self.model.random diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py index be4ef43..2c58125 100644 --- a/mesa_frames/concrete/model.py +++ b/mesa_frames/concrete/model.py @@ -4,6 +4,7 @@ import numpy as np from typing_extensions import Any +from mesa_frames.abstract.space import SpaceDF from mesa_frames.concrete.agents import AgentsDF if TYPE_CHECKING: @@ -59,6 +60,7 @@ class ModelDF: _seed: int | Sequence[int] running: bool _agents: AgentsDF + _space: SpaceDF | None # This will be a MultiSpaceDF object def __new__( cls, seed: int | Sequence[int] | None = None, *args: Any, **kwargs: Any @@ -77,6 +79,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.schedule = None self.current_id = 0 self._agents = AgentsDF(self) + self._space = None def get_agents_of_type(self, agent_type: type) -> "AgentSetDF": """Retrieve the AgentSetDF of a specified type. @@ -147,3 +150,15 @@ 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 diff --git a/mesa_frames/types_.py b/mesa_frames/types_.py index c34e792..b1b4ddf 100644 --- a/mesa_frames/types_.py +++ b/mesa_frames/types_.py @@ -1,6 +1,10 @@ from collections.abc import Collection from typing import Literal +from collections.abc import Sequence + +import geopandas as gpd +import geopolars as gpl import pandas as pd import polars as pl from numpy import ndarray @@ -15,20 +19,55 @@ AnyArrayLike = ArrayLike | pd.Index | pd.Series PandasMaskLike = AgnosticMask | pd.Series | pd.DataFrame | AnyArrayLike PandasIdsLike = AgnosticIds | pd.Series | pd.Index +PandasGridCapacity = ndarray ###----- Polars Types -----### PolarsMaskLike = AgnosticMask | pl.Expr | pl.Series | pl.DataFrame | Collection[int] PolarsIdsLike = AgnosticIds | pl.Series +PolarsGridCapacity = list[pl.Expr] ###----- Generic -----### - +GeoDataFrame = gpd.GeoDataFrame | gpl.GeoDataFrame DataFrame = pd.DataFrame | pl.DataFrame Series = pd.Series | pl.Series +Series = pd.Series | pl.Series Index = pd.Index | pl.Series BoolSeries = pd.Series | pl.Series MaskLike = AgnosticMask | PandasMaskLike | PolarsMaskLike IdsLike = AgnosticIds | PandasIdsLike | PolarsIdsLike + ###----- Time ------### TimeT = float | int + +###----- Space -----### + +NetworkCoordinate = int | DataFrame + +GridCoordinate = int | Sequence[int] | DataFrame + +DiscreteCoordinate = NetworkCoordinate | GridCoordinate +ContinousCoordinate = float | Sequence[float] | DataFrame + +SpaceCoordinate = DiscreteCoordinate | ContinousCoordinate + + +NetworkCoordinates = NetworkCoordinate | Collection[NetworkCoordinate] +GridCoordinates = ( + GridCoordinate | Sequence[int | slice | Sequence[int]] | Collection[GridCoordinate] +) + +DiscreteCoordinates = NetworkCoordinates | GridCoordinates +ContinousCoordinates = ( + ContinousCoordinate + | Sequence[float | Sequence[float]] + | Collection[ContinousCoordinate] +) + +SpaceCoordinates = DiscreteCoordinates | ContinousCoordinates + +GridCapacity = PandasGridCapacity | PolarsGridCapacity +NetworkCapacity = DataFrame + +DiscreteSpaceCapacity = GridCapacity | NetworkCapacity diff --git a/pyproject.toml b/pyproject.toml index 137e98e..f7091cd 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 + "polars>=1.0.0", #polars._typing (see mesa_frames.types) added in 1.0.0 + "geopolars" ] + dev = [ "mesa_frames[pandas,polars]", "perfplot", #readme_script