diff --git a/mesa_frames/abstract/agents.py b/mesa_frames/abstract/agents.py index 9edec7c..cc66287 100644 --- a/mesa_frames/abstract/agents.py +++ b/mesa_frames/abstract/agents.py @@ -1,5 +1,3 @@ - - from __future__ import annotations # PEP 563: postponed evaluation of type annotations from abc import ABC, abstractmethod @@ -20,7 +18,7 @@ from numpy.random import Generator -from mesa_frames.types import BoolSeries, DataFrame, MaskLike, Series +from mesa_frames.types import BoolSeries, DataFrame, IdsLike, MaskLike, Series if TYPE_CHECKING: from mesa_frames.concrete.model import ModelDF @@ -42,17 +40,17 @@ class AgentContainer(ABC): ------- copy(deep: bool = False, memo: dict | None = None) -> Self Create a copy of the AgentContainer. - discard(ids: MaskLike, inplace: bool = True) -> Self + discard(ids: IdsLike, inplace: bool = True) -> Self Removes an agent from the AgentContainer. Does not raise an error if the agent is not found. add(other: Any, inplace: bool = True) -> Self Add agents to the AgentContainer. - contains(ids: Hashable | Collection[Hashable]) -> bool | BoolSeries + contains(ids: IdsLike) -> bool | BoolSeries Check if agents with the specified IDs are in the AgentContainer. do(method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any | dict[str, Any] Invoke a method on the AgentContainer. get(attr_names: str | Collection[str] | None = None, mask: MaskLike | None = None) -> Series | DataFrame | dict[str, Series] | dict[str, DataFrame] Retrieve the value of a specified attribute for each agent in the AgentContainer. - remove(ids: MaskLike, inplace: bool = True) -> Self + remove(ids: IdsLike, inplace: bool = True) -> Self Removes an agent from the AgentContainer. select(mask: MaskLike | None = None, filter_func: Callable[[Self], MaskLike] | None = None, n: int | None = None, negate: bool = False, inplace: bool = True) -> Self Select agents in the AgentContainer based on the given criteria. @@ -144,7 +142,7 @@ def copy( return obj - def discard(self, ids: MaskLike, inplace: bool = True) -> Self: + def discard(self, ids: IdsLike, inplace: bool = True) -> Self: """Removes an agent from the AgentContainer. Does not raise an error if the agent is not found. Parameters @@ -182,19 +180,19 @@ def add(self, other, inplace: bool = True) -> Self: @overload @abstractmethod - def contains(self, ids: Collection[Hashable]) -> BoolSeries: ... + def contains(self, ids: int) -> bool: ... @overload @abstractmethod - def contains(self, ids: Hashable) -> bool: ... + def contains(self, ids: IdsLike) -> BoolSeries: ... @abstractmethod - def contains(self, ids: Hashable | Collection[Hashable]) -> bool | BoolSeries: + def contains(self, ids: IdsLike) -> bool | BoolSeries: """Check if agents with the specified IDs are in the AgentContainer. Parameters ---------- - ids : Hashable | Collection[Any] + ids : IdsLike The ID(s) to check for. Returns @@ -283,7 +281,7 @@ def get( ... @abstractmethod - def remove(self, ids: MaskLike, inplace: bool = True) -> Self: + def remove(self, ids: IdsLike, inplace: bool = True) -> Self: """Removes an agent from the AgentContainer. Parameters @@ -446,7 +444,7 @@ def _get_obj(self, inplace: bool) -> Self: def __add__(self, other) -> Self: return self.add(other=other, inplace=False) - def __contains__(self, id: Hashable) -> bool: + def __contains__(self, id: int) -> bool: """Check if an agent is in the AgentContainer. Parameters @@ -459,13 +457,9 @@ def __contains__(self, id: Hashable) -> bool: bool True if the agent is in the AgentContainer, False otherwise. """ - bool_series = self.contains(ids=id) - if isinstance(bool_series, bool): - return bool_series - elif len(bool_series) == 1: - return bool_series[0].value - else: - raise ValueError("The in operator can only be used with a single ID.") + if not isinstance(id, int): + raise TypeError("id must be an integer") + return self.contains(ids=id) def __copy__(self) -> Self: """Create a shallow copy of the AgentContainer. @@ -543,7 +537,7 @@ def __iadd__(self, other) -> Self: """ return self.add(other=other, inplace=True) - def __isub__(self, other: MaskLike) -> Self: + def __isub__(self, other: IdsLike) -> Self: """Remove agents from the AgentContainer through the -= operator. Parameters @@ -558,7 +552,7 @@ def __isub__(self, other: MaskLike) -> Self: """ return self.discard(other, inplace=True) - def __sub__(self, other: MaskLike) -> Self: + def __sub__(self, other: IdsLike) -> Self: """Remove agents from a new AgentContainer through the - operator. Parameters @@ -1038,4 +1032,4 @@ def active_agents(self) -> DataFrame: ... @property @abstractmethod - def inactive_agents(self) -> DataFrame: ... \ No newline at end of file + def inactive_agents(self) -> DataFrame: ... diff --git a/mesa_frames/concrete/agents.py b/mesa_frames/concrete/agents.py index 4aceaba..34cfe65 100644 --- a/mesa_frames/concrete/agents.py +++ b/mesa_frames/concrete/agents.py @@ -1,7 +1,12 @@ +from operator import ne from typing import Any, Callable, Iterable, Iterator, Literal, Self, Sequence, overload +import polars as pl + from mesa_frames.abstract.agents import AgentContainer, AgentSetDF, Collection, Hashable -from mesa_frames.types import BoolSeries, DataFrame, MaskLike, Series +from mesa_frames.concrete.agentset_pandas import AgentSetPandas +from mesa_frames.concrete.agentset_polars import AgentSetPolars +from mesa_frames.types import BoolSeries, DataFrame, IdsLike, MaskLike, Series class AgentsDF(AgentContainer): @@ -10,6 +15,7 @@ class AgentsDF(AgentContainer): "_agentsets": ("copy", []), } _backend: str + _ids: pl.Series """A collection of AgentSetDFs. All agents of the model are stored here. Attributes @@ -38,19 +44,19 @@ class AgentsDF(AgentContainer): ------- __init__(self) -> None Initialize a new AgentsDF. - add(self, other: AgentSetDF | list[AgentSetDF], inplace: bool = True) -> Self + add(self, other: AgentSetDF | Iterable[AgentSetDF], inplace: bool = True) -> Self Add agents to the AgentsDF. - contains(self, ids: Hashable | Collection[Hashable]) -> bool | dict[str, pd.Series] + contains(self, ids: IdsLike) -> bool | pl.Series Check if agents with the specified IDs are in the AgentsDF. copy(self, deep: bool = False, memo: dict | None = None) -> Self Create a copy of the AgentsDF. - discard(self, ids: MaskLike, inplace: bool = True) -> Self + discard(self, ids: IdsLike, inplace: bool = True) -> Self Remove an agent from the AgentsDF. Does not raise an error if the agent is not found. do(self, method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any Invoke a method on the AgentsDF. - get(self, attr_names: str | Collection[str] | None = None, mask: MaskLike = None) -> dict[str, pd.Series] | dict[str, pd.DataFrame] + get(self, attr_names: str | Collection[str] | None = None, mask: MaskLike = None) -> dict[str, Series] | dict[str, DataFrame] Retrieve the value of a specified attribute for each agent in the AgentsDF. - remove(self, ids: MaskLike, inplace: bool = True) -> Self + remove(self, ids: IdsLike, inplace: bool = True) -> Self Remove agents from the AgentsDF. select(self, mask: MaskLike = None, filter_func: Callable[[Self], MaskLike] | None = None, n: int | None = None, negate: bool = False, inplace: bool = True) -> Self Select agents in the AgentsDF based on the given criteria. @@ -60,8 +66,14 @@ class AgentsDF(AgentContainer): Shuffle the order of agents in the AgentsDF. sort(self, by: str | Sequence[str], ascending: bool | Sequence[bool] = True, inplace: bool = True, **kwargs) -> Self Sort the agents in the AgentsDF based on the given criteria. + _check_ids(self, other: AgentSetDF | Iterable[AgentSetDF]) -> None + Check if the IDs of the agents to be added are unique. + __add__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self + Add AgentSetDFs to a new AgentsDF through the + operator. __getattr__(self, key: str) -> Any Retrieve an attribute of the underlying agent sets. + __iadd__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self + Add AgentSetDFs to the AgentsDF through the += operator. __iter__(self) -> Iterator Get an iterator for the agents in the AgentsDF. __len__(self) -> int @@ -76,8 +88,11 @@ class AgentsDF(AgentContainer): def __init__(self) -> None: self._agentsets = [] + self._ids = pl.Series(name="unique_id", dtype=pl.Int64) - def add(self, other: AgentSetDF | list[AgentSetDF], inplace: bool = True) -> Self: + def add( + self, other: AgentSetDF | Iterable[AgentSetDF], inplace: bool = True + ) -> Self: """Add an AgentSetDF to the AgentsDF. Parameters @@ -93,23 +108,24 @@ def add(self, other: AgentSetDF | list[AgentSetDF], inplace: bool = True) -> Sel The updated AgentsDF. """ obj = self._get_obj(inplace) - if isinstance(other, list): + self._check_ids(other) + if isinstance(other, Iterable): obj._agentsets += other else: obj._agentsets.append(other) return self @overload - def contains(self, ids: Collection[Hashable]) -> BoolSeries: ... + def contains(self, ids: int) -> bool: ... @overload - def contains(self, ids: Hashable) -> bool: ... + def contains(self, ids: IdsLike) -> pl.Series: ... - def contains(self, ids: Hashable | Collection[Hashable]) -> bool | BoolSeries: - bool_series = self._agentsets[0].contains(ids) - for agentset in self._agentsets[1:]: - bool_series = bool_series | agentset.contains(ids) - return bool_series + def contains(self, ids: IdsLike) -> bool | pl.Series: + if isinstance(ids, int): + return ids in self._ids + else: + return pl.Series(ids).is_in(self._ids) @overload def do( @@ -174,7 +190,7 @@ def get( for agentset in self._agentsets } - def remove(self, ids: MaskLike, inplace: bool = True) -> Self: + def remove(self, ids: IdsLike, inplace: bool = True) -> Self: obj = self._get_obj(inplace) deleted = 0 for agentset in obj._agentsets: @@ -204,7 +220,7 @@ def set( def select( self, mask: MaskLike | None = None, - filter_func: Callable[[DataFrame], MaskLike] | None = None, + filter_func: Callable[[AgentSetDF], MaskLike] | None = None, n: int | None = None, inplace: bool = True, negate: bool = False, @@ -237,12 +253,35 @@ def sort( ] return obj - def __add__(self, other: AgentSetDF | list[AgentSetDF]) -> Self: + def _check_ids(self, other: AgentSetDF | Iterable[AgentSetDF]) -> None: + """Check if the IDs of the agents to be added are unique. + + Parameters + ---------- + other : AgentSetDF | Iterable[AgentSetDF] + The AgentSetDFs to check. + + Raises + ------ + ValueError + If the agent set contains IDs already present in agents. + """ + for agentset in other if isinstance(other, Iterable) else [other]: + if isinstance(agentset, AgentSetPandas): + new_ids = pl.Series(agentset._agents.index) + elif isinstance(agentset, AgentSetPolars): + new_ids = agentset._agents["unique_id"] + if new_ids.is_in(self._ids).any(): + raise ValueError( + "The agent set contains IDs already present in agents." + ) + + def __add__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self: """Add AgentSetDFs to a new AgentsDF through the + operator. Parameters ---------- - other : AgentSetDF | list[AgentSetDF] + other : AgentSetDF | Iterable[AgentSetDF] The AgentSetDFs to add. Returns @@ -258,12 +297,12 @@ def __getattr__(self, name: str) -> dict[str, Any]: for agentset in self._agentsets } - def __iadd__(self, other: AgentSetDF | list[AgentSetDF]) -> Self: + def __iadd__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self: """Add AgentSetDFs to the AgentsDF through the += operator. Parameters ---------- - other : Self | AgentSetDF | list[AgentSetDF] + other : Self | AgentSetDF | Iterable[AgentSetDF] The AgentSetDFs to add. Returns diff --git a/mesa_frames/concrete/agentset_pandas.py b/mesa_frames/concrete/agentset_pandas.py index d6481f3..d2e7c13 100644 --- a/mesa_frames/concrete/agentset_pandas.py +++ b/mesa_frames/concrete/agentset_pandas.py @@ -15,7 +15,7 @@ from mesa_frames.abstract.agents import AgentSetDF from mesa_frames.concrete.agentset_polars import AgentSetPolars from mesa_frames.concrete.model import ModelDF -from mesa_frames.types import PandasMaskLike +from mesa_frames.types import PandasIdsLike, PandasMaskLike class AgentSetPandas(AgentSetDF): @@ -62,17 +62,17 @@ class AgentSetPandas(AgentSetDF): Initialize a new AgentSetPandas. add(self, other: pd.DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True) -> Self Add agents to the AgentSetPandas. - contains(self, ids: Hashable | Collection[Hashable]) -> bool | pd.Series + contains(self, ids: PandasIdsLike) -> bool | pd.Series Check if agents with the specified IDs are in the AgentSetPandas. copy(self, deep: bool = False, memo: dict | None = None) -> Self Create a copy of the AgentSetPandas. - discard(self, ids: PandasMaskLike, inplace: bool = True) -> Self + discard(self, ids: PandasIdsLike, inplace: bool = True) -> Self Remove an agent from the AgentSetPandas. Does not raise an error if the agent is not found. do(self, method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any Invoke a method on the AgentSetPandas. get(self, attr_names: str | Collection[str] | None, mask: PandasMaskLike = None) -> pd.Series | pd.DataFrame Retrieve the value of a specified attribute for each agent in the AgentSetPandas. - remove(self, ids: PandasMaskLike, inplace: bool = True) -> Self + remove(self, ids: PandasIdsLike, inplace: bool = True) -> Self Remove agents from the AgentSetPandas. select(self, mask: PandasMaskLike = None, filter_func: Callable[[Self], PandasMaskLike] | None = None, n: int | None = None, negate: bool = False, inplace: bool = True) -> Self Select agents in the AgentSetPandas based on the given criteria. @@ -104,8 +104,10 @@ class AgentSetPandas(AgentSetDF): def __init__(self, model: ModelDF) -> None: self._model = model - self._agents = pd.DataFrame(columns=["unique_id"]).set_index("unique_id") - self._mask = pd.Series(True, index=self._agents.index) + self._agents = pd.DataFrame( + columns=["unique_id"], dtype={"unique_id": pd.Int64Dtype} + ).set_index("unique_id") + self._mask = pd.Series(True, index=self._agents.index, dtype=pd.BooleanDtype) def add( self, @@ -140,17 +142,21 @@ def add( return obj @overload - def contains(self, ids: Collection[Hashable]) -> pd.Series: ... + def contains(self, ids: int) -> bool: ... @overload - def contains(self, ids: Hashable) -> bool: ... + def contains(self, ids: PandasIdsLike) -> pd.Series: ... def contains( self, - ids: Hashable | Collection[Hashable], + ids: int | Collection[int] | pd.Series[int] | pd.Index, ) -> bool | pd.Series: if isinstance(ids, pd.Series): return ids.isin(self._agents.index) + elif isinstance(ids, pd.Index): + return pd.Series( + ids.isin(self._agents.index), index=ids, dtype=pd.BooleanDtype + ) elif isinstance(ids, Collection): return pd.Series(list(ids), index=list(ids)).isin(self._agents.index) else: @@ -172,7 +178,7 @@ def get( def remove( self, - ids: PandasMaskLike, + ids: PandasIdsLike, inplace: bool = True, ) -> Self: obj = self._get_obj(inplace) diff --git a/mesa_frames/concrete/agentset_polars.py b/mesa_frames/concrete/agentset_polars.py index a4939d2..1b02094 100644 --- a/mesa_frames/concrete/agentset_polars.py +++ b/mesa_frames/concrete/agentset_polars.py @@ -16,7 +16,7 @@ from polars.type_aliases import IntoExpr from mesa_frames.abstract.agents import AgentSetDF -from mesa_frames.types import PolarsMaskLike +from mesa_frames.types import PolarsIdsLike, PolarsMaskLike if TYPE_CHECKING: from mesa_frames.concrete.agentset_pandas import AgentSetPandas @@ -70,17 +70,17 @@ class AgentSetPolars(AgentSetDF): Initialize a new AgentSetPolars. add(self, other: pl.DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True) -> Self Add agents to the AgentSetPolars. - contains(self, ids: Hashable | Collection[Hashable]) -> bool | pl.Series + contains(self, ids: PolarsIdsLike) -> bool | pl.Series Check if agents with the specified IDs are in the AgentSetPolars. copy(self, deep: bool = False, memo: dict | None = None) -> Self Create a copy of the AgentSetPolars. - discard(self, ids: PolarsMaskLike, inplace: bool = True) -> Self + discard(self, ids: PolarsIdsLike, inplace: bool = True) -> Self Remove an agent from the AgentSetPolars. Does not raise an error if the agent is not found. do(self, method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any Invoke a method on the AgentSetPolars. get(self, attr_names: IntoExpr | Iterable[IntoExpr] | None, mask: PolarsMaskLike = None) -> pl.Series | pl.DataFrame Retrieve the value of a specified attribute for each agent in the AgentSetPolars. - remove(self, ids: PolarsMaskLike, inplace: bool = True) -> Self + remove(self, ids: PolarsIdsLike, inplace: bool = True) -> Self Remove agents from the AgentSetPolars. select(self, mask: PolarsMaskLike = None, filter_func: Callable[[Self], PolarsMaskLike] | None = None, n: int | None = None, negate: bool = False, inplace: bool = True) -> Self Select agents in the AgentSetPolars based on the given criteria. @@ -165,14 +165,14 @@ def add( return obj @overload - def contains(self, ids: Collection[Hashable]) -> pl.Series: ... + def contains(self, ids: int) -> bool: ... @overload - def contains(self, ids: Hashable) -> bool: ... + def contains(self, ids: PolarsIdsLike) -> pl.Series: ... def contains( self, - ids: Hashable | Collection[Hashable], + ids: PolarsIdsLike, ) -> bool | pl.Series: if isinstance(ids, pl.Series): return ids.is_in(self._agents["unique_id"]) @@ -181,7 +181,7 @@ def contains( else: return ids in self._agents["unique_id"] - def discard(self, ids: PolarsMaskLike, inplace: bool = True) -> Self: + def discard(self, ids: PolarsIdsLike, inplace: bool = True) -> Self: """Remove an agent from the AgentSetPolars. Does not raise an error if the agent is not found. Parameters diff --git a/mesa_frames/types.py b/mesa_frames/types.py index 8ca8726..a5da994 100644 --- a/mesa_frames/types.py +++ b/mesa_frames/types.py @@ -2,7 +2,7 @@ ####----- Agnostic Types -----#### AgnosticMask = Literal["all", "active"] | Hashable | None - +AgnosticIds = int | Collection[int] ###----- Pandas Types -----### import pandas as pd @@ -11,17 +11,18 @@ ArrayLike = pd.api.extensions.ExtensionArray | ndarray AnyArrayLike = ArrayLike | pd.Index | pd.Series PandasMaskLike = AgnosticMask | pd.Series | pd.DataFrame | AnyArrayLike - +PandasIdsLike = AgnosticIds | pd.Series[int] | pd.Index ###----- Polars Types -----### import polars as pl PolarsMaskLike = AgnosticMask | pl.Expr | pl.Series | pl.DataFrame | Collection[int] - +PolarsIdsLike = AgnosticIds | pl.Series ###----- Generic -----### DataFrame = pd.DataFrame | pl.DataFrame Series = pd.Series | pl.Series -BoolSeries = pd.Series | pl.Series | pl.Expr +BoolSeries = pd.Series | pl.Series MaskLike = AgnosticMask | PandasMaskLike | PolarsMaskLike +IdsLike = AgnosticIds | pd.Series[int] | pd.Index | pl.Series