diff --git a/docs/scripts/readme_plot.py b/docs/scripts/readme_plot.py index 2cca434..31ecacc 100644 --- a/docs/scripts/readme_plot.py +++ b/docs/scripts/readme_plot.py @@ -1,16 +1,15 @@ -from hmac import new -from random import seed -from typing import TYPE_CHECKING - +import matplotlib.pyplot as plt import mesa import numpy as np import pandas as pd import perfplot +import polars as pl +import seaborn as sns -from mesa_frames import AgentSetDF, ModelDF +from mesa_frames import AgentSetPandas, AgentSetPolars, ModelDF -# Mesa implementation +### ---------- Mesa implementation ---------- ### def mesa_implementation(n_agents: int) -> None: model = MoneyModel(n_agents) model.run_model(100) @@ -39,7 +38,7 @@ class MoneyModel(mesa.Model): """A model with some number of agents.""" def __init__(self, N): - super().__init__ + super().__init__() self.num_agents = N # Create scheduler and assign it to the model self.schedule = mesa.time.RandomActivation(self) @@ -69,58 +68,178 @@ def run_model(self, n_steps) -> None: return 1 + (1 / N) - 2 * B""" -# Mesa Frames implementation -def mesa_frames_implementation(n_agents: int) -> None: - model = MoneyModelDF(n_agents) - model.run_model(100) +### ---------- Mesa-frames implementation ---------- ### + + +class MoneyAgentPolars(AgentSetPolars): + def __init__(self, n: int, model: ModelDF): + super().__init__(model) + ## Adding the agents to the agent set + # 1. Changing the agents attribute directly (not reccomended, if other agents were added before, they will be lost) + """self.agents = pl.DataFrame( + {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)} + )""" + # 2. Adding the dataframe with add + """self.add( + pl.DataFrame( + { + "unique_id": pl.arange(n, eager=True), + "wealth": pl.ones(n, eager=True), + } + ) + )""" + # 3. Adding the dataframe with __iadd__ + self += pl.DataFrame( + {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)} + ) + + def step(self) -> None: + # The give_money method is called + # self.give_money() + self.do("give_money") + + def give_money(self): + ## Active agents are changed to wealthy agents + # 1. Using a native expression + # self.select(pl.col("wealth") > 0) + # 2. Using the __getitem__ method + # self.select(self["wealth"] > 0) + # 3. Using the fallback __getattr__ method + self.select(self.wealth > 0) + + # Receiving agents are sampled (only native expressions currently supported) + other_agents = self.agents.sample( + n=len(self.active_agents), with_replacement=True + ) + + # Wealth of wealthy is decreased by 1 + # 1. Using a native expression + """self.agents = self.agents.with_columns( + wealth=pl.when(pl.col("unique_id").is_in(self.active_agents["unique_id"])) + .then(pl.col("wealth") - 1) + .otherwise(pl.col("wealth")) + )""" + # 2. Using the __setitem__ method with self.active_agents mask + # self[self.active_agents, "wealth"] -= 1 + # 3. Using the __setitem__ method with "active" mask + self["active", "wealth"] -= 1 + + # Compute the income of the other agents (only native expressions currently supported) + new_wealth = other_agents.group_by("unique_id").len() + + # Add the income to the other agents + # 1. Using native expressions + """self.agents = self.agents.with_columns( + pl.when(pl.col("unique_id").is_in(new_wealth["unique_id"])) + .then(pl.col("wealth") + new_wealth["wealth"]) + .otherwise(pl.col("wealth")) + )""" + + # 2. Using the set method + """self.set( + attr_names="wealth", + values=pl.col("wealth") + new_wealth["len"], + mask=new_wealth, + )""" + + # 3. Using the __setitem__ method + self[new_wealth, "wealth"] += new_wealth["len"] + + +class MoneyAgentPandas(AgentSetPandas): + def __init__(self, n: int, model: ModelDF) -> None: + super().__init__(model) + ## Adding the agents to the agent set + # 1. Changing the agents attribute directly (not reccomended, if other agents were added before, they will be lost) + # self.agents = pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)}) + # 2. Adding the dataframe with add + # self.add(pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)})) + # 3. Adding the dataframe with __iadd__ + self += pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)}) + + def step(self) -> None: + # The give_money method is called + self.do("give_money") + + def give_money(self): + ## Active agents are changed to wealthy agents + # 1. Using a native expression + # self.select(self.agents['wealth'] > 0) + # 2. Using the __getitem__ method + # self.select(self["wealth"] > 0) + # 3. Using the fallback __getattr__ method + self.select(self.wealth > 0) + + # Receiving agents are sampled (only native expressions currently supported) + other_agents = self.agents.sample(n=len(self.active_agents), replace=True) + + # Wealth of wealthy is decreased by 1 + # 1. Using a native expression + """b_mask = self.active_agents.index.isin(self.agents) + self.agents.loc[b_mask, "wealth"] -= 1""" + # 2. Using the __setitem__ method with self.active_agents mask + # self[self.active_agents, "wealth"] -= 1 + # 3. Using the __setitem__ method with "active" mask + self["active", "wealth"] -= 1 + + # Compute the income of the other agents (only native expressions currently supported) + new_wealth = other_agents.groupby("unique_id").count() + + # Add the income to the other agents + # 1. Using native expressions + """merged = pd.merge( + self.agents, new_wealth, on="unique_id", how="left", suffixes=("", "_new") + ) + merged["wealth"] = merged["wealth"] + merged["wealth_new"].fillna(0) + self.agents = merged.drop(columns=["wealth_new"])""" + + # 2. Using the set method + # self.set(attr_names="wealth", values=self["wealth"] + new_wealth["wealth"], mask=new_wealth) + + # 3. Using the __setitem__ method + self[new_wealth, "wealth"] += new_wealth["wealth"] class MoneyModelDF(ModelDF): - def __init__(self, N): + def __init__(self, N: int, agents_cls): super().__init__() - self.num_agents = N - self.agents = self.agents.add(MoneyAgentsDF(N, model=self)) + self.n_agents = N + self.agents += agents_cls(N, self) def step(self): - self.agents = self.agents.do("step") + # Executes the step method for every agentset in self.agents + self.agents.do("step") def run_model(self, n): for _ in range(n): self.step() -class MoneyAgentsDF(AgentSetDF): - def __init__(self, n: int, model: MoneyModelDF): - super().__init__(model=model) - self.add(n, data={"wealth": np.ones(n)}) +def mesa_frames_polars(n_agents: int) -> None: + model = MoneyModelDF(n_agents, MoneyAgentPolars) + model.run_model(100) - def step(self): - wealthy_agents = self.agents["wealth"] > 0 - self.select(wealthy_agents).do("give_money") - def give_money(self): - other_agents = self.agents.sample(len(self.active_agents), replace=True) - new_wealth = ( - other_agents.index.value_counts() - .reindex(self.active_agents.index) - .fillna(-1) - ) - self.set_attribute("wealth", self.get_attribute("wealth") + new_wealth) +def mesa_frames_pandas(n_agents: int) -> None: + model = MoneyModelDF(n_agents, MoneyAgentPandas) + model.run_model(100) def main(): - mesa_frames_implementation(100) + + sns.set_theme(style="whitegrid") + out = perfplot.bench( setup=lambda n: n, - kernels=[mesa_implementation, mesa_frames_implementation], - labels=["mesa", "mesa-frames"], - n_range=[k for k in range(100, 1000, 100)], + kernels=[mesa_implementation, mesa_frames_polars, mesa_frames_pandas], + labels=["mesa", "mesa-frames (polars)", "mesa-frames (pandas)"], + n_range=[k for k in range(100, 10000, 1000)], xlabel="Number of agents", equality_check=None, - title="100 steps of the Boltzmann Wealth model", + title="100 steps of the Boltzmann Wealth model: mesa vs mesa-frames (polars) vs mesa-frames (pandas)", ) - out.show() - # out.save("docs/images/readme_plot.png") + plt.ylabel("Execution time (s)") + out.save("docs/images/readme_plot_1.png") if __name__ == "__main__": diff --git a/mesa_frames/__init__.py b/mesa_frames/__init__.py index c48f87d..cad5fbe 100644 --- a/mesa_frames/__init__.py +++ b/mesa_frames/__init__.py @@ -1,2 +1,4 @@ -from mesa_frames.agent import AgentSetPandas, AgentSetPolars, AgentsPandas, AgentsPolars -from mesa_frames.model import ModelDF +from mesa_frames.concrete.agents import AgentsDF +from mesa_frames.concrete.agentset_pandas import AgentSetPandas +from mesa_frames.concrete.agentset_polars import AgentSetPolars +from mesa_frames.concrete.model import ModelDF diff --git a/mesa_frames/_decorator.py b/mesa_frames/_decorator.py deleted file mode 100644 index 248429e..0000000 --- a/mesa_frames/_decorator.py +++ /dev/null @@ -1,51 +0,0 @@ -from textwrap import dedent -from typing import Callable, Union - - -# doc decorator function ported with modifications from Pandas -# https://github.com/pandas-dev/pandas/blob/master/pandas/util/_decorators.py - - -def doc(*docstrings: Union[str, Callable], **params) -> Callable: - """ - A decorator take docstring templates, concatenate them and perform string - substitution on it. - This decorator will add a variable "_docstring_components" to the wrapped - callable to keep track the original docstring template for potential usage. - If it should be consider as a template, it will be saved as a string. - Otherwise, it will be saved as callable, and later user __doc__ and dedent - to get docstring. - - Parameters - ---------- - *docstrings : str or callable - The string / docstring / docstring template to be appended in order - after default docstring under callable. - **params - The string which would be used to format docstring template. - """ - - def decorator(decorated: Callable) -> Callable: - # collecting docstring and docstring templates - docstring_components: list[Union[str, Callable]] = [] - if decorated.__doc__: - docstring_components.append(dedent(decorated.__doc__)) - - for docstring in docstrings: - if hasattr(docstring, "_docstring_components"): - docstring_components.extend(docstring._docstring_components) - elif isinstance(docstring, str) or docstring.__doc__: - docstring_components.append(docstring) - - # formatting templates and concatenating docstring - decorated.__doc__ = "".join( - component.format(**params) - if isinstance(component, str) - else dedent(component.__doc__ or "") - for component in docstring_components - ) - - decorated._docstring_components = docstring_components - return decorated - - return decorator \ No newline at end of file diff --git a/mesa_frames/base/__init__.py b/mesa_frames/abstract/__init__.py similarity index 100% rename from mesa_frames/base/__init__.py rename to mesa_frames/abstract/__init__.py diff --git a/mesa_frames/abstract/agents.py b/mesa_frames/abstract/agents.py new file mode 100644 index 0000000..9edec7c --- /dev/null +++ b/mesa_frames/abstract/agents.py @@ -0,0 +1,1041 @@ + + +from __future__ import annotations # PEP 563: postponed evaluation of type annotations + +from abc import ABC, abstractmethod +from contextlib import suppress +from copy import copy, deepcopy +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Hashable, + Iterator, + Literal, + Self, + Sequence, + overload, +) + +from numpy.random import Generator + +from mesa_frames.types import BoolSeries, DataFrame, MaskLike, Series + +if TYPE_CHECKING: + from mesa_frames.concrete.model import ModelDF + + +class AgentContainer(ABC): + """An abstract class for containing agents. Defines the common interface for AgentSetDF and AgentsDF. + + Attributes + ---------- + _copy_with_method : dict[str, tuple[str, list[str]]] + A dictionary of attributes to copy with a specified method and arguments. + _copy_only_reference : list[str] + A list of attributes to copy with a reference only. + _model : ModelDF + The model that the AgentContainer belongs to. + + Methods + ------- + copy(deep: bool = False, memo: dict | None = None) -> Self + Create a copy of the AgentContainer. + discard(ids: MaskLike, 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 + 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 + 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. + set(attr_names: str | dict[str, Any] | Collection[str], values: Any | None = None, mask: MaskLike | None = None, inplace: bool = True) -> Self + Sets the value of a specified attribute or attributes for each agent in the mask in AgentContainer. + shuffle(inplace: bool = False) -> Self + Shuffles the order of agents in the AgentContainer. + sort(by: str | Sequence[str], ascending: bool | Sequence[bool] = True, inplace: bool = True, **kwargs) -> Self + Sorts the agents in the agent set based on the given criteria. + + Properties + ---------- + model : ModelDF + Get the model associated with the AgentContainer. + random : Generator + Get the random number generator associated with the model. + agents : DataFrame | dict[str, DataFrame] + Get or set the agents in the AgentContainer. + active_agents : DataFrame | dict[str, DataFrame] + Get or set the active agents in the AgentContainer. + inactive_agents : DataFrame | dict[str, DataFrame] + Get the inactive agents in the AgentContainer. + """ + + _copy_with_method: dict[str, tuple[str, list[str]]] + _copy_only_reference: list[str] = [ + "_model", + ] + _model: ModelDF + + @abstractmethod + def __init__(self) -> None: ... + + def copy( + self, + deep: bool = False, + memo: dict | None = None, + ) -> Self: + """Create a copy of the AgentContainer. + + Parameters + ---------- + deep : bool, optional + Flag indicating whether to perform a deep copy of the AgentContainer. + If True, all attributes of the AgentContainer will be recursively copied (except attributes in self._copy_reference_only). + If False, only the top-level attributes will be copied. + Defaults to False. + + memo : dict | None, optional + A dictionary used to track already copied objects during deep copy. + Defaults to None. + + Returns + ------- + Self + A new instance of the AgentContainer class that is a copy of the original instance. + """ + cls = self.__class__ + obj = cls.__new__(cls) + + if deep: + if not memo: + memo = {} + memo[id(self)] = obj + attributes = self.__dict__.copy() + [ + setattr(obj, k, deepcopy(v, memo)) + for k, v in attributes.items() + if k not in self._copy_with_method + and k not in self._copy_only_reference + ] + else: + [ + setattr(obj, k, copy(v)) + for k, v in self.__dict__.items() + if k not in self._copy_with_method + and k not in self._copy_only_reference + ] + + # Copy attributes with a reference only + for attr in self._copy_only_reference: + setattr(obj, attr, getattr(self, attr)) + + # Copy attributes with a specified method + for attr in self._copy_with_method: + attr_obj = getattr(self, attr) + attr_copy_method, attr_copy_args = self._copy_with_method[attr] + setattr(obj, attr, getattr(attr_obj, attr_copy_method)(*attr_copy_args)) + + return obj + + def discard(self, ids: MaskLike, inplace: bool = True) -> Self: + """Removes an agent from the AgentContainer. Does not raise an error if the agent is not found. + + Parameters + ---------- + ids : MaskLike + The MaskLike of the agents to remove. + inplace : bool + Whether to remove the agent in place. Defaults to False. + + Returns + ---------- + Self + """ + with suppress(KeyError): + return self.remove(ids, inplace=inplace) + return self._get_obj(inplace) + + @abstractmethod + def add(self, other, inplace: bool = True) -> Self: + """Add agents to the AgentContainer. + + Parameters + ---------- + other : Any + The agents to add. + inplace : bool + Whether to add the agents in place. Defaults to True. + + Returns + ------- + Self + The updated AgentContainer. + """ + ... + + @overload + @abstractmethod + def contains(self, ids: Collection[Hashable]) -> BoolSeries: ... + + @overload + @abstractmethod + def contains(self, ids: Hashable) -> bool: ... + + @abstractmethod + def contains(self, ids: Hashable | Collection[Hashable]) -> bool | BoolSeries: + """Check if agents with the specified IDs are in the AgentContainer. + + Parameters + ---------- + ids : Hashable | Collection[Any] + The ID(s) to check for. + + Returns + ------- + bool | BoolSeries + True if the agent is in the AgentContainer, False otherwise. + """ + + @overload + @abstractmethod + def do( + self, + method_name: str, + *args, + return_results: Literal[False] = False, + inplace: bool = True, + **kwargs, + ) -> Self: ... + + @overload + @abstractmethod + def do( + self, + method_name: str, + *args, + return_results: Literal[True], + inplace: bool = True, + **kwargs, + ) -> Any | dict[str, Any]: ... + + @abstractmethod + def do( + self, + method_name: str, + *args, + return_results: bool = False, + inplace: bool = True, + **kwargs, + ) -> Self | Any | dict[str, Any]: + """Invoke a method on the AgentContainer. + + Parameters + ---------- + method_name : str + The name of the method to invoke. + return_results : bool, optional + Whether to return the result of the method, by default False + inplace : bool, optional + Whether the operation should be done inplace, by default False + + Returns + ------- + Self | Any + The updated AgentContainer or the result of the method. + """ + ... + + @abstractmethod + @overload + def get(self, attr_names: str) -> Series | dict[str, Series]: ... + + @abstractmethod + @overload + def get(self, attr_names: Collection[str]) -> DataFrame | dict[str, DataFrame]: ... + + @abstractmethod + def get( + self, + attr_names: str | Collection[str] | None = None, + mask: MaskLike | None = None, + ) -> Series | DataFrame | dict[str, Series] | dict[str, DataFrame]: + """Retrieves the value of a specified attribute for each agent in the AgentContainer. + + Parameters + ---------- + attr_names : str | Collection[str] | None + The attributes to retrieve. If None, all attributes are retrieved. Defaults to None. + MaskLike : MaskLike | None + The MaskLike of agents to retrieve the attribute for. If None, attributes of all agents are returned. Defaults to None. + + Returns + ---------- + Series | DataFrame | dict[str, Series | DataFrame] + The attribute values. + """ + ... + + @abstractmethod + def remove(self, ids: MaskLike, inplace: bool = True) -> Self: + """Removes an agent from the AgentContainer. + + Parameters + ---------- + id : MaskLike + The ID of the agent to remove. + inplace : bool + Whether to remove the agent in place. + + Returns + ---------- + Self + The updated AgentContainer. + """ + ... + + @abstractmethod + def select( + self, + 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. + + Parameters + ---------- + mask : MaskLike | None, optional + The MaskLike of agents to be selected, by default None + filter_func : Callable[[Self], MaskLike] | None, optional + A function which takes as input the AgentContainer and returns a MaskLike, by default None + n : int, optional + The maximum number of agents to be selected, by default None + negate : bool, optional + If the selection should be negated, by default False + inplace : bool, optional + If the operation should be performed on the same object, by default True + + Returns + ------- + Self + A new or updated AgentContainer. + """ + ... + + @abstractmethod + @overload + def set( + self, + attr_names: dict[str, Any], + values: None, + mask: MaskLike | None = None, + inplace: bool = True, + ) -> Self: ... + + @abstractmethod + @overload + def set( + self, + attr_names: str | Collection[str], + values: Any, + mask: MaskLike | None = None, + inplace: bool = True, + ) -> Self: ... + + @abstractmethod + def set( + self, + attr_names: str | dict[str, Any] | Collection[str], + values: Any | None = None, + mask: MaskLike | None = None, + inplace: bool = True, + ) -> Self: + """Sets the value of a specified attribute or attributes for each agent in the mask in AgentContainer. + + Parameters + ---------- + attr_names : str | dict[str, Any] | Collection[str] | None + The key can be: + - A string: sets the specified column of the agents in the AgentContainer. + - A collection of strings: sets the specified columns of the agents in the AgentContainer. + - A dictionary: keys should be attributes and values should be the values to set. Value should be None. + value : Any | None + The value to set the attribute to. If None, attr_names must be a dictionary. + mask : MaskLike | None + The MaskLike of agents to set the attribute for. + inplace : bool + Whether to set the attribute in place. + + Returns + ---------- + Self + The updated agent set. + """ + ... + + @abstractmethod + def shuffle(self, inplace: bool = False) -> Self: + """Shuffles the order of agents in the AgentContainer. + + Parameters + ---------- + inplace : bool + Whether to shuffle the agents in place. + + Returns + ---------- + Self + A new or updated AgentContainer. + """ + + @abstractmethod + def sort( + self, + by: str | Sequence[str], + ascending: bool | Sequence[bool] = True, + inplace: bool = True, + **kwargs, + ) -> Self: + """ + Sorts the agents in the agent set based on the given criteria. + + Parameters + ---------- + by : str | Sequence[str] + The attribute(s) to sort by. + ascending : bool | Sequence[bool] + Whether to sort in ascending order. + inplace : bool + Whether to sort the agents in place. + **kwargs + Keyword arguments to pass to the sort + + Returns + ---------- + Self + A new or updated AgentContainer. + """ + + def _get_obj(self, inplace: bool) -> Self: + """Get the object to perform operations on. + + Parameters + ---------- + inplace : bool + If inplace, return self. Otherwise, return a copy. + + Returns + ---------- + Self + The object to perform operations on. + """ + if inplace: + return self + else: + return deepcopy(self) + + def __add__(self, other) -> Self: + return self.add(other=other, inplace=False) + + def __contains__(self, id: Hashable) -> bool: + """Check if an agent is in the AgentContainer. + + Parameters + ---------- + id : Hashable + The ID(s) to check for. + + Returns + ------- + 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.") + + def __copy__(self) -> Self: + """Create a shallow copy of the AgentContainer. + + Returns + ------- + Self + A shallow copy of the AgentContainer. + """ + return self.copy(deep=False) + + def __deepcopy__(self, memo: dict) -> Self: + """Create a deep copy of the AgentContainer. + + Parameters + ---------- + memo : dict + A dictionary to store the copied objects. + + Returns + ------- + Self + A deep copy of the AgentContainer. + """ + return self.copy(deep=True, memo=memo) + + def __getitem__( + self, + key: ( + str + | Collection[str] + | MaskLike + | tuple[MaskLike, str] + | tuple[MaskLike, Collection[str]] + ), + ) -> Series | DataFrame | dict[str, Series] | dict[str, DataFrame]: + """Implements the [] operator for the AgentContainer. + + The key can be: + - An attribute or collection of attributes (eg. AgentContainer["str"], AgentContainer[["str1", "str2"]]): returns the specified column(s) of the agents in the AgentContainer. + - A MaskLike (eg. AgentContainer[MaskLike]): returns the agents in the AgentContainer that satisfy the MaskLike. + - A tuple (eg. AgentContainer[MaskLike, "str"]): returns the specified column of the agents in the AgentContainer that satisfy the MaskLike. + + Parameters + ---------- + key : Attributes | MaskLike | tuple[MaskLike, Attributes] + The key to retrieve. + + Returns + ------- + Series | DataFrame + The attribute values. + """ + # TODO: fix types + if isinstance(key, tuple): + return self.get(mask=key[0], attr_names=key[1]) + else: + try: + return self.get(attr_names=key) + except: + return self.get(mask=key) + + def __iadd__(self, other) -> Self: + """Add agents to the AgentContainer through the += operator. + + Parameters + ---------- + other + The agents to add. + + Returns + ------- + Self + The updated AgentContainer. + """ + return self.add(other=other, inplace=True) + + def __isub__(self, other: MaskLike) -> Self: + """Remove agents from the AgentContainer through the -= operator. + + Parameters + ---------- + other : MaskLike + The agents to remove. + + Returns + ------- + Self + The updated AgentContainer. + """ + return self.discard(other, inplace=True) + + def __sub__(self, other: MaskLike) -> Self: + """Remove agents from a new AgentContainer through the - operator. + + Parameters + ---------- + other : DataFrame | ListLike + The agents to remove. + + Returns + ------- + Self + A new AgentContainer with the removed agents. + """ + return self.discard(other, inplace=False) + + def __setitem__( + self, + key: str | Collection[str] | MaskLike | tuple[MaskLike, str | Collection[str]], + values: Any, + ) -> None: + """Implement the [] operator for setting values in the AgentContainer. + + The key can be: + - A string (eg. AgentContainer["str"]): sets the specified column of the agents in the AgentContainer. + - A list of strings(eg. AgentContainer[["str1", "str2"]]): sets the specified columns of the agents in the AgentContainer. + - A tuple (eg. AgentContainer[MaskLike, "str"]): sets the specified column of the agents in the AgentContainer that satisfy the MaskLike. + - A MaskLike (eg. AgentContainer[MaskLike]): sets the attributes of the agents in the AgentContainer that satisfy the MaskLike. + + Parameters + ---------- + key : str | list[str] | MaskLike | tuple[MaskLike, str | list[str]] + The key to set. + values : Any + The values to set for the specified key. + """ + # TODO: fix types as in __getitem__ + if isinstance(key, tuple): + self.set(mask=key[0], attr_names=key[1], values=values) + else: + if isinstance(key, str) or ( + isinstance(key, Collection) and all(isinstance(k, str) for k in key) + ): + try: + self.set(attr_names=key, values=values) + except: # key=MaskLike + self.set(attr_names=None, mask=key, values=values) + else: + self.set(attr_names=None, mask=key, values=values) + + @abstractmethod + def __getattr__(self, name: str) -> Any | dict[str, Any]: + """Fallback for retrieving attributes of the AgentContainer. Retrieves an attribute of the underlying DataFrame(s). + + Parameters + ---------- + name : str + The name of the attribute to retrieve. + + Returns + ------- + Any | dict[str, Any] + The attribute value + """ + + @abstractmethod + def __iter__(self) -> Iterator: + """Iterate over the agents in the AgentContainer. + + Returns + ------- + Iterator + An iterator over the agents. + """ + ... + + @abstractmethod + def __len__(self) -> int: + """Get the number of agents in the AgentContainer. + + Returns + ------- + int + The number of agents in the AgentContainer. + """ + ... + + @abstractmethod + def __repr__(self) -> str: + """Get a string representation of the DataFrame in the AgentContainer. + + Returns + ------- + str + A string representation of the DataFrame in the AgentContainer. + """ + pass + + @abstractmethod + def __reversed__(self) -> Iterator: + """Iterate over the agents in the AgentContainer in reverse order. + + Returns + ------- + Iterator + An iterator over the agents in reverse order. + """ + ... + + @abstractmethod + def __str__(self) -> str: + """Get a string representation of the agents in the AgentContainer. + + Returns + ------- + str + A string representation of the agents in the AgentContainer. + """ + ... + + @property + def model(self) -> ModelDF: + """The model that the AgentContainer belongs to. + + Returns + ------- + ModelDF + """ + return self._model + + @property + def random(self) -> Generator: + """The random number generator of the model. + + Returns + ------- + Generator""" + return self.model.random + + @property + @abstractmethod + def agents(self) -> DataFrame | dict[str, DataFrame]: + """The agents in the AgentContainer. + + Returns + ------- + DataFrame | dict[str, DataFrame] + """ + + @agents.setter + @abstractmethod + def agents(self, agents: DataFrame | list["AgentSetDF"]) -> None: + """Set the agents in the AgentContainer. + + Parameters + ---------- + agents : DataFrame | list[AgentSetDF] + """ + + @property + @abstractmethod + def active_agents(self) -> DataFrame | dict[str, DataFrame]: + """The active agents in the AgentContainer. + + Returns + ------- + DataFrame + """ + + @active_agents.setter + @abstractmethod + def active_agents( + self, + mask: MaskLike, + ) -> None: + """Set the active agents in the AgentContainer. + + Parameters + ---------- + mask : MaskLike + The mask to apply. + """ + self.select(mask=mask, inplace=True) + + @property + @abstractmethod + def inactive_agents(self) -> DataFrame | dict[str, DataFrame]: + """The inactive agents in the AgentContainer. + + Returns + ------- + DataFrame + """ + + +class AgentSetDF(AgentContainer): + """The AgentSetDF class is a container for agents of the same type. + + Attributes + ---------- + _agents : DataFrame + The agents in the AgentSetDF. + _copy_only_reference : list[str] + A list of attributes to copy with a reference only. + _copy_with_method : dict[str, tuple[str, list[str]]] + A dictionary of attributes to copy with a specified method and arguments. + _mask : MaskLike + The underlying mask used for the active agents in the AgentSetDF. + _model : ModelDF + The model that the AgentSetDF belongs to. + + Methods + ------- + __init__(self, model: ModelDF) -> None + Create a new AgentSetDF. + add(self, other: DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True) -> Self + Add agents to the AgentSetDF. + contains(self, ids: Hashable | Collection[Hashable]) -> bool | BoolSeries + Check if agents with the specified IDs are in the AgentSetDF. + copy(self, deep: bool = False, memo: dict | None = None) -> Self + Create a copy of the AgentSetDF. + discard(self, ids: MaskLike, inplace: bool = True) -> Self + Removes an agent from the AgentSetDF. 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 AgentSetDF. + get(self, attr_names: str | Collection[str] | None = None, mask: MaskLike | None = None) -> Series | DataFrame + Retrieve the value of a specified attribute for each agent in the AgentSetDF. + remove(self, ids: MaskLike, inplace: bool = True) -> Self + Removes an agent from the AgentSetDF. + select(self, 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 AgentSetDF based on the given criteria. + set(self, attr_names: str | dict[str, Any] | Collection[str], values: Any | None = None, mask: MaskLike | None = None, inplace: bool = True) -> Self + Sets the value of a specified attribute or attributes for each agent in the mask in AgentSetDF. + shuffle(self, inplace: bool = False) -> Self + Shuffles the order of agents in the AgentSetDF. + sort(self, by: str | Sequence[str], ascending: bool | Sequence[bool] = True, inplace: bool = True, **kwargs) -> Self + Sorts the agents in the AgentSetDF based on the given criteria. + _get_obj(self, inplace: bool) -> Self + Get the appropriate object, either the current instance or a copy, based on the `inplace` parameter. + __add__(self, other: Self | DataFrame | Sequence[Any] | dict[str, Any]) -> Self + Add agents to a new AgentSetDF through the + operator. + __iadd__(self, other: Self | DataFrame | Sequence[Any] | dict[str, Any]) -> Self + Add agents to the AgentSetDF through the += operator. + __getattr__(self, name: str) -> Any + Retrieve an attribute of the AgentSetDF. + __getitem__(self, key: str | Collection[str] | MaskLike | tuple[MaskLike, str] | tuple[MaskLike, Collection[str]]) -> Series | DataFrame + Retrieve an item from the AgentSetDF. + __iter__(self) -> Iterator + Get an iterator for the agents in the AgentSetDF. + __len__(self) -> int + Get the number of agents in the AgentSetDF. + __repr__(self) -> str + Get the string representation of the AgentSetDF. + __reversed__(self) -> Iterator + Get a reversed iterator for the agents in the AgentSetDF. + __str__(self) -> str + Get the string representation of the AgentSetDF. + + Properties + ---------- + active_agents(self) -> DataFrame + Get the active agents in the AgentSetDF. + agents(self) -> DataFrame + Get or set the agents in the AgentSetDF. + inactive_agents(self) -> DataFrame + Get the inactive agents in the AgentSetDF. + model(self) -> ModelDF + Get the model associated with the AgentSetDF. + random(self) -> Generator + Get the random number generator associated with the model. + """ + + _agents: DataFrame + _mask: MaskLike + _model: ModelDF + + @abstractmethod + def __init__(self, model: ModelDF) -> None: + """Create a new AgentSetDF. + + Parameters + ---------- + model : ModelDF + The model that the agent set belongs to. + + Returns + ------- + None + """ + ... + + @abstractmethod + def add( + self, other: DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True + ) -> Self: + """Add agents to the AgentSetDF + + Other can be: + - A DataFrame: adds the agents from the DataFrame. + - A Sequence[Any]: should be one single agent to add. + - A dictionary: keys should be attributes and values should be the values to add. + + Parameters + ---------- + other : DataFrame | Sequence[Any] | dict[str, Any] + The agents to add. + inplace : bool, optional + If True, perform the operation in place, by default True + + Returns + ------- + Self + A new AgentContainer with the added agents. + """ + ... + + @overload + def do( + self, + method_name: str, + *args, + return_results: Literal[False] = False, + inplace: bool = True, + **kwargs, + ) -> Self: ... + + @overload + def do( + self, + method_name: str, + *args, + return_results: Literal[True], + inplace: bool = True, + **kwargs, + ) -> Any: ... + + def do( + self, + method_name: str, + *args, + return_results: bool = False, + inplace: bool = True, + **kwargs, + ) -> Self | Any: + obj = self._get_obj(inplace) + method = getattr(obj, method_name) + if return_results: + return method(*args, **kwargs) + else: + method(*args, **kwargs) + return obj + + @abstractmethod + @overload + def get( + self, + attr_names: str, + mask: MaskLike | None = None, + ) -> Series: ... + + @abstractmethod + @overload + def get( + self, + attr_names: Collection[str] | None = None, + mask: MaskLike | None = None, + ) -> DataFrame: ... + + @abstractmethod + def get( + self, + attr_names: str | Collection[str] | None = None, + mask: MaskLike | None = None, + ) -> Series | DataFrame: ... + + def __add__(self, other: DataFrame | Sequence[Any] | dict[str, Any]) -> Self: + """Add agents to a new AgentSetDF through the + operator. + + Other can be: + - A DataFrame: adds the agents from the DataFrame. + - A Sequence[Any]: should be one single agent to add. + - A dictionary: keys should be attributes and values should be the values to add. + + Parameters + ---------- + other : DataFrame | Sequence[Any] | dict[str, Any] + The agents to add. + + Returns + ------- + Self + A new AgentContainer with the added agents. + """ + return super().__add__(other) + + def __iadd__(self, other: DataFrame | Sequence[Any] | dict[str, Any]) -> Self: + """ + Add agents to the AgentSetDF through the += operator. + + Other can be: + - A DataFrame: adds the agents from the DataFrame. + - A Sequence[Any]: should be one single agent to add. + - A dictionary: keys should be attributes and values should be the values to add. + + Parameters + ---------- + other : DataFrame | Sequence[Any] | dict[str, Any] + The agents to add. + + Returns + ------- + Self + The updated AgentContainer. + """ + return super().__iadd__(other) + + @abstractmethod + def __getattr__(self, name: str) -> Any: + if name == "_agents": + raise RuntimeError( + "The _agents attribute is not set. You probably forgot to call super().__init__ in the __init__ method." + ) + + @overload + def __getitem__(self, key: str | tuple[MaskLike, str]) -> Series | DataFrame: ... + + @overload + def __getitem__( + self, key: MaskLike | Collection[str] | tuple[MaskLike, Collection[str]] + ) -> DataFrame: ... + + def __getitem__( + self, + key: ( + str + | Collection[str] + | MaskLike + | tuple[MaskLike, str] + | tuple[MaskLike, Collection[str]] + ), + ) -> Series | DataFrame: + attr = super().__getitem__(key) + assert isinstance(attr, (Series, DataFrame)) + return attr + + def __len__(self) -> int: + return len(self._agents) + + def __iter__(self) -> Iterator: + return iter(self._agents) + + def __repr__(self) -> str: + return repr(self._agents) + + def __str__(self) -> str: + return str(self._agents) + + def __reversed__(self) -> Iterator: + return reversed(self._agents) + + @property + def agents(self) -> DataFrame: + return self._agents + + @agents.setter + def agents(self, agents: DataFrame) -> None: + """Set the agents in the AgentSetDF. + + Parameters + ---------- + agents : DataFrame + The agents to set. + """ + self._agents = agents + + @property + @abstractmethod + def active_agents(self) -> DataFrame: ... + + @property + @abstractmethod + def inactive_agents(self) -> DataFrame: ... \ No newline at end of file diff --git a/mesa_frames/agent.py b/mesa_frames/agent.py deleted file mode 100644 index 1084f79..0000000 --- a/mesa_frames/agent.py +++ /dev/null @@ -1,1543 +0,0 @@ -from __future__ import annotations # PEP 563: postponed evaluation of type annotations - -from abc import ABC, abstractmethod -from contextlib import suppress -from copy import copy, deepcopy -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Hashable, - Iterable, - Literal, - Self, - Sequence, - overload, -) - -import pandas as pd -import polars as pl -from numpy import int64, ndarray -from pandas.core.arrays.base import ExtensionArray -from polars.datatypes import N_INFER_DEFAULT - -from mesa_frames.model import ModelDF - -# For AgentSetPandas.select -ArrayLike = ExtensionArray | ndarray -AnyArrayLike = ArrayLike | pd.Index | pd.Series -ListLike = AnyArrayLike | list | range - -# For AgentSetPandas.drop -IndexLabel = Hashable | Sequence[Hashable] - -# For AgentContainer.__getitem__ and AgentContainer.__setitem__ -DataFrame = pd.DataFrame | pl.DataFrame - -Series = pd.Series | pl.Series - -BoolSeries = pd.Series | pl.Expr | pl.Series - -PandasMaskLike = ( - Literal["active"] | Literal["all"] | pd.Series | pd.DataFrame | ListLike | Hashable -) - -PolarsMaskLike = ( - Literal["active"] - | Literal["all"] - | pl.Expr - | pl.Series - | pl.DataFrame - | ListLike - | Hashable -) - -MaskLike = PandasMaskLike | PolarsMaskLike - -if TYPE_CHECKING: - - # For AgentSetDF - from numpy.random import Generator - - ValueKeyFunc = Callable[[pd.Series], pd.Series | AnyArrayLike] | None - - # For AgentSetPolars - from polars.type_aliases import ( - FrameInitTypes, - IntoExpr, - Orientation, - SchemaDefinition, - SchemaDict, - ) - - -### The AgentContainer class defines the interface for AgentSetDF and AgentsDF. It contains methods for selecting, shuffling, sorting, and manipulating agents. ### - - -class AgentContainer(ABC): - model: ModelDF - _mask: BoolSeries - _skip_copy: list[str] = ["model", "_mask"] - """An abstract class for containing agents. Defines the common interface for AgentSetDF and AgentsDF. - - Attributes - ---------- - model : ModelDF - The model to which the AgentContainer belongs. - _mask : Series - A boolean mask indicating which agents are active. - _skip_copy : list[str] - A list of attributes to skip during the copy process. - """ - - def __new__(cls, model: ModelDF) -> Self: - """Create a new AgentContainer object. - - Parameters - ---------- - model : ModelDF - The model to which the AgentContainer belongs. - - Returns - ------- - Self - A new AgentContainer object. - """ - obj = super().__new__(cls) - obj.model = model - return obj - - def __add__(self, other: Self | DataFrame | ListLike | dict[str, Any]) -> Self: - """Add agents to a new AgentContainer through the + operator. - - Other can be: - - A Self: adds the agents from the other AgentContainer. - - A DataFrame: adds the agents from the DataFrame. - - A ListLike: should be one single agent to add. - - A dictionary: keys should be attributes and values should be the values to add. - - Parameters - ---------- - other : Self | DataFrame | ListLike | dict[str, Any] - The agents to add. - - Returns - ------- - Self - A new AgentContainer with the added agents. - """ - new_obj = deepcopy(self) - return new_obj.add(other) - - def __iadd__(self, other: Self | DataFrame | ListLike | dict[str, Any]) -> Self: - """Add agents to the AgentContainer through the += operator. - - Other can be: - - A Self: adds the agents from the other AgentContainer. - - A DataFrame: adds the agents from the DataFrame. - - A ListLike: should be one single agent to add. - - A dictionary: keys should be attributes and values should be the values to add. - - Parameters - ---------- - other : Self | DataFrame | ListLike | dict[str, Any] - The agents to add. - - Returns - ------- - Self - The updated AgentContainer. - """ - return self.add(other) - - @abstractmethod - def __contains__(self, id: Hashable) -> bool: - """Check if an agent is in the AgentContainer. - - Parameters - ---------- - id : Hashable - The ID(s) to check for. - - Returns - ------- - bool - True if the agent is in the AgentContainer, False otherwise. - """ - - def __copy__(self) -> Self: - """Create a shallow copy of the AgentContainer. - - Returns - ------- - Self - A shallow copy of the AgentContainer. - """ - return self.copy(deep=False) - - def __deepcopy__(self, memo: dict) -> Self: - """Create a deep copy of the AgentContainer. - - Parameters - ---------- - memo : dict - A dictionary to store the copied objects. - - Returns - ------- - Self - A deep copy of the AgentContainer. - """ - return self.copy(deep=True, memo=memo) - - def __getattr__(self, name: str) -> Series: - """Fallback for retrieving attributes of the AgentContainer. Retrieves an attribute column of the agents in the AgentContainer. - - Parameters - ---------- - name : str - The name of the attribute to retrieve. - - Returns - ------- - Series - The attribute values. - """ - return self.get_attribute(name) - - @overload - def __getitem__(self, key: str | tuple[MaskLike, str]) -> Series: ... - - @overload - def __getitem__(self, key: list[str]) -> DataFrame: ... - - def __getitem__( - self, key: str | list[str] | MaskLike | tuple[MaskLike, str | list[str]] - ) -> Series | DataFrame: # tuple is not generic so it is not type hintable - """Implement the [] operator for the AgentContainer. - - The key can be: - - A string (eg. AgentContainer["str"]): returns the specified column of the agents in the AgentContainer. - - A list of strings(eg. AgentContainer[["str1", "str2"]]): returns the specified columns of the agents in the AgentContainer. - - A tuple (eg. AgentContainer[mask, "str"]): returns the specified column of the agents in the AgentContainer that satisfy the mask. - - A mask (eg. AgentContainer[mask]): returns the agents in the AgentContainer that satisfy the mask. - - Parameters - ---------- - key : str | list[str] | MaskLike | tuple[MaskLike, str | list[str]] - The key to retrieve. - - Returns - ------- - Series | DataFrame - The attribute values. - """ - if isinstance(key, (str, list)): - return self.get_attribute(attr_names=key) - - elif isinstance(key, tuple): - return self.get_attribute(mask=key[0], attr_names=key[1]) - - else: # MaskLike - return self.get_attribute(mask=key) - - @abstractmethod - def __iter__(self) -> Iterable: - """Iterate over the agents in the AgentContainer. - - Returns - ------- - Iterable - An iterator over the agents. - """ - - def __isub__(self, other: MaskLike) -> Self: - """Remove agents from the AgentContainer through the -= operator. - - Parameters - ---------- - other : Self | DataFrame | ListLike - The agents to remove. - - Returns - ------- - Self - The updated AgentContainer. - """ - return self.discard(other) - - @abstractmethod - def __len__(self) -> int | dict[str, int]: - """Get the number of agents in the AgentContainer. - - Returns - ------- - int | dict[str, int] - The number of agents in the AgentContainer. - """ - - @abstractmethod - def __repr__(self) -> str: - """Get a string representation of the DataFrame in the AgentContainer. - - Returns - ------- - str - A string representation of the DataFrame in the AgentContainer. - """ - return repr(self.agents) - - def __setitem__( - self, - key: str | list[str] | MaskLike | tuple[MaskLike, str | list[str]], - value: Any, - ) -> None: - """Implement the [] operator for setting values in the AgentContainer. - - The key can be: - - A string (eg. AgentContainer["str"]): sets the specified column of the agents in the AgentContainer. - - A list of strings(eg. AgentContainer[["str1", "str2"]]): sets the specified columns of the agents in the AgentContainer. - - A tuple (eg. AgentContainer[mask, "str"]): sets the specified column of the agents in the AgentContainer that satisfy the mask. - - A mask (eg. AgentContainer[mask]): sets the attributes of the agents in the AgentContainer that satisfy the mask. - - Parameters - ---------- - key : str | list[str] | MaskLike | tuple[MaskLike, str | list[str] - The key to set. - """ - if isinstance(key, (str, list)): - self.set_attribute(attr_names=key, value=value) - - elif isinstance(key, tuple): - self.set_attribute(mask=key[0], attr_names=key[1], value=value) - - else: # key=MaskLike - self.set_attribute(mask=key, value=value) - - @abstractmethod - def __str__(self) -> str: - """Get a string representation of the DataFrame in the AgentContainer. - - Returns - ------- - str - A string representation of the DataFrame in the AgentContainer. - """ - - def __sub__(self, other: MaskLike) -> Self: - """Remove agents from a new AgentContainer through the - operator. - - Parameters - ---------- - other : DataFrame | ListLike - The agents to remove. - - Returns - ------- - Self - A new AgentContainer with the removed agents. - """ - new_obj = deepcopy(self) - return new_obj.discard(other) - - @abstractmethod - def __reversed__(self) -> Iterable: - """Iterate over the agents in the AgentContainer in reverse order. - - Returns - ------- - Iterable - An iterator over the agents in reverse order. - """ - - @property - @abstractmethod - def agents(self) -> DataFrame: - """The agents in the AgentContainer. - - Returns - ------- - DataFrame - """ - - @property - @abstractmethod - def active_agents(self) -> DataFrame: - """The active agents in the AgentContainer. - - Returns - ------- - DataFrame - """ - - @active_agents.setter - def active_agents(self, mask: MaskLike) -> None: - """Set the active agents in the AgentContainer. - - Parameters - ---------- - mask : MaskLike - The mask to apply. - """ - self.select(mask=mask) - - @property - @abstractmethod - def inactive_agents(self) -> DataFrame: - """The inactive agents in the AgentContainer. - - Returns - ------- - DataFrame - """ - - @property - def random(self) -> Generator: - """ - Provide access to the model's random number generator. - - Returns - ------- - np.Generator - """ - return self.model.random - - def _get_obj(self, inplace: bool) -> Self: - """Get the object to perform operations on. - - Parameters - ---------- - inplace : bool - If inplace, return self. Otherwise, return a copy. - - Returns - ---------- - Self - The object to perform operations on. - """ - if inplace: - return self - else: - return deepcopy(self) - - @abstractmethod - def contains(self, ids: MaskLike) -> BoolSeries: - """Check if agents with the specified IDs are in the AgentContainer. - - Parameters - ---------- - id : MaskLike - The ID(s) to check for. - - Returns - ------- - BoolSeries - """ - - def copy( - self, - deep: bool = False, - skip: list[str] | str | None = None, - memo: dict | None = None, - ) -> Self: - """Create a copy of the AgentContainer. - - Parameters - ---------- - deep : bool, optional - Flag indicating whether to perform a deep copy of the AgentContainer. - If True, all attributes of the AgentContainer will be recursively copied (except self.agents, check Pandas/Polars documentation). - If False, only the top-level attributes will be copied. - Defaults to False. - - skip : list[str] | str | None, optional - A list of attribute names or a single attribute name to skip during the copy process. - If an attribute name is specified, it will be skipped for all levels of the copy. - If a list of attribute names is specified, they will be skipped for all levels of the copy. - If None, no attributes will be skipped. - Defaults to None. - - memo : dict | None, optional - A dictionary used to track already copied objects during deep copy. - Defaults to None. - - Returns - ------- - Self - A new instance of the AgentContainer class that is a copy of the original instance. - """ - skip_list = self._skip_copy.copy() - cls = self.__class__ - obj = cls.__new__(cls, self.model) - if isinstance(skip, str): - skip_list.append(skip) - elif isinstance(skip, list): - skip_list += skip - if deep: - if not memo: - memo = {} - memo[id(self)] = obj - attributes = self.__dict__.copy() - setattr(obj, "model", attributes.pop("model")) - [ - setattr(obj, k, deepcopy(v, memo)) - for k, v in attributes.items() - if k not in skip_list - ] - else: - [ - setattr(obj, k, copy(v)) - for k, v in self.__dict__.items() - if k not in skip_list - ] - return obj - - @abstractmethod - def select( - self, - mask: MaskLike | None = None, - filter_func: Callable[[Self], MaskLike] | None = None, - n: int | None = None, - inplace: bool = True, - ) -> Self: - """Select agents in the AgentContainer based on the given criteria. - - Parameters - ---------- - mask : MaskLike | None, optional - The mask of agents to be selected, by default None - filter_func : Callable[[Self], MaskLike] | None, optional - A function which takes as input the AgentContainer and returns a MaskLike, by default None - n : int, optional - The maximum number of agents to be selected, by default None - inplace : bool, optional - If the operation should be performed on the same object, by default True - - Returns - ------- - Self - A new or updated AgentContainer. - """ - - @abstractmethod - def shuffle(self, inplace: bool = True) -> Self: - """ - Shuffles the order of agents in the AgentContainer. - - Parameters - ---------- - inplace : bool - Whether to shuffle the agents in place. - - Returns - ---------- - Self - A new or updated AgentContainer. - """ - - @abstractmethod - def sort(self, *args, inplace: bool = True, **kwargs) -> Self: - """ - Sorts the agents in the agent set based on the given criteria. - - Parameters - ---------- - *args - Positional arguments to pass to the sort method. - inplace : bool - Whether to sort the agents in place. - **kwargs - Keyword arguments to pass to the sort - - Returns - ---------- - Self - A new or updated AgentContainer. - """ - - @overload - def do( - self, - method_name: str, - *args, - return_results: Literal[False] = False, - inplace: bool = True, - **kwargs, - ) -> Self: ... - - @overload - def do( - self, - method_name: str, - *args, - return_results: Literal[True], - inplace: bool = True, - **kwargs, - ) -> Any: ... - - def do( - self, - method_name: str, - *args, - return_results: bool = False, - inplace: bool = True, - **kwargs, - ) -> Self | Any: - """Invoke a method on the AgentContainer. - - Parameters - ---------- - method_name : str - The name of the method to invoke. - return_results : bool, optional - Whether to return the result of the method, by default False - inplace : bool, optional - Whether the operation should be done inplace, by default True - - Returns - ------- - Self | Any - The updated AgentContainer or the result of the method. - """ - obj = self._get_obj(inplace) - method = getattr(obj, method_name) - if return_results: - return method(*args, **kwargs) - else: - method(*args, **kwargs) - return obj - - @abstractmethod - @overload - def get_attribute( - self, - attr_names: list[str] | None = None, - mask: MaskLike | None = None, - ) -> DataFrame: ... - - @abstractmethod - @overload - def get_attribute( - self, - attr_names: str, - mask: MaskLike | None = None, - ) -> Series: ... - - @abstractmethod - def get_attribute( - self, - attr_names: str | list[str] | None = None, - mask: MaskLike | None = None, - ) -> Series | DataFrame: - """ - Retrieves the value of a specified attribute for each agent in the AgentContainer. - - Parameters - ---------- - attr_names : str | list[str] | None - The name of the attribute to retrieve. If None, all attributes are retrieved. Defaults to None. - mask : MaskLike | None - The mask of agents to retrieve the attribute for. If None, attributes of all agents are returned. Defaults to None. - - Returns - ---------- - Series | DataFrame - The attribute values. - """ - - @abstractmethod - @overload - def set_attribute( - self, - attr_names: None = None, - value: Any = Any, - mask: MaskLike = MaskLike, - inplace: bool = True, - ) -> Self: ... - - @abstractmethod - @overload - def set_attribute( - self, - attr_names: dict[str, Any], - value: None, - mask: MaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - @abstractmethod - @overload - def set_attribute( - self, - attr_names: str | list[str], - value: Any, - mask: MaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - @abstractmethod - def set_attribute( - self, - attr_names: str | dict[str, Any] | list[str] | None = None, - value: Any | None = None, - mask: MaskLike | None = None, - inplace: bool = True, - ) -> Self: - """ - Sets the value of a specified attribute or attributes for each agent in the AgentContainer. - - The key can be: - - A string: sets the specified column of the agents in the AgentContainer. - - A list of strings: sets the specified columns of the agents in the AgentContainer. - - A dictionary: keys should be attributes and values should be the values to set. Value should be None. - - Parameters - ---------- - attr_names : str | dict[str, Any] - The name of the attribute to set. - value : Any | None - The value to set the attribute to. If None, attr_names must be a dictionary. - mask : MaskLike | None - The mask of agents to set the attribute for. - inplace : bool - Whether to set the attribute in place. - - Returns - ---------- - AgentContainer - The updated agent set. - """ - - @abstractmethod - def add( - self, other: Self | DataFrame | ListLike | dict[str, Any], inplace: bool = True - ) -> Self: - """Adds agents to the AgentContainer. - - Other can be: - - A Self: adds the agents from the other AgentContainer. - - A DataFrame: adds the agents from the DataFrame. - - A ListLike: should be one single agent to add. - - A dictionary: keys should be attributes and values should be the values to add. - - Parameters - ---------- - other : Self | DataFrame | ListLike | dict[str, Any] - The agents to add. - inplace : bool, optional - Whether the operation is done into place, by default True - - Returns - ------- - Self - The updated AgentContainer. - """ - - def discard(self, id: MaskLike, inplace: bool = True) -> Self: - """ - Removes an agent from the AgentContainer. Does not raise an error if the agent is not found. - - Parameters - ---------- - id : ListLike | Any - The ID of the agent to remove. - inplace : bool - Whether to remove the agent in place. - - Returns - ---------- - AgentContainer - The updated AgentContainer. - """ - with suppress(KeyError): - return self.remove(id, inplace=inplace) - - @abstractmethod - def remove(self, id: MaskLike, inplace: bool = True) -> Self: - """ - Removes an agent from the AgentContainer. - - Parameters - ---------- - id : ListLike | Any - The ID of the agent to remove. - inplace : bool - Whether to remove the agent in place. - - Returns - ---------- - AgentContainer - The updated AgentContainer. - """ - - -### The AgentSetDF class is a container for agents of the same type. It has an implementation with Pandas and Polars ### - - -class AgentSetDF(AgentContainer): - _agents: DataFrame - _skip_copy = ["model", "_mask", "_agents"] - """A container for agents of the same type. - - Attributes - ---------- - model : ModelDF - The model to which the AgentSetDF belongs. - _mask : Series - A boolean mask indicating which agents are active. - _agents : DataFrame - The agents in the AgentSetDF. - _skip_copy : list[str] - A list of attributes to skip during the copy process. - """ - - @property - def agents(self) -> DataFrame: - """The agents in the AgentSetDF.""" - return self._agents - - @agents.setter - def agents_setter(self, agents: DataFrame) -> None: - """Set the agents in the AgentSetDF. - - Parameters - ---------- - agents : DataFrame - The agents to set. - """ - self._agents = agents - - def __len__(self) -> int: - return len(self._agents) - - def __repr__(self) -> str: - return repr(self._agents) - - def __str__(self) -> str: - return str(self._agents) - - def contains(self, ids: MaskLike) -> BoolSeries | bool: - - if isinstance( - ids, - (ListLike, Series, pl.Expr, DataFrame), - ) or ids == "all" or ids == "active": - return self._get_bool_mask(ids) - else: - return ids in self - - @abstractmethod - def _get_bool_mask(self, mask: MaskLike) -> BoolSeries: - """Get a boolean mask for the agents in the AgentSet. - - The mask can be: - - "all": all agents are selected. - - "active": only active agents are selected. - - A ListLike of IDs: only agents with the specified IDs are selected. - - A DataFrame: only agents with indices in the DataFrame are selected. - - A BoolSeries: only agents with True values are selected. - - Any other value: only the agent with the specified ID value is selected. - - Parameters - ---------- - mask : MaskLike - The mask to apply. - - Returns - ------- - BoolSeries - The boolean mask for the agents. - """ - - -class AgentSetPandas(AgentSetDF): - _agents: pd.DataFrame - _mask: pd.Series[bool] - """A pandas-based implementation of the AgentSet. - - Attributes - ---------- - model : ModelDF - The model to which the AgentSet belongs. - _mask : pd.Series[bool] - A boolean mask indicating which agents are active. - _agents : pd.DataFrame - The agents in the AgentSet. - _skip_copy : list[str] - A list of attributes to skip during the copy process. - """ - - def __new__(cls, model: ModelDF) -> Self: - obj = super().__new__(cls, model) - obj._agents = pd.DataFrame(columns=["unique_id"]).set_index("unique_id") - obj._mask = pd.Series(True, index=obj._agents.index) - return obj - - def __contains__(self, id: Hashable) -> bool: - return id in self._agents.index - - def __deepcopy__(self, memo: dict) -> Self: - obj = super().__deepcopy__(memo) - obj._agents = self._agents.copy(deep=True) - return obj - - def __iter__(self): - return self._agents.iterrows() - - def __reversed__(self) -> Iterable: - return self._agents[::-1].iterrows() - - @property - def agents(self) -> pd.DataFrame: - return self._agents - - @property - def active_agents(self) -> pd.DataFrame: - return self._agents.loc[self._mask] - - @active_agents.setter # When a property is overriden, so it is the getter - def active_agents(self, mask: PandasMaskLike) -> None: - return AgentContainer.active_agents.fset(self, mask) # type: ignore - - @property - def inactive_agents(self) -> pd.DataFrame: - return self._agents.loc[~self._mask] - - def _get_bool_mask( - self, - mask: PandasMaskLike | None = None, - ) -> pd.Series: - if isinstance(mask, pd.Series) and mask.dtype == bool: - return mask - elif isinstance(mask, self.__class__): - return pd.Series( - self._agents.index.isin(mask.agents.index), index=self._agents.index - ) - elif isinstance(mask, pd.DataFrame): - return pd.Series( - self._agents.index.isin(mask.index), index=self._agents.index - ) - elif isinstance(mask, list): - return pd.Series(self._agents.index.isin(mask), index=self._agents.index) - elif mask is None or mask == "all": - return pd.Series(True, index=self._agents.index) - elif mask == "active": - return self._mask - else: - return pd.Series(self._agents.index.isin([mask]), index=self._agents.index) - - def copy( - self, - deep: bool = False, - skip: list[str] | str | None = None, - memo: dict | None = None, - ) -> Self: - obj = super().copy(deep, skip, memo) - obj._agents = self._agents.copy(deep=deep) - obj._mask = self._mask.copy(deep=deep) - return obj - - def select( - self, - mask: PandasMaskLike | None = None, - filter_func: Callable[[Self], PandasMaskLike] | None = None, - n: int | None = None, - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) - bool_mask = obj._get_bool_mask(mask) - if n != None: - bool_mask = pd.Series( - obj._agents.index.isin(obj._agents[bool_mask].sample(n).index), - index=obj._agents.index, - ) - if filter_func: - bool_mask = bool_mask & obj._get_bool_mask(filter_func(obj)) - obj._mask = bool_mask - return obj - - def shuffle(self, inplace: bool = True) -> Self: - obj = self._get_obj(inplace) - obj._agents = obj._agents.sample(frac=1) - return obj - - def sort( - self, - by: str | Sequence[str], - key: ValueKeyFunc | None = None, - ascending: bool | Sequence[bool] = True, - inplace: bool = True, - ) -> Self: - """ - Sort the agents in the agent set based on the given criteria. - - Parameters - ---------- - by : str | Sequence[str] - The attribute(s) to sort by. - key : ValueKeyFunc | None - A function to use for sorting. - ascending : bool | Sequence[bool] - Whether to sort in ascending order. - - Returns - ---------- - AgentSetDF: The sorted agent set. - """ - obj = self._get_obj(inplace) - obj._agents.sort_values(by=by, key=key, ascending=ascending, inplace=True) - return obj - - @overload - def set_attribute( - self, - attr_names: None = None, - value: Any = Any, - mask: PandasMaskLike = PandasMaskLike, - inplace: bool = True, - ) -> Self: ... - - @overload - def set_attribute( - self, - attr_names: dict[str, Any], - value: None, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - @overload - def set_attribute( - self, - attr_names: str | list[str], - value: Any, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - def set_attribute( - self, - attr_names: str | list[str] | dict[str, Any] | None = None, - value: Any | None = None, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) - mask = obj._get_bool_mask(mask) - if attr_names is None: - attr_names = obj._agents.columns.values.tolist() - if isinstance(attr_names, (str, list)) and value is not None: - obj._agents.loc[mask, attr_names] = value - elif isinstance(attr_names, dict): - for key, value in attr_names.items(): - obj._agents.loc[mask, key] = value - else: - raise ValueError( - "attr_names must be a string or a dictionary with columns as keys and values." - ) - return obj - - @overload - def get_attribute( - self, - attr_names: list[str] | None = None, - mask: PandasMaskLike | None = None, - ) -> pd.DataFrame: ... - - @overload - def get_attribute( - self, - attr_names: str, - mask: PandasMaskLike | None = None, - ) -> pd.Series: ... - - def get_attribute( - self, - attr_names: str | list[str] | None = None, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> pd.Series | pd.DataFrame: - obj = self._get_obj(inplace) - mask = obj._get_bool_mask(mask) - if attr_names is None: - return obj._agents.loc[mask] - else: - return obj._agents.loc[mask, attr_names] - - def add( - self, - other: Self | pd.DataFrame | ListLike | dict[str, Any], - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) - if isinstance(other, obj.__class__): - new_agents = other.agents - elif isinstance(other, pd.DataFrame): - new_agents = other - if "unique_id" != other.index.name: - try: - new_agents.set_index("unique_id", inplace=True, drop=True) - except KeyError: - new_agents["unique_id"] = obj.random.random(len(other)) * 10**8 - elif isinstance(other, dict): - if "unique_id" not in other: - index = obj.random.random(len(other)) * 10**8 - if not isinstance(other["unique_id"], ListLike): - index = [other["unique_id"]] - else: - index = other["unique_id"] - new_agents = ( - pd.DataFrame(other, index=pd.Index(index)) - .reset_index(drop=True) - .set_index("unique_id") - ) - else: # ListLike - if len(other) == len(obj._agents.columns): - # data missing unique_id - new_agents = pd.DataFrame([other], columns=obj._agents.columns) - new_agents["unique_id"] = obj.random.random(1) * 10**8 - elif len(other) == len(obj._agents.columns) + 1: - new_agents = pd.DataFrame( - [other], columns=["unique_id"] + obj._agents.columns.values.tolist() - ) - else: - raise ValueError( - "Length of data must match the number of columns in the AgentSet if being added as a ListLike." - ) - new_agents.set_index("unique_id", inplace=True, drop=True) - obj._agents = pd.concat([obj._agents, new_agents]) - return obj - - def remove(self, id: PandasMaskLike, inplace: bool = True) -> Self: - initial_len = len(self._agents) - obj = self._get_obj(inplace) - mask = obj._get_bool_mask(id) - remove_ids = obj._agents[mask].index - obj._agents.drop(remove_ids, inplace=True) - if len(obj._agents) == initial_len: - raise KeyError(f"IDs {id} not found in agent set.") - return obj - - -class AgentSetPolars(AgentSetDF): - _agents: pl.DataFrame - _mask: pl.Expr | pl.Series - """A polars-based implementation of the AgentSet.""" - - def __init__(self, model: ModelDF): - """Create a new AgentSetDF. - - Parameters - ---------- - model : ModelDF - The model to which the AgentSetDF belongs. - - Attributes - ---------- - agents : pl.DataFrame - The agents in the AgentSetDF. - model : ModelDF - The model to which the AgentSetDF belongs. - """ - super().__init__(model) - self._agents = pl.DataFrame(schema={"unique_id": pl.String}) - self._mask = pl.repeat(True, len(self.agents)) - - @property - def agents(self) -> pl.DataFrame: - if self._agents is None: - self._agents = pl.DataFrame(schema={"unique_id": pl.String}) - return self._agents - - @property - def active_agents(self) -> pl.DataFrame: - return self.agents.filter(self._mask) - - @property - def inactive_agents(self) -> pl.DataFrame: - return self.agents.filter(~self._mask) - - def select( - self, - mask: pl.Expr | pl.Series | pl.DataFrame | None = None, - filter_func: Callable[[Self], pl.Series] | None = None, - n: int = 0, - ) -> Self: - if mask is None: # if not mask doesn't work - mask = pl.repeat(True, len(self.agents)) - elif isinstance(mask, pl.DataFrame): - mask = self.agents["unique_id"].is_in(mask["unique_id"]) - if filter_func: - mask = mask & filter_func(self) - if n != 0: - mask = ( - self.agents.filter(mask) - .sample(n)["unique_id"] - .is_in(self.agents["unique_id"]) - ) - self._mask = mask - return self - - def shuffle(self) -> Self: - self.agents = self.agents.sample(fraction=1) - return self - - def sort( - self, - by: IntoExpr | Iterable[IntoExpr], - *more_by: IntoExpr, - descending: bool | Sequence[bool] = False, - nulls_last: bool = False, - ) -> Self: - """Sort the agents in the agent set based on the given criteria. - - Parameters - ---------- - by (IntoExpr | Iterable[IntoExpr]): The attribute(s) to sort by. - more_by (IntoExpr): Additional attributes to sort by. - descending (bool | Sequence[bool]): Whether to sort in descending order. - nulls_last (bool): Whether to place null values last. - - Returns - ---------- - AgentSetDF: The sorted agent set. - """ - self.agents = self.agents.sort( - by=by, *more_by, descending=descending, nulls_last=nulls_last - ) - return self - - def get_attribute(self, attr_names: str) -> pl.Series: - return self.agents.filter(self._mask)[attr_names] - - def set_attribute(self, attr_names: str, value: Any) -> Self: - if type(value) == pl.Series: - self.agents.filter(self._mask).with_columns(**{attr_names: value}) - else: - self.agents.filter(self._mask).with_columns(**{attr_names: pl.lit(value)}) - return self - - def add( - self, - n: int, - data: FrameInitTypes | None = None, - schema: SchemaDefinition | None = None, - schema_overrides: SchemaDict | None = None, - orient: Orientation | None = None, - infer_schema_length: int | None = N_INFER_DEFAULT, - nan_to_null: bool = False, - ) -> Self: - """Adds new agents to the agent set. - - Parameters - ---------- - n : int - The number of agents to add. - data : dict, Sequence, ndarray, Series, or pandas.DataFrame - Two-dimensional data in various forms; dict input must contain Sequences, Generators, or a range. Sequence may contain Series or other Sequences. - schema : Sequence of str, (str,DataType) pairs, or a {str:DataType,} dict - The DataFrame schema may be declared in several ways: - - As a dict of {name:type} pairs; if type is None, it will be auto-inferred. - - As a list of column names; in this case types are automatically inferred. - - As a list of (name,type) pairs; this is equivalent to the dictionary form. - If you supply a list of column names that does not match the names in the underlying data, the names given here will overwrite them. The number of names given in the schema should match the underlying data dimensions. - schema_overrides : dict, default None - Support type specification or override of one or more columns; note that any dtypes inferred from the schema param will be overridden. - The number of entries in the schema should match the underlying data dimensions, unless a sequence of dictionaries is being passed, in which case a *partial* schema can be declared to prevent specific fields from being loaded. - orient : {'col', 'row'}, default None - Whether to interpret two-dimensional data as columns or as rows. - If None, the orientation is inferred by matching the columns and data dimensions. - If this does not yield conclusive results, column orientation is used. - infer_schema_length : int or None - The maximum number of rows to scan for schema inference. If set to None, the full data may be scanned *(this is slow)*. - This parameter only applies if the input data is a sequence or generator of rows; other input is read as-is. - nan_to_null : bool, default False - If the data comes from one or more numpy arrays, can optionally convert input data np.nan values to null instead. This is a no-op for all other input data. - - Returns - ---------- - AgentSetPolars: The updated agent set. - """ - new_df = pl.DataFrame( - data=data, - schema=schema, - schema_overrides=schema_overrides, - orient=orient, - infer_schema_length=infer_schema_length, - nan_to_null=nan_to_null, - ) - - if "unique_id" not in new_df.columns: - new_df = new_df.with_columns( - unique_id=pl.Series( - values=self.random.random(n) * 10**8, dtype=pl.Int64 - ) - ) - - old_active_agents = self.agents.filter(self._mask)["unique_id"] - self.agents = pl.concat([self.agents, new_df]) - self._mask = self.agents["unique_id"].is_in(old_active_agents) | self.agents[ - "unique_id" - ].is_in(new_df["unique_id"]) - return self - - def discard(self, id: int) -> Self: - with suppress(KeyError): - self.agents = self.agents.filter(self.agents["unique_id"] != id) - return self - - def remove(self, id: int) -> Self: - self.agents = self.agents.filter(self.agents["unique_id"] != id) - return self - - -### The AgentsDF class is a container for AgentSetDFs. It has an implementation with Pandas and Polars ### - - -class AgentsDF(AgentContainer): - agentsets: list[AgentSetDF] - """A collection of AgentSetDFs. All agents of the model are stored here.""" - - def __init__(self, model: ModelDF) -> None: - super().__init__(model) - self.agentsets = [] - - def __len__(self) -> int: - return sum(len(agentset.agents) for agentset in self.agentsets) - - def __repr__(self): - return self.agentsets.__repr__() - - def __str__(self) -> str: - return self.agentsets.__str__() - - def sort( - self, - by: str | Sequence[str], - key: ValueKeyFunc | None, - ascending: bool | Sequence[bool] = True, - ) -> Self: - self.agentsets = [ - agentset.sort(by, key, ascending) for agentset in self.agentsets - ] - return self - - @overload - def do( - self, - method_name: str, - return_results: Literal[False] = False, - *args, - **kwargs, - ) -> Self: ... - - @overload - def do( - self, - method_name: str, - return_results: Literal[True], - *args, - **kwargs, - ) -> list[Any]: ... - - def do( - self, - method_name: str, - return_results: bool = False, - *args, - **kwargs, - ) -> Self | list[Any]: - if return_results: - return [ - agentset.do(method_name, return_results, *args, **kwargs) - for agentset in self.agentsets - ] - else: - self.agentsets = [ - agentset.do(method_name, return_results, *args, **kwargs) - for agentset in self.agentsets - ] - return self - - def add(self, agentsets: AgentSetDF | list[AgentSetDF]) -> Self: - if isinstance(agentsets, list): - self.agentsets += agentsets - else: - self.agentsets.append(agentsets) - return self - - @abstractmethod - def to_frame(self) -> DataFrame: - """Convert the AgentsDF to a single DataFrame. - - Returns - ------- - DataFrame - A DataFrame containing all agents from all AgentSetDFs. - """ - pass - - def get_agents_of_type(self, agent_type: type[AgentSetDF]) -> AgentSetDF: - """Retrieve the AgentSetDF of a specified type. - - Parameters - ---------- - agent_type : type - The type of AgentSetDF to retrieve. - - Returns - ------- - AgentSetDF - The AgentSetDF of the specified type. - """ - for agentset in self.agentsets: - if isinstance(agentset, agent_type): - return agentset - raise ValueError(f"No AgentSetDF of type {agent_type} found.") - - def set_attribute(self, attr_names: str, value: Any) -> Self: - self.agentsets = [ - agentset.set_attribute(attr_names, value) for agentset in self.agentsets - ] - return self - - def shuffle(self) -> Self: - self.agentsets = [agentset.shuffle() for agentset in self.agentsets] - return self - - def discard(self, id: int) -> Self: - self.agentsets = [agentset.discard(id) for agentset in self.agentsets] - return self - - def remove(self, id: int) -> Self: - for i, agentset in enumerate(self.agentsets): - original_size = len(agentset.agents) - self.agentsets[i] = agentset.discard(id) - if original_size != len(self.agentsets[i].agents): - return self - raise KeyError(f"Agent with id {id} not found in any agentset.") - - -class AgentsPandas(AgentsDF): - agentsets: list[AgentSetPandas] - """A pandas implementation of a collection of AgentSetDF. All agents of the model are stored here.""" - - def __init__(self, model: ModelDF): - """Create a new AgentsDF object. - - Parameters - ---------- - model : ModelDF - The model to which the AgentsDF object belongs. - - Attributes - ---------- - agentsets : list[AgentSetDF] - The AgentSetDFs that make up the AgentsDF object. - model : ModelDF - The model to which the AgentSetDF belongs. - """ - super().__init__(model) - - @property - def active_agents(self) -> pd.DataFrame: - return pd.concat([agentset.active_agents for agentset in self.agentsets]) - - @property - def inactive_agents(self) -> pd.DataFrame: - return pd.concat([agentset.inactive_agents for agentset in self.agentsets]) - - def select( - self, - mask: pd.Series[bool] | pd.DataFrame | None = None, - filter_func: Callable[[AgentSetDF], pd.Series[bool]] | None = None, - n: int = 0, - ) -> Self: - n, r = int(n / len(self.agentsets)), n % len(self.agentsets) - new_agentsets: list[AgentSetPandas] = [] - for agentset in self.agentsets: - if mask is None: - agentset_mask = mask - elif isinstance(mask, pd.DataFrame): - agentset_mask = pd.Series( - agentset.agents.index.isin(mask), index=agentset.agents.index - ) - else: - agentset_mask = pd.Series( - agentset.agents.index.isin(mask[mask].index), - index=agentset.agents.index, - ) - agentset.select(mask=agentset_mask, filter_func=filter_func, n=n + r) - if len(agentset.active_agents) > n: - r = len(agentset.active_agents) - n - new_agentsets.append(agentset) - self.agentsets = new_agentsets - return self - - def get_attribute(self, attr_names: str) -> pd.Series[Any]: - return pd.concat( - [agentset.get_attribute(attr_names) for agentset in self.agentsets] - ) - - def add(self, agentsets: AgentSetPandas | list[AgentSetPandas]) -> Self: - return super().add(agentsets) # type: ignore - - -class AgentsPolars(AgentsDF): - agentsets: list[AgentSetPolars] - """A polars implementation of a collection of AgentSetDF. All agents of the model are stored here.""" - - def __init__(self, model: ModelDF): - """Create a new AgentsDF object. - - Parameters - ---------- - model : ModelDF - The model to which the AgentsDF object belongs. - - Attributes - ---------- - agentsets : list[AgentSetDF] - The AgentSetDFs that make up the AgentsDF object. - model : ModelDF - The model to which the AgentSetDF belongs. - """ - super().__init__(model) - - @property - def active_agents(self) -> pl.DataFrame: - return pl.concat([agentset.active_agents for agentset in self.agentsets]) - - @property - def inactive_agents(self) -> pl.DataFrame: - return pl.concat([agentset.inactive_agents for agentset in self.agentsets]) - - def select( - self, - mask: pl.Expr | pl.Series | pl.DataFrame | None = None, - filter_func: Callable[[AgentSetDF], pl.Series] | None = None, - n: int = 0, - ) -> Self: - n, r = int(n / len(self.agentsets)), n % len(self.agentsets) - new_agentsets: list[AgentSetPolars] = [] - for agentset in self.agentsets: - if mask is None: - agentset_mask = mask - elif isinstance(mask, pl.DataFrame): - agentset_mask = agentset.agents["unique_id"].is_in(mask["unique_id"]) - elif isinstance(mask, pl.Series): - agentset_mask = agentset.agents["unique_id"].is_in(mask) - agentset.select(mask=agentset_mask, filter_func=filter_func, n=n + r) - if len(agentset.active_agents) > n: - r = len(agentset.active_agents) - n - new_agentsets.append(agentset) - self.agentsets = new_agentsets - return self - - def get_attribute(self, attr_names: str) -> pl.Series: - return pl.concat( - [agentset.get_attribute(attr_names) for agentset in self.agentsets] - ) - - def add(self, agentsets: AgentSetPolars | list[AgentSetPolars]) -> Self: - return super().add(agentsets) # type: ignore #child classes are not checked? diff --git a/mesa_frames/base/agent.py b/mesa_frames/base/agent.py deleted file mode 100644 index 4c52528..0000000 --- a/mesa_frames/base/agent.py +++ /dev/null @@ -1,993 +0,0 @@ -from __future__ import annotations # PEP 563: postponed evaluation of type annotations - -from abc import ABC, abstractmethod -from contextlib import suppress -from copy import copy, deepcopy -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Hashable, - Iterable, - Literal, - Self, - Sequence, - overload, -) - -import pandas as pd -import polars as pl - -# from .model import ModelDF - -# For AgentContainer.__getitem__ and AgentContainer.__setitem__ -DataFrame = pd.DataFrame | pl.DataFrame - -Series = pd.Series | pl.Series - -BoolSeries = pd.Series | pl.Expr | pl.Series - -# For AgentSetPandas.select -ArrayLike = ExtensionArray | ndarray -AnyArrayLike = ArrayLike | pd.Index | pd.Series -ListLike = AnyArrayLike | list | range - -# For AgentSetPandas.drop -IndexLabel = Hashable | Sequence[Hashable] - -PolarsMaskLike = ( - Literal["active"] - | Literal["all"] - | pl.Expr - | pl.Series - | pl.DataFrame - | ListLike - | Hashable -) - -PandasMaskLike = ( - Literal["active"] | Literal["all"] | pd.Series | pd.DataFrame | ListLike | Hashable -) - -MaskLike = PandasMaskLike | PolarsMaskLike - -if TYPE_CHECKING: - - # For AgentSetDF - from numpy.random import Generator - - ValueKeyFunc = Callable[[pd.Series], pd.Series | AnyArrayLike] | None - - # For AgentSetPolars - - -### The AgentContainer class defines the interface for AgentSetDF and AgentsDF. It contains methods for selecting, shuffling, sorting, and manipulating agents. ### - - -class DataFrameAccessor: - def __init__(self, agent_set: AgentSetDF): - self._agent_set = agent_set - - def __getattr__(self, name): - obj = deepcopy(self._agent_set) - attr = getattr(obj._agents, name) - - if callable(attr): - - def method(*args, **kwargs): - result = attr(*args, **kwargs) - # If the result is a DataFrame, wrap it in the AgentSet class - if isinstance(result, (pd.DataFrame, pl.DataFrame)): - obj._agents = result - return obj - return result - - return method - return attr - - -class AgentContainer(ABC): - model: ModelDF - _mask: BoolSeries - _skip_copy: list[str] = ["model", "_mask"] - """An abstract class for containing agents. Defines the common interface for AgentSetDF and AgentsDF. - - Attributes - ---------- - model : ModelDF - The model to which the AgentContainer belongs. - _mask : Series - A boolean mask indicating which agents are active. - _skip_copy : list[str] - A list of attributes to skip during the copy process. - """ - - def __new__(cls, model: ModelDF) -> Self: - """Create a new AgentContainer object. - - Parameters - ---------- - model : ModelDF - The model to which the AgentContainer belongs. - - Returns - ------- - Self - A new AgentContainer object. - """ - obj = super().__new__(cls) - obj.model = model - return obj - - def __add__(self, other: Self | DataFrame | ListLike | dict[str, Any]) -> Self: - """Add agents to a new AgentContainer through the + operator. - - Other can be: - - A Self: adds the agents from the other AgentContainer. - - A DataFrame: adds the agents from the DataFrame. - - A ListLike: should be one single agent to add. - - A dictionary: keys should be attributes and values should be the values to add. - - Parameters - ---------- - other : Self | DataFrame | ListLike | dict[str, Any] - The agents to add. - - Returns - ------- - Self - A new AgentContainer with the added agents. - """ - new_obj = deepcopy(self) - return new_obj.add(other) - - def __iadd__(self, other: Self | DataFrame | ListLike | dict[str, Any]) -> Self: - """Add agents to the AgentContainer through the += operator. - - Other can be: - - A Self: adds the agents from the other AgentContainer. - - A DataFrame: adds the agents from the DataFrame. - - A ListLike: should be one single agent to add. - - A dictionary: keys should be attributes and values should be the values to add. - - Parameters - ---------- - other : Self | DataFrame | ListLike | dict[str, Any] - The agents to add. - - Returns - ------- - Self - The updated AgentContainer. - """ - return self.add(other) - - @abstractmethod - def __contains__(self, id: Hashable) -> bool: - """Check if an agent is in the AgentContainer. - - Parameters - ---------- - id : Hashable - The ID(s) to check for. - - Returns - ------- - bool - True if the agent is in the AgentContainer, False otherwise. - """ - - def __copy__(self) -> Self: - """Create a shallow copy of the AgentContainer. - - Returns - ------- - Self - A shallow copy of the AgentContainer. - """ - return self.copy(deep=False) - - def __deepcopy__(self, memo: dict) -> Self: - """Create a deep copy of the AgentContainer. - - Parameters - ---------- - memo : dict - A dictionary to store the copied objects. - - Returns - ------- - Self - A deep copy of the AgentContainer. - """ - return self.copy(deep=True, memo=memo) - - def __getattr__(self, name: str) -> Series: - """Fallback for retrieving attributes of the AgentContainer. Retrieves an attribute column of the agents in the AgentContainer. - - Parameters - ---------- - name : str - The name of the attribute to retrieve. - - Returns - ------- - Series - The attribute values. - """ - return self.get_attribute(name) - - @overload - def __getitem__(self, key: str | tuple[MaskLike, str]) -> Series: ... - - @overload - def __getitem__(self, key: list[str]) -> DataFrame: ... - - def __getitem__( - self, key: str | list[str] | MaskLike | tuple[MaskLike, str | list[str]] - ) -> Series | DataFrame: # tuple is not generic so it is not type hintable - """Implement the [] operator for the AgentContainer. - - The key can be: - - A string (eg. AgentContainer["str"]): returns the specified column of the agents in the AgentContainer. - - A list of strings(eg. AgentContainer[["str1", "str2"]]): returns the specified columns of the agents in the AgentContainer. - - A tuple (eg. AgentContainer[mask, "str"]): returns the specified column of the agents in the AgentContainer that satisfy the mask. - - A mask (eg. AgentContainer[mask]): returns the agents in the AgentContainer that satisfy the mask. - - Parameters - ---------- - key : str | list[str] | MaskLike | tuple[MaskLike, str | list[str]] - The key to retrieve. - - Returns - ------- - Series | DataFrame - The attribute values. - """ - if isinstance(key, (str, list)): - return self.get_attribute(attr_names=key) - - elif isinstance(key, tuple): - return self.get_attribute(mask=key[0], attr_names=key[1]) - - else: # MaskLike - return self.get_attribute(mask=key) - - @abstractmethod - def __iter__(self) -> Iterable: - """Iterate over the agents in the AgentContainer. - - Returns - ------- - Iterable - An iterator over the agents. - """ - - def __isub__(self, other: MaskLike) -> Self: - """Remove agents from the AgentContainer through the -= operator. - - Parameters - ---------- - other : Self | DataFrame | ListLike - The agents to remove. - - Returns - ------- - Self - The updated AgentContainer. - """ - return self.discard(other) - - @abstractmethod - def __len__(self) -> int | dict[str, int]: - """Get the number of agents in the AgentContainer. - - Returns - ------- - int | dict[str, int] - The number of agents in the AgentContainer. - """ - - @abstractmethod - def __repr__(self) -> str: - """Get a string representation of the DataFrame in the AgentContainer. - - Returns - ------- - str - A string representation of the DataFrame in the AgentContainer. - """ - return repr(self.agents) - - def __setitem__( - self, - key: str | list[str] | MaskLike | tuple[MaskLike, str | list[str]], - value: Any, - ) -> None: - """Implement the [] operator for setting values in the AgentContainer. - - The key can be: - - A string (eg. AgentContainer["str"]): sets the specified column of the agents in the AgentContainer. - - A list of strings(eg. AgentContainer[["str1", "str2"]]): sets the specified columns of the agents in the AgentContainer. - - A tuple (eg. AgentContainer[mask, "str"]): sets the specified column of the agents in the AgentContainer that satisfy the mask. - - A mask (eg. AgentContainer[mask]): sets the attributes of the agents in the AgentContainer that satisfy the mask. - - Parameters - ---------- - key : str | list[str] | MaskLike | tuple[MaskLike, str | list[str] - The key to set. - """ - if isinstance(key, (str, list)): - self.set_attribute(attr_names=key, value=value) - - elif isinstance(key, tuple): - self.set_attribute(mask=key[0], attr_names=key[1], value=value) - - else: # key=MaskLike - self.set_attribute(mask=key, value=value) - - @abstractmethod - def __str__(self) -> str: - """Get a string representation of the DataFrame in the AgentContainer. - - Returns - ------- - str - A string representation of the DataFrame in the AgentContainer. - """ - - def __sub__(self, other: MaskLike) -> Self: - """Remove agents from a new AgentContainer through the - operator. - - Parameters - ---------- - other : DataFrame | ListLike - The agents to remove. - - Returns - ------- - Self - A new AgentContainer with the removed agents. - """ - new_obj = deepcopy(self) - return new_obj.discard(other) - - @abstractmethod - def __reversed__(self) -> Iterable: - """Iterate over the agents in the AgentContainer in reverse order. - - Returns - ------- - Iterable - An iterator over the agents in reverse order. - """ - - @property - @abstractmethod - def agents(self) -> DataFrame: - """The agents in the AgentContainer. - - Returns - ------- - DataFrame - """ - - @property - @abstractmethod - def active_agents(self) -> DataFrame: - """The active agents in the AgentContainer. - - Returns - ------- - DataFrame - """ - - @active_agents.setter - def active_agents(self, mask: MaskLike) -> None: - """Set the active agents in the AgentContainer. - - Parameters - ---------- - mask : MaskLike - The mask to apply. - """ - self.select(mask=mask) - - @property - @abstractmethod - def inactive_agents(self) -> DataFrame: - """The inactive agents in the AgentContainer. - - Returns - ------- - DataFrame - """ - - @property - def random(self) -> Generator: - """ - Provide access to the model's random number generator. - - Returns - ------- - np.Generator - """ - return self.model.random - - def _get_obj(self, inplace: bool) -> Self: - """Get the object to perform operations on. - - Parameters - ---------- - inplace : bool - If inplace, return self. Otherwise, return a copy. - - Returns - ---------- - Self - The object to perform operations on. - """ - if inplace: - return self - else: - return deepcopy(self) - - @abstractmethod - def contains(self, ids: MaskLike) -> BoolSeries: - """Check if agents with the specified IDs are in the AgentContainer. - - Parameters - ---------- - id : MaskLike - The ID(s) to check for. - - Returns - ------- - BoolSeries - """ - - def copy( - self, - deep: bool = False, - skip: list[str] | str | None = None, - memo: dict | None = None, - ) -> Self: - """Create a copy of the AgentContainer. - - Parameters - ---------- - deep : bool, optional - Flag indicating whether to perform a deep copy of the AgentContainer. - If True, all attributes of the AgentContainer will be recursively copied (except self.agents, check Pandas/Polars documentation). - If False, only the top-level attributes will be copied. - Defaults to False. - - skip : list[str] | str | None, optional - A list of attribute names or a single attribute name to skip during the copy process. - If an attribute name is specified, it will be skipped for all levels of the copy. - If a list of attribute names is specified, they will be skipped for all levels of the copy. - If None, no attributes will be skipped. - Defaults to None. - - memo : dict | None, optional - A dictionary used to track already copied objects during deep copy. - Defaults to None. - - Returns - ------- - Self - A new instance of the AgentContainer class that is a copy of the original instance. - """ - skip_list = self._skip_copy.copy() - cls = self.__class__ - obj = cls.__new__(cls, self.model) - if isinstance(skip, str): - skip_list.append(skip) - elif isinstance(skip, list): - skip_list += skip - if deep: - if not memo: - memo = {} - memo[id(self)] = obj - attributes = self.__dict__.copy() - setattr(obj, "model", attributes.pop("model")) - [ - setattr(obj, k, deepcopy(v, memo)) - for k, v in attributes.items() - if k not in skip_list - ] - else: - [ - setattr(obj, k, copy(v)) - for k, v in self.__dict__.items() - if k not in skip_list - ] - return obj - - @abstractmethod - def select( - self, - mask: MaskLike | None = None, - filter_func: Callable[[Self], MaskLike] | None = None, - n: int | None = None, - inplace: bool = True, - ) -> Self: - """Select agents in the AgentContainer based on the given criteria. - - Parameters - ---------- - mask : MaskLike | None, optional - The mask of agents to be selected, by default None - filter_func : Callable[[Self], MaskLike] | None, optional - A function which takes as input the AgentContainer and returns a MaskLike, by default None - n : int, optional - The maximum number of agents to be selected, by default None - inplace : bool, optional - If the operation should be performed on the same object, by default True - - Returns - ------- - Self - A new or updated AgentContainer. - """ - - @abstractmethod - def shuffle(self, inplace: bool = True) -> Self: - """ - Shuffles the order of agents in the AgentContainer. - - Parameters - ---------- - inplace : bool - Whether to shuffle the agents in place. - - Returns - ---------- - Self - A new or updated AgentContainer. - """ - - @abstractmethod - def sort(self, *args, inplace: bool = True, **kwargs) -> Self: - """ - Sorts the agents in the agent set based on the given criteria. - - Parameters - ---------- - *args - Positional arguments to pass to the sort method. - inplace : bool - Whether to sort the agents in place. - **kwargs - Keyword arguments to pass to the sort - - Returns - ---------- - Self - A new or updated AgentContainer. - """ - - @overload - def do( - self, - method_name: str, - *args, - return_results: Literal[False] = False, - inplace: bool = True, - **kwargs, - ) -> Self: ... - - @overload - def do( - self, - method_name: str, - *args, - return_results: Literal[True], - inplace: bool = True, - **kwargs, - ) -> Any: ... - - def do( - self, - method_name: str, - *args, - return_results: bool = False, - inplace: bool = True, - **kwargs, - ) -> Self | Any: - """Invoke a method on the AgentContainer. - - Parameters - ---------- - method_name : str - The name of the method to invoke. - return_results : bool, optional - Whether to return the result of the method, by default False - inplace : bool, optional - Whether the operation should be done inplace, by default True - - Returns - ------- - Self | Any - The updated AgentContainer or the result of the method. - """ - obj = self._get_obj(inplace) - method = getattr(obj, method_name) - if return_results: - return method(*args, **kwargs) - else: - method(*args, **kwargs) - return obj - - @abstractmethod - @overload - def get_attribute( - self, - attr_names: list[str] | None = None, - mask: MaskLike | None = None, - ) -> DataFrame: ... - - @abstractmethod - @overload - def get_attribute( - self, - attr_names: str, - mask: MaskLike | None = None, - ) -> Series: ... - - @abstractmethod - def get_attribute( - self, - attr_names: str | list[str] | None = None, - mask: MaskLike | None = None, - ) -> Series | DataFrame: - """ - Retrieves the value of a specified attribute for each agent in the AgentContainer. - - Parameters - ---------- - attr_names : str | list[str] | None - The name of the attribute to retrieve. If None, all attributes are retrieved. Defaults to None. - mask : MaskLike | None - The mask of agents to retrieve the attribute for. If None, attributes of all agents are returned. Defaults to None. - - Returns - ---------- - Series | DataFrame - The attribute values. - """ - - @abstractmethod - @overload - def set_attribute( - self, - attr_names: None = None, - value: Any = Any, - mask: MaskLike = MaskLike, - inplace: bool = True, - ) -> Self: ... - - @abstractmethod - @overload - def set_attribute( - self, - attr_names: dict[str, Any], - value: None, - mask: MaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - @abstractmethod - @overload - def set_attribute( - self, - attr_names: str | list[str], - value: Any, - mask: MaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - @abstractmethod - def set_attribute( - self, - attr_names: str | dict[str, Any] | list[str] | None = None, - value: Any | None = None, - mask: MaskLike | None = None, - inplace: bool = True, - ) -> Self: - """ - Sets the value of a specified attribute or attributes for each agent in the AgentContainer. - - The key can be: - - A string: sets the specified column of the agents in the AgentContainer. - - A list of strings: sets the specified columns of the agents in the AgentContainer. - - A dictionary: keys should be attributes and values should be the values to set. Value should be None. - - Parameters - ---------- - attr_names : str | dict[str, Any] - The name of the attribute to set. - value : Any | None - The value to set the attribute to. If None, attr_names must be a dictionary. - mask : MaskLike | None - The mask of agents to set the attribute for. - inplace : bool - Whether to set the attribute in place. - - Returns - ---------- - AgentContainer - The updated agent set. - """ - - @abstractmethod - def add( - self, other: Self | DataFrame | ListLike | dict[str, Any], inplace: bool = True - ) -> Self: - """Adds agents to the AgentContainer. - - Other can be: - - A Self: adds the agents from the other AgentContainer. - - A DataFrame: adds the agents from the DataFrame. - - A ListLike: should be one single agent to add. - - A dictionary: keys should be attributes and values should be the values to add. - - Parameters - ---------- - other : Self | DataFrame | ListLike | dict[str, Any] - The agents to add. - inplace : bool, optional - Whether the operation is done into place, by default True - - Returns - ------- - Self - The updated AgentContainer. - """ - - def discard(self, id: MaskLike, inplace: bool = True) -> Self: - """ - Removes an agent from the AgentContainer. Does not raise an error if the agent is not found. - - Parameters - ---------- - id : ListLike | Any - The ID of the agent to remove. - inplace : bool - Whether to remove the agent in place. - - Returns - ---------- - AgentContainer - The updated AgentContainer. - """ - with suppress(KeyError): - return self.remove(id, inplace=inplace) - - @abstractmethod - def remove(self, id: MaskLike, inplace: bool = True) -> Self: - """ - Removes an agent from the AgentContainer. - - Parameters - ---------- - id : ListLike | Any - The ID of the agent to remove. - inplace : bool - Whether to remove the agent in place. - - Returns - ---------- - AgentContainer - The updated AgentContainer. - """ - - -### The AgentSetDF class is a container for agents of the same type. It has an implementation with Pandas and Polars ### - - -class AgentSetDF(AgentContainer): - _agents: DataFrame - _skip_copy = ["model", "_mask", "_agents"] - """A container for agents of the same type. - - Attributes - ---------- - model : ModelDF - The model to which the AgentSetDF belongs. - _mask : Series - A boolean mask indicating which agents are active. - _agents : DataFrame - The agents in the AgentSetDF. - _skip_copy : list[str] - A list of attributes to skip during the copy process. - """ - - @property - def agents(self) -> DataFrame: - """The agents in the AgentSetDF.""" - return self._agents - - @agents.setter - def agents_setter(self, agents: DataFrame) -> None: - """Set the agents in the AgentSetDF. - - Parameters - ---------- - agents : DataFrame - The agents to set. - """ - self._agents = agents - - def __len__(self) -> int: - return len(self._agents) - - def __repr__(self) -> str: - return repr(self._agents) - - def __str__(self) -> str: - return str(self._agents) - - def contains(self, ids: MaskLike) -> BoolSeries | bool: - - if ( - isinstance( - ids, - (ListLike, Series, pl.Expr, DataFrame), - ) - or ids == "all" - or ids == "active" - ): - return self._get_bool_mask(ids) - else: - return ids in self - - @abstractmethod - def _get_bool_mask(self, mask: MaskLike) -> BoolSeries: - """Get a boolean mask for the agents in the AgentSet. - - The mask can be: - - "all": all agents are selected. - - "active": only active agents are selected. - - A ListLike of IDs: only agents with the specified IDs are selected. - - A DataFrame: only agents with indices in the DataFrame are selected. - - A BoolSeries: only agents with True values are selected. - - Any other value: only the agent with the specified ID value is selected. - - Parameters - ---------- - mask : MaskLike - The mask to apply. - - Returns - ------- - BoolSeries - The boolean mask for the agents. - """ - - -### The AgentsDF class is a container for AgentSetDFs. It has an implementation with Pandas and Polars ### - - -class AgentsDF(AgentContainer): - agentsets: list[AgentSetDF] - """A collection of AgentSetDFs. All agents of the model are stored here.""" - - def __init__(self, model: ModelDF) -> None: - super().__init__(model) - self.agentsets = [] - - def __len__(self) -> int: - return sum(len(agentset.agents) for agentset in self.agentsets) - - def __repr__(self): - return self.agentsets.__repr__() - - def __str__(self) -> str: - return self.agentsets.__str__() - - def sort( - self, - by: str | Sequence[str], - key: ValueKeyFunc | None, - ascending: bool | Sequence[bool] = True, - ) -> Self: - self.agentsets = [ - agentset.sort(by, key, ascending) for agentset in self.agentsets - ] - return self - - @overload - def do( - self, - method_name: str, - return_results: Literal[False] = False, - *args, - **kwargs, - ) -> Self: ... - - @overload - def do( - self, - method_name: str, - return_results: Literal[True], - *args, - **kwargs, - ) -> list[Any]: ... - - def do( - self, - method_name: str, - return_results: bool = False, - *args, - **kwargs, - ) -> Self | list[Any]: - if return_results: - return [ - agentset.do(method_name, return_results, *args, **kwargs) - for agentset in self.agentsets - ] - else: - self.agentsets = [ - agentset.do(method_name, return_results, *args, **kwargs) - for agentset in self.agentsets - ] - return self - - def add(self, agentsets: AgentSetDF | list[AgentSetDF]) -> Self: - if isinstance(agentsets, list): - self.agentsets += agentsets - else: - self.agentsets.append(agentsets) - return self - - @abstractmethod - def to_frame(self) -> DataFrame: - """Convert the AgentsDF to a single DataFrame. - - Returns - ------- - DataFrame - A DataFrame containing all agents from all AgentSetDFs. - """ - pass - - def get_agents_of_type(self, agent_type: type[AgentSetDF]) -> AgentSetDF: - """Retrieve the AgentSetDF of a specified type. - - Parameters - ---------- - agent_type : type - The type of AgentSetDF to retrieve. - - Returns - ------- - AgentSetDF - The AgentSetDF of the specified type. - """ - for agentset in self.agentsets: - if isinstance(agentset, agent_type): - return agentset - raise ValueError(f"No AgentSetDF of type {agent_type} found.") - - def set_attribute(self, attr_names: str, value: Any) -> Self: - self.agentsets = [ - agentset.set_attribute(attr_names, value) for agentset in self.agentsets - ] - return self - - def shuffle(self) -> Self: - self.agentsets = [agentset.shuffle() for agentset in self.agentsets] - return self - - def discard(self, id: int) -> Self: - self.agentsets = [agentset.discard(id) for agentset in self.agentsets] - return self - - def remove(self, id: int) -> Self: - for i, agentset in enumerate(self.agentsets): - original_size = len(agentset.agents) - self.agentsets[i] = agentset.discard(id) - if original_size != len(self.agentsets[i].agents): - return self - raise KeyError(f"Agent with id {id} not found in any agentset.") diff --git a/mesa_frames/base/model.py b/mesa_frames/base/model.py deleted file mode 100644 index a18d10d..0000000 --- a/mesa_frames/base/model.py +++ /dev/null @@ -1,397 +0,0 @@ -from logging import warn -from typing import Any - -import numpy as np - -from .base.agent import AgentsDF, AgentSetDF - - -class ModelDF: - random: np.random.Generator - _seed: int - running: bool - _agents: AgentsDF | None - - def __new__(cls, seed: int = 0, *args: Any, **kwargs: Any) -> Any: - """Create a new model object and instantiate its RNG automatically.""" - obj = object.__new__(cls) - if seed == 0: - seed = np.random.SeedSequence().entropy # type: ignore - obj._seed = seed - obj.random = np.random.default_rng(seed=obj._seed) - return obj - - def __init__(self, *args: Any, **kwargs: Any) -> None: - """Create a new model. Overload this method with the actual code to - start the model. Always start with super().__init__() to initialize the - model object properly. - """ - self.running: bool = True - self._agents: AgentsDF | None = None - - @property - def agents(self) -> AgentsDF | None: - return self._agents - - @agents.setter - def agents(self, agents: AgentsDF) -> None: - if not isinstance(agents, AgentsDF): - raise TypeError("agents must be an instance of AgentsDF") - if type(agents) != type(self._agents): - warn( - f"Changing the backend from {type(self._agents)} to {type(agents)}", - RuntimeWarning, - ) - self._agents = agents - - def get_agents_of_type(self, agent_type: type) -> AgentSetDF: - """Retrieve the AgentSetDF of a specified type. - - Parameters - ---------- - agent_type : type - The type of AgentSetDF to retrieve. - - Returns - ------- - AgentSetDF - The AgentSetDF of the specified type. - """ - return self.agents.get_agents_of_type(agent_type) - - def next_id(self) -> int: - raise NotImplementedError("next_id() method not implemented for ModelDF") - - def reset_randomizer(self, seed: int | None = None) -> None: - """Reset the model random number generator. - - Parameters: - ---------- - seed : int | None - A new seed for the RNG; if None, reset using the current seed - """ - - self._seed = np.random.SeedSequence(seed=seed).entropy # type: ignore - - """'def initialize_data_collector( - self, model_reporters=None, agent_reporters=None, tables=None - ) -> None: - if not self._agents: - raise RuntimeError( - "You must create agents before initializing the data collector." - ) - self.datacollector = DataCollectorDF( - model_reporters=model_reporters, - agent_reporters=agent_reporters, - tables=tables, - ) - self.datacollector.collect(self)""" - - -''' -OLD IMPLEMENTATION (HAS TO BE DELETED) - -class ModelDF: - """The base class for all models - - Attributes - ---------- - unique_id : int - The unique_id of the model. - running : bool - Indicates if the model is running or not. - agents : pd.DataFrame | gpd.GeoDataFrame | None - The dataframe containing the agents of the model. - agent_types : list[tuple[type[AgentDF], float]] | None - The list of agent types and their proportions. - p_agents : dict[type[AgentDF], float] | None - The dictionary of agents to create. The keys are the types of agents, - the values are the percentages of each agent type. The sum of the values should be 1. - space - The space where the agents will be placed. Can be None if model does not have a space. - """ - - def __new__(cls, *args, **kwargs): - """Create a new model object and instantiate its RNG automatically - (adds supports to numpy with respect to base model).""" - obj = object.__new__(cls) - obj._seed = kwargs.get("seed") - if obj._seed is None: - # We explicitly specify the seed here so that we know its value in - # advance. - obj._seed = np.random.SeedSequence().entropy - # Use default_rng to get a new Generator instance - obj.random = np.random.default_rng(seed = obj._seed) - return obj - - def __init__(self, unique_id: int | None = None, space=None): - """Create a new model. Overload this method with the actual code to - start the model. Always start with super().__init__() to initialize the - model object properly. - - Parameters - ---------- - unique_id : int | None - The unique_id of the model. - If None, a random unique_id is assigned using a 64-bit random integer. - space - The space where the agents will be placed. Can be None if model does not have a space. - """ - # Initialize default attributes - self.running: bool = True - self.agents: pd.DataFrame | gpd.GeoDataFrame | None = None - self.agent_types: list[tuple[type[AgentDF], float]] | None = None - self.p_agents: dict[type[AgentDF], float] | None = None - # self.schedule : BaseScheduler = None - - # Initialize optional parameters - if not unique_id: - self.unique_id = np.random.randint( - low=-9223372036854775808, high=9223372036854775807, dtype="int64" - ) - else: - self.unique_id = unique_id - self.space = space - - # Initialize data collection - # self.initialize_data_collector(data_collection) - - def get_agents_of_type(self, agent_type: type[AgentDF]) -> pd.Series: - """Returns a boolean mask of the agents dataframe of the model which corresponds to the agent_type. - - Parameters - ---------- - agent_type : type[AgentDF] - The agent_type to get the mask for. - """ - if self.agents is None: - raise RuntimeError( - "You must create agents before getting their masks. Use create_agents() method." - ) - return self.agents["type"].str.contains(agent_type.__name__) # type: ignore - - def run_model(self, n_steps: int | None = None, merged_mro: bool = False) -> None: - """If n_steps are specified, executes model.step() until n_steps are reached. - Otherwise, until self.running is false (as the default mesa.Model.run_model). - - Parameters - ---------- - n_steps : int | None - The number of steps which the model will execute. - Can be None if a running condition turning false is used. - merged_mro: bool - If False, the model will execute one step for each class in p_agent. This is the default behaviour. - If True, the model will execute one step for each inherited agent type in the order of a "merged" MRO. - This may increase performance if there are multiple and complex inheritance as each agent_type (even if parents of different classes), - will be executed only once. Not a viable option if the behavior of a class depends on another. - """ - if n_steps: - if not (isinstance(n_steps, int) and n_steps > 0): - raise TypeError( - "n_steps should be an integer greater than 0 or None if a running condition is used" - ) - for _ in range(n_steps): - self.step(merged_mro) - else: - while self.running: - self.step(merged_mro) - - def step(self, merged_mro: bool = False) -> None: - """Executes one step of the model. - - Parameters - ---------- - merged_mro : bool - If False, the model will execute one step for each class in p_agent. This is the default behaviour. - If True, the model will execute one step for each inherited agent type in the order of a "merged" MRO. - This may increase performance if there are multiple and complex inheritance as each agent_type (even if parents of different classes), - will be executed only once. Not a viable option if the behavior of a class depends on another. - """ - if self.agent_types is None or self.p_agents is None: - raise RuntimeError( - "You must create agents before running the model. Use create_agents() method." - ) - if merged_mro: - for agent_type in self.agent_types: - agent_type[0].step() - else: - for agent in self.p_agents: - agent.step() - - def reset_randomizer(self, seed: int | None = None) -> None: - """Reset the model random number generator. - - Parameters - ---------- - seed : int | None - A new seed for the RNG; if None, reset using the current seed - """ - if seed is None: - seed = self._seed - self.random = np.random.default_rng(seed) - self._seed = seed - - def create_agents( - self, n_agents: int, p_agents: dict[type[AgentDF], float] - ) -> None: - """Populate the self.agents dataframe. - - Parameters - ---------- - n_agents : int | None - The number of agents which the model will create. - p_agents : dict[type[AgentDF], float] - The dictionary of agents to create. The keys are the types of agents, - the values are the percentages of each agent type. The sum of the values should be 1. - """ - - # Verify parameters - if not (isinstance(n_agents, int) and n_agents > 0): - raise TypeError("n_agents should be an integer greater than 0") - if sum(p_agents.values()) != 1: - raise ValueError("Sum of proportions of agents should be 1") - if any(p < 0 or p > 1 for p in p_agents.values()): - raise ValueError("Proportions of agents should be between 0 and 1") - - self.p_agents = p_agents - - start_time = time() - print("Creating agents: ...") - - mros = [[agent.__mro__[:-1], p] for agent, p in p_agents.items()] - mros_copy = deepcopy(mros) - agent_types = [] - - # Create a "merged MRO" (inspired by C3 linearization algorithm) - while True: - candunique_idate_added = False - # if all mros are empty, the merged mro is done - if not any(mro[0] for mro in mros): - break - for mro in mros: - # If mro is empty, continue - if not mro[0]: - continue - # candunique_idate = head - candunique_idate = mro[0][0] - # If candunique_idate appears in the tail of another MRO, skip it for now (because other agent_types depend on it, will be added later) - if any( - candunique_idate in other_mro[0][1:] - for other_mro in mros - if other_mro is not mro - ): - continue - else: - p = 0 - for i, other_mro in enumerate(mros): - if other_mro[0][0] == candunique_idate: - p += other_mro[1] - mros[i][0] = other_mro[0][1:] - else: - continue - agent_types.append((candunique_idate, p)) # Safe to add it - candunique_idate_added = True - # If there wasn't any good head, there is an inconsistent hierarchy - if not candunique_idate_added: - raise ValueError("Inconsistent hierarchy") - self.agent_types = list(agent_types) - - # Create a single DF using vars and values for every class - columns: set[str] = set() - dtypes: dict[str, str] = {} - for agent_type in self.agent_types: - for key, val in agent_type[0].dtypes.items(): - if key not in columns: - columns.add(key) - dtypes[key] = val - - if "geometry" in columns: - if not (self.space and hasattr(self.space, "crs")): - raise ValueError( - "You must specify a space with a crs attribute if you want to create GeoAgents" - ) - self.agents = gpd.GeoDataFrame( - index=pd.RangeIndex(0, n_agents), - columns=list(columns), - crs=self.space.crs, - ) - else: - self.agents = pd.DataFrame( - index=pd.RangeIndex(0, n_agents), columns=list(columns) - ) - - # Populate agents type - start_index = 0 - for i, (_, p) in enumerate(p_agents.items()): - self.agents.loc[ - start_index : start_index + int(n_agents * p) - 1, "type" - ] = str(mros_copy[i][0]) - start_index += int(n_agents * p) - - # Initialize agents - AgentDF.model = self - self.update_agents_masks() - for agent in p_agents: - agent.__init__() - - # Set dtypes - for col, dtype in dtypes.items(): - if "int" in dtype and self.agents[col].isna().sum() > 0: # type: ignore - warn( - f"Pandas does not support NaN values for int{dtype[-2:]} dtypes. Changing dtype to float{dtype[-2:]} for {col}", - RuntimeWarning, - ) - dtypes[col] = "float" + dtype[-2:] - self.agents = self.agents.astype(dtypes) - - # Set agents' unique_id as index (Have to reassign masks because index changed) - self.agents.set_index("id", inplace=True) - self.update_agents_masks() - - print("Created agents: " + "--- %s seconds ---" % (time() - start_time)) - - def _initialize_data_collection(self, how="2d") -> None: - """Initializes the data collection of the model. - - Parameters - ---------- - how : str - The frequency of the data collection. It can be 'xd', 'xd', 'xh', 'weekly', 'daily', 'hourly'. - """ - # TODO: finish implementation of different data collections - if how == "2d": - return - - def update_agents_masks(self) -> None: - """Updates the masks attributes of each agent in self.agent_types. - Useful after agents are created/deleted or index changes. - """ - if self.agent_types is None: - raise RuntimeError( - "You must create agents before updating their masks. Use create_agents() method." - ) - for agent_type in self.agent_types: - agent_type[0].mask = self.get_agents_of_type(agent_type[0]) - - # TODO: implement different data collection frequencies (xw, xd, xh, weekly, daily, hourly, per every step): - """def initialize_data_collector( - self, - model_reporters=None, - agent_reporters=None, - tables=None, - ) -> None: - if not hasattr(self, "schedule") or self.schedule is None: - raise RuntimeError( - "You must initialize the scheduler (self.schedule) before initializing the data collector." - ) - if self.schedule.get_agent_count() == 0: - raise RuntimeError( - "You must add agents to the scheduler before initializing the data collector." - ) - self.datacollector = DataCollector( - model_reporters=model_reporters, - agent_reporters=agent_reporters, - tables=tables, - ) - # Collect data for the first time during initialization. - self.datacollector.collect(self)''' diff --git a/mesa_frames/ciccio.py b/mesa_frames/ciccio.py deleted file mode 100644 index 45316e4..0000000 --- a/mesa_frames/ciccio.py +++ /dev/null @@ -1,26 +0,0 @@ -from copy import deepcopy - -import pandas as pd - -from mesa_frames import AgentSetPandas, ModelDF - - -class CiccioAgentSet(AgentSetPandas): - def __init__(self, model: ModelDF): - self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth") - - def add_wealth(self, amount: int) -> None: - self.agents["wealth"] += amount - - -example_model = ModelDF() -example = CiccioAgentSet(example_model) - -example._agents = pd.DataFrame({"unique_id": [0, 1, 2, 3, 4]}).set_index("unique_id") -example.add({"unique_id": [0, 1, 2, 3]}) - - - -deepcopy(example) -print(example.model) -print(dir(example)) diff --git a/mesa_frames/concrete/__init__.py b/mesa_frames/concrete/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mesa_frames/concrete/agents.py b/mesa_frames/concrete/agents.py new file mode 100644 index 0000000..4aceaba --- /dev/null +++ b/mesa_frames/concrete/agents.py @@ -0,0 +1,333 @@ +from typing import Any, Callable, Iterable, Iterator, Literal, Self, Sequence, overload + +from mesa_frames.abstract.agents import AgentContainer, AgentSetDF, Collection, Hashable +from mesa_frames.types import BoolSeries, DataFrame, MaskLike, Series + + +class AgentsDF(AgentContainer): + _agentsets: list[AgentSetDF] + _copy_with_method: dict[str, tuple[str, list[str]]] = { + "_agentsets": ("copy", []), + } + _backend: str + """A collection of AgentSetDFs. All agents of the model are stored here. + + Attributes + ---------- + _agentsets : list[AgentSetDF] + The agent sets contained in this collection. + _copy_with_method : dict[str, tuple[str, list[str]]] + A dictionary of attributes to copy with a specified method and arguments. + _backend : str + The backend used for data operations. + + Properties + ---------- + active_agents(self) -> dict[str, pd.DataFrame] + Get the active agents in the AgentsDF. + agents(self) -> dict[str, pd.DataFrame] + Get or set the agents in the AgentsDF. + inactive_agents(self) -> dict[str, pd.DataFrame] + Get the inactive agents in the AgentsDF. + model(self) -> ModelDF + Get the model associated with the AgentsDF. + random(self) -> np.random.Generator + Get the random number generator associated with the model. + + Methods + ------- + __init__(self) -> None + Initialize a new AgentsDF. + add(self, other: AgentSetDF | list[AgentSetDF], inplace: bool = True) -> Self + Add agents to the AgentsDF. + contains(self, ids: Hashable | Collection[Hashable]) -> bool | dict[str, pd.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 + 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] + Retrieve the value of a specified attribute for each agent in the AgentsDF. + remove(self, ids: MaskLike, 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. + set(self, attr_names: str | Collection[str] | dict[str, Any] | None = None, values: Any | None = None, mask: MaskLike | None = None, inplace: bool = True) -> Self + Set the value of a specified attribute or attributes for each agent in the mask in the AgentsDF. + shuffle(self, inplace: bool = True) -> Self + 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. + __getattr__(self, key: str) -> Any + Retrieve an attribute of the underlying agent sets. + __iter__(self) -> Iterator + Get an iterator for the agents in the AgentsDF. + __len__(self) -> int + Get the number of agents in the AgentsDF. + __repr__(self) -> str + Get the string representation of the AgentsDF. + __reversed__(self) -> Iterator + Get a reversed iterator for the agents in the AgentsDF. + __str__(self) -> str + Get the string representation of the AgentsDF. + """ + + def __init__(self) -> None: + self._agentsets = [] + + def add(self, other: AgentSetDF | list[AgentSetDF], inplace: bool = True) -> Self: + """Add an AgentSetDF to the AgentsDF. + + Parameters + ---------- + other : AgentSetDF + The AgentSetDF to add. + inplace : bool + Whether to add the AgentSetDF in place. + + Returns + ---------- + Self + The updated AgentsDF. + """ + obj = self._get_obj(inplace) + if isinstance(other, list): + obj._agentsets += other + else: + obj._agentsets.append(other) + return self + + @overload + def contains(self, ids: Collection[Hashable]) -> BoolSeries: ... + + @overload + def contains(self, ids: Hashable) -> bool: ... + + 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 + + @overload + def do( + self, + method_name: str, + *args, + return_results: Literal[False] = False, + inplace: bool = True, + **kwargs, + ) -> Self: ... + + @overload + def do( + self, + method_name: str, + *args, + return_results: Literal[True], + inplace: bool = True, + **kwargs, + ) -> dict[str, Any]: ... + + def do( + self, + method_name: str, + *args, + return_results: bool = False, + inplace: bool = True, + **kwargs, + ) -> Self | Any: + obj = self._get_obj(inplace) + if return_results: + return { + agentset.__class__.__name__: agentset.do( + method_name, + *args, + return_results=return_results, + **kwargs, + inplace=inplace, + ) + for agentset in obj._agentsets + } + else: + obj._agentsets = [ + agentset.do( + method_name, + *args, + return_results=return_results, + **kwargs, + inplace=inplace, + ) + for agentset in obj._agentsets + ] + return obj + + def get( + self, + attr_names: str | list[str] | None = None, + mask: MaskLike | None = None, + ) -> dict[str, Series] | dict[str, DataFrame]: + return { + agentset.__class__.__name__: agentset.get(attr_names, mask) + for agentset in self._agentsets + } + + def remove(self, ids: MaskLike, inplace: bool = True) -> Self: + obj = self._get_obj(inplace) + deleted = 0 + for agentset in obj._agentsets: + initial_len = len(agentset) + agentset.discard(ids, inplace=True) + deleted += initial_len - len(agentset) + if deleted < len(list(ids)): # TODO: fix type hint + raise KeyError(f"None of the agentsets contain the ID {MaskLike}.") + return obj + + def set( + self, + attr_names: str | dict[str, Any] | Collection[str], + values: Any | None = None, + mask: MaskLike | None = None, + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + obj._agentsets = [ + agentset.set( + attr_names=attr_names, values=values, mask=mask, inplace=inplace + ) + for agentset in obj._agentsets + ] + return obj + + def select( + self, + mask: MaskLike | None = None, + filter_func: Callable[[DataFrame], MaskLike] | None = None, + n: int | None = None, + inplace: bool = True, + negate: bool = False, + ) -> Self: + obj = self._get_obj(inplace) + obj._agentsets = [ + agentset.select( + mask=mask, filter_func=filter_func, n=n, negate=negate, inplace=inplace + ) + for agentset in obj._agentsets + ] + return obj + + def shuffle(self, inplace: bool = True) -> Self: + obj = self._get_obj(inplace) + obj._agentsets = [agentset.shuffle(inplace) for agentset in obj._agentsets] + return obj + + def sort( + self, + by: str | Sequence[str], + ascending: bool | Sequence[bool] = True, + inplace: bool = True, + **kwargs, + ) -> Self: + obj = self._get_obj(inplace) + obj._agentsets = [ + agentset.sort(by=by, ascending=ascending, inplace=inplace, **kwargs) + for agentset in obj._agentsets + ] + return obj + + def __add__(self, other: AgentSetDF | list[AgentSetDF]) -> Self: + """Add AgentSetDFs to a new AgentsDF through the + operator. + + Parameters + ---------- + other : AgentSetDF | list[AgentSetDF] + The AgentSetDFs to add. + + Returns + ------- + AgentsDF + A new AgentsDF with the added AgentSetDFs. + """ + return super().__add__(other) + + def __getattr__(self, name: str) -> dict[str, Any]: + return { + agentset.__class__.__name__: getattr(agentset, name) + for agentset in self._agentsets + } + + def __iadd__(self, other: AgentSetDF | list[AgentSetDF]) -> Self: + """Add AgentSetDFs to the AgentsDF through the += operator. + + Parameters + ---------- + other : Self | AgentSetDF | list[AgentSetDF] + The AgentSetDFs to add. + + Returns + ------- + AgentsDF + The updated AgentsDF. + """ + return super().__iadd__(other) + + def __iter__(self) -> Iterator: + return ( + agent for agentset in self._agentsets for agent in iter(agentset._backend) + ) + + def __repr__(self) -> str: + return str( + { + agentset.__class__.__name__: repr(agentset) + for agentset in self._agentsets + } + ) + + def __str__(self) -> str: + return str( + {agentset.__class__.__name__: str(agentset) for agentset in self._agentsets} + ) + + def __reversed__(self) -> Iterator: + return ( + agent + for agentset in self._agentsets + for agent in reversed(agentset._backend) + ) + + def __len__(self) -> int: + return sum(len(agentset._agents) for agentset in self._agentsets) + + @property + def agents(self) -> dict[str, DataFrame]: + return { + agentset.__class__.__name__: agentset.agents for agentset in self._agentsets + } + + @agents.setter + def agents(self, other: Iterable[AgentSetDF]) -> None: + """Set the agents in the AgentsDF. + + Parameters + ---------- + other : Iterable[AgentSetDF] + The AgentSetDFs to set. + """ + self._agentsets = list(other) + + @property + def active_agents(self) -> dict[str, DataFrame]: + return { + agentset.__class__.__name__: agentset.active_agents + for agentset in self._agentsets + } + + @property + def inactive_agents(self): + return { + agentset.__class__.__name__: agentset.inactive_agents + for agentset in self._agentsets + } diff --git a/mesa_frames/concrete/agentset_pandas.py b/mesa_frames/concrete/agentset_pandas.py new file mode 100644 index 0000000..d6481f3 --- /dev/null +++ b/mesa_frames/concrete/agentset_pandas.py @@ -0,0 +1,371 @@ +from typing import ( + Any, + Callable, + Collection, + Hashable, + Iterator, + Self, + Sequence, + overload, +) + +import pandas as pd +import polars as pl + +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 + + +class AgentSetPandas(AgentSetDF): + _agents: pd.DataFrame + _mask: pd.Series + _copy_with_method: dict[str, tuple[str, list[str]]] = { + "_agents": ("copy", ["deep"]), + "_mask": ("copy", ["deep"]), + } + """A pandas-based implementation of the AgentSet. + + Attributes + ---------- + _agents : pd.DataFrame + The agents in the AgentSet. + _copy_only_reference : list[str] = ['_model'] + A list of attributes to copy with a reference only. + _copy_with_method: dict[str, tuple[str, list[str]]] = { + "_agents": ("copy", ["deep"]), + "_mask": ("copy", ["deep"]), + } + A dictionary of attributes to copy with a specified method and arguments. + _mask : pd.Series + A boolean mask indicating which agents are active. + _model : ModelDF + The model that the AgentSetDF belongs to. + + Properties + ---------- + active_agents(self) -> pd.DataFrame + Get the active agents in the AgentSetPandas. + agents(self) -> pd.DataFrame + Get or set the agents in the AgentSetPandas. + inactive_agents(self) -> pd.DataFrame + Get the inactive agents in the AgentSetPandas. + model(self) -> ModelDF + Get the model associated with the AgentSetPandas. + random(self) -> Generator + Get the random number generator associated with the model. + + Methods + ------- + __init__(self, model: ModelDF) -> None + 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 + 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 + 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 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. + set(self, attr_names: str | Collection[str] | dict[str, Any] | None = None, values: Any | None = None, mask: PandasMaskLike | None = None, inplace: bool = True) -> Self + Set the value of a specified attribute or attributes for each agent in the mask in the AgentSetPandas. + shuffle(self, inplace: bool = True) -> Self + Shuffle the order of agents in the AgentSetPandas. + sort(self, by: str | Sequence[str], ascending: bool | Sequence[bool] = True, inplace: bool = True, **kwargs) -> Self + Sort the agents in the AgentSetPandas based on the given criteria. + to_polars(self) -> "AgentSetPolars" + Convert the AgentSetPandas to an AgentSetPolars. + _get_bool_mask(self, mask: PandasMaskLike = None) -> pd.Series + Get a boolean mask for selecting agents. + _get_masked_df(self, mask: PandasMaskLike = None) -> pd.DataFrame + Get a DataFrame of agents that match the mask. + __getattr__(self, key: str) -> pd.Series + Retrieve an attribute of the underlying DataFrame. + __iter__(self) -> Iterator + Get an iterator for the agents in the AgentSetPandas. + __len__(self) -> int + Get the number of agents in the AgentSetPandas. + __repr__(self) -> str + Get the string representation of the AgentSetPandas. + __reversed__(self) -> Iterator + Get a reversed iterator for the agents in the AgentSetPandas. + __str__(self) -> str + Get the string representation of the AgentSetPandas. + """ + + 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) + + def add( + self, + other: pd.DataFrame | Sequence[Any] | dict[str, Any], + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + if isinstance(other, pd.DataFrame): + new_agents = other + if "unique_id" != other.index.name: + try: + new_agents.set_index("unique_id", inplace=True, drop=True) + except KeyError: + raise KeyError("DataFrame must have a unique_id column/index.") + elif isinstance(other, dict): + if "unique_id" not in other: + raise KeyError("Dictionary must have a unique_id key.") + index = other.pop("unique_id") + if not isinstance(index, list): + index = [index] + new_agents = pd.DataFrame(other, index=pd.Index(index, name="unique_id")) + else: + if len(other) != len(obj._agents.columns) + 1: + raise ValueError( + "Length of data must match the number of columns in the AgentSet if being added as a Collection." + ) + columns = pd.Index(["unique_id"]).append(obj._agents.columns.copy()) + new_agents = pd.DataFrame([other], columns=columns).set_index( + "unique_id", drop=True + ) + obj._agents = pd.concat([obj._agents, new_agents]) + return obj + + @overload + def contains(self, ids: Collection[Hashable]) -> pd.Series: ... + + @overload + def contains(self, ids: Hashable) -> bool: ... + + def contains( + self, + ids: Hashable | Collection[Hashable], + ) -> bool | pd.Series: + if isinstance(ids, pd.Series): + return ids.isin(self._agents.index) + elif isinstance(ids, Collection): + return pd.Series(list(ids), index=list(ids)).isin(self._agents.index) + else: + return ids in self._agents.index + + def get( + self, + attr_names: str | Collection[str] | None = None, + mask: PandasMaskLike = None, + ) -> pd.Series | pd.DataFrame: + mask = self._get_bool_mask(mask) + if attr_names is None: + return self._agents.loc[mask] + else: + if isinstance(attr_names, str): + return self._agents.loc[mask, attr_names] + if isinstance(attr_names, Collection): + return self._agents.loc[mask, list(attr_names)] + + def remove( + self, + ids: PandasMaskLike, + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + initial_len = len(obj._agents) + mask = obj._get_bool_mask(ids) + remove_ids = obj._agents[mask].index + obj._agents.drop(remove_ids, inplace=True) + if len(obj._agents) == initial_len: + raise KeyError(f"IDs {ids} not found in agent set.") + return obj + + def set( + self, + attr_names: str | dict[str, Any] | Collection[str] | None = None, + values: Any | None = None, + mask: PandasMaskLike = None, + inplace: bool = True, + ) -> Self: + + obj = self._get_obj(inplace) + b_mask = obj._get_bool_mask(mask) + masked_df = obj._get_masked_df(mask) + + if not attr_names: + attr_names = masked_df.columns + + if isinstance(attr_names, dict): + for key, val in attr_names.items(): + masked_df.loc[:, key] = val + elif ( + isinstance(attr_names, str) + or ( + isinstance(attr_names, Collection) + and all(isinstance(n, str) for n in attr_names) + ) + ) and values is not None: + if not isinstance(attr_names, str): # isinstance(attr_names, Collection) + attr_names = list(attr_names) + masked_df.loc[:, attr_names] = values + else: + raise ValueError( + "Either attr_names must be a dictionary with columns as keys and values or values must be provided." + ) + + non_masked_df = obj._agents[~b_mask] + original_index = obj._agents.index + obj._agents = pd.concat([non_masked_df, masked_df]) + obj._agents = obj._agents.reindex(original_index) + return obj + + def select( + self, + mask: PandasMaskLike = None, + filter_func: Callable[[Self], PandasMaskLike] | None = None, + n: int | None = None, + negate: bool = False, + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + bool_mask = obj._get_bool_mask(mask) + if n != None: + bool_mask = pd.Series( + obj._agents.index.isin(obj._agents[bool_mask].sample(n).index), + index=obj._agents.index, + ) + if filter_func: + bool_mask = bool_mask & obj._get_bool_mask(filter_func(obj)) + if negate: + bool_mask = ~bool_mask + obj._mask = bool_mask + return obj + + def shuffle(self, inplace: bool = True) -> Self: + obj = self._get_obj(inplace) + obj._agents = obj._agents.sample(frac=1) + return obj + + def sort( + self, + by: str | Sequence[str], + ascending: bool | Sequence[bool] = True, + inplace: bool = True, + **kwargs, + ) -> Self: + obj = self._get_obj(inplace) + obj._agents.sort_values(by=by, ascending=ascending, **kwargs, inplace=True) + return obj + + def to_polars(self) -> AgentSetPolars: + new_obj = AgentSetPolars(self._model) + new_obj._agents = pl.DataFrame(self._agents) + new_obj._mask = pl.Series(self._mask) + return new_obj + + def _get_bool_mask( + self, + mask: PandasMaskLike = None, + ) -> pd.Series: + if isinstance(mask, pd.Series) and mask.dtype == bool: + return mask + elif isinstance(mask, pd.DataFrame): + return pd.Series( + self._agents.index.isin(mask.index), index=self._agents.index + ) + elif isinstance(mask, list): + return pd.Series(self._agents.index.isin(mask), index=self._agents.index) + elif mask is None or mask == "all": + return pd.Series(True, index=self._agents.index) + elif mask == "active": + return self._mask + else: + return pd.Series(self._agents.index.isin([mask]), index=self._agents.index) + + def _get_masked_df( + self, + mask: PandasMaskLike = None, + ) -> pd.DataFrame: + if isinstance(mask, pd.Series) and mask.dtype == bool: + return self._agents.loc[mask] + elif isinstance(mask, pd.DataFrame): + if not mask.index.isin(self._agents.index).all(): + raise KeyError( + "Some 'unique_id' of mask are not present in DataFrame 'unique_id'." + ) + if mask.index.name != "unique_id": + if "unique_id" in mask.columns: + mask.set_index("unique_id", inplace=True, drop=True) + else: + raise KeyError("DataFrame must have a unique_id column/index.") + return pd.DataFrame(index=mask.index).join( + self._agents, on="unique_id", how="left" + ) + elif isinstance(mask, pd.Series): + if not mask.isin(self._agents.index).all(): + raise KeyError( + "Some 'unique_id' of mask are not present in DataFrame 'unique_id'." + ) + mask_df = mask.to_frame("unique_id").set_index("unique_id") + return mask_df.join(self._agents, on="unique_id", how="left") + elif mask is None or mask == "all": + return self._agents + elif mask == "active": + return self._agents.loc[self._mask] + else: + mask_series = pd.Series(mask) + if not mask_series.isin(self._agents.index).all(): + raise KeyError( + "Some 'unique_id' of mask are not present in DataFrame 'unique_id'." + ) + mask_df = mask_series.to_frame("unique_id").set_index("unique_id") + return mask_df.join(self._agents, on="unique_id", how="left") + + def __getattr__(self, name: str) -> Any: + super().__getattr__(name) + return getattr(self._agents, name) + + def __iter__(self) -> Iterator: + return iter(self._agents.iterrows()) + + def __len__(self) -> int: + return len(self._agents) + + def __repr__(self) -> str: + return repr(self._agents) + + def __reversed__(self) -> Iterator: + return iter(self._agents[::-1].iterrows()) + + def __str__(self) -> str: + return str(self._agents) + + @property + def agents(self) -> pd.DataFrame: + return self._agents + + @agents.setter + def agents(self, new_agents: pd.DataFrame) -> None: + if new_agents.index.name == "unique_id": + pass + elif "unique_id" in new_agents.columns: + new_agents.set_index("unique_id", inplace=True, drop=True) + else: + raise KeyError("The DataFrame should have a 'unique_id' index/column") + self._agents = new_agents + + @property + def active_agents(self) -> pd.DataFrame: + return self._agents.loc[self._mask] + + @active_agents.setter + def active_agents(self, mask: PandasMaskLike) -> None: + self.select(mask=mask, inplace=True) + + @property + def inactive_agents(self) -> pd.DataFrame: + return self._agents.loc[~self._mask] diff --git a/mesa_frames/concrete/agentset_polars.py b/mesa_frames/concrete/agentset_polars.py new file mode 100644 index 0000000..a4939d2 --- /dev/null +++ b/mesa_frames/concrete/agentset_polars.py @@ -0,0 +1,489 @@ +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Hashable, + Iterable, + Iterator, + Literal, + Self, + Sequence, + overload, +) + +import polars as pl +from polars.type_aliases import IntoExpr + +from mesa_frames.abstract.agents import AgentSetDF +from mesa_frames.types import PolarsMaskLike + +if TYPE_CHECKING: + from mesa_frames.concrete.agentset_pandas import AgentSetPandas + from mesa_frames.concrete.model import ModelDF + + +class AgentSetPolars(AgentSetDF): + _agents: pl.DataFrame + _copy_with_method: dict[str, tuple[str, list[str]]] = { + "_agents": ("clone", []), + } + _copy_only_reference: list[str] = ["_model", "_mask"] + _mask: pl.Expr | pl.Series + _model: "ModelDF" + + """A polars-based implementation of the AgentSet. + + Attributes + ---------- + _agents : pl.DataFrame + The agents in the AgentSet. + _copy_only_reference : list[str] = ["_model", "_mask"] + A list of attributes to copy with a reference only. + _copy_with_method: dict[str, tuple[str, list[str]]] = { + "_agents": ("copy", ["deep"]), + "_mask": ("copy", ["deep"]), + } + A dictionary of attributes to copy with a specified method and arguments. + model : ModelDF + The model to which the AgentSet belongs. + _mask : pl.Series + A boolean mask indicating which agents are active. + + Properties + ---------- + active_agents(self) -> pl.DataFrame + Get the active agents in the AgentSetPolars. + agents(self) -> pl.DataFrame + Get or set the agents in the AgentSetPolars. + inactive_agents(self) -> pl.DataFrame + Get the inactive agents in the AgentSetPolars. + model(self) -> ModelDF + Get the model associated with the AgentSetPolars. + random(self) -> Generator + Get the random number generator associated with the model. + + + Methods + ------- + __init__(self, model: ModelDF) -> None + 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 + 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 + 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 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. + set(self, attr_names: str | Collection[str] | dict[str, Any] | None = None, values: Any | None = None, mask: PolarsMaskLike | None = None, inplace: bool = True) -> Self + Set the value of a specified attribute or attributes for each agent in the mask in the AgentSetPolars. + shuffle(self, inplace: bool = True) -> Self + Shuffle the order of agents in the AgentSetPolars. + sort(self, by: str | Sequence[str], ascending: bool | Sequence[bool] = True, inplace: bool = True, **kwargs) -> Self + Sort the agents in the AgentSetPolars based on the given criteria. + to_pandas(self) -> "AgentSetPandas" + Convert the AgentSetPolars to an AgentSetPandas. + _get_bool_mask(self, mask: PolarsMaskLike = None) -> pl.Series | pl.Expr + Get a boolean mask for selecting agents. + _get_masked_df(self, mask: PolarsMaskLike = None) -> pl.DataFrame + Get a DataFrame of agents that match the mask. + __getattr__(self, key: str) -> pl.Series + Retrieve an attribute of the underlying DataFrame. + __iter__(self) -> Iterator + Get an iterator for the agents in the AgentSetPolars. + __len__(self) -> int + Get the number of agents in the AgentSetPolars. + __repr__(self) -> str + Get the string representation of the AgentSetPolars. + __reversed__(self) -> Iterator + Get a reversed iterator for the agents in the AgentSetPolars. + __str__(self) -> str + Get the string representation of the AgentSetPolars. + + """ + + def __init__(self, model: "ModelDF") -> None: + """Initialize a new AgentSetPolars. + + Parameters + ---------- + model : ModelDF + The model that the agent set belongs to. + + Returns + ------- + None + """ + self._model = model + self._agents = pl.DataFrame(schema={"unique_id": pl.Int64}) + self._mask = pl.repeat(True, len(self._agents)) + + def add( + self, + other: pl.DataFrame | Sequence[Any] | dict[str, Any], + inplace: bool = True, + ) -> Self: + """Add agents to the AgentSetPolars. + + Parameters + ---------- + other : pl.DataFrame | Sequence[Any] | dict[str, Any] + The agents to add. + inplace : bool, optional + Whether to add the agents in place, by default True. + + Returns + ------- + Self + The updated AgentSetPolars. + """ + obj = self._get_obj(inplace) + if isinstance(other, pl.DataFrame): + if "unique_id" not in other.columns: + raise KeyError("DataFrame must have a unique_id column.") + new_agents = other + elif isinstance(other, dict): + if "unique_id" not in other: + raise KeyError("Dictionary must have a unique_id key.") + new_agents = pl.DataFrame(other) + else: + if len(other) != len(obj._agents.columns): + raise ValueError( + "Length of data must match the number of columns in the AgentSet if being added as a Collection." + ) + new_agents = pl.DataFrame([other], schema=obj._agents.schema) + obj._agents = pl.concat([obj._agents, new_agents], how="diagonal_relaxed") + return obj + + @overload + def contains(self, ids: Collection[Hashable]) -> pl.Series: ... + + @overload + def contains(self, ids: Hashable) -> bool: ... + + def contains( + self, + ids: Hashable | Collection[Hashable], + ) -> bool | pl.Series: + if isinstance(ids, pl.Series): + return ids.is_in(self._agents["unique_id"]) + elif isinstance(ids, Collection): + return pl.Series(ids).is_in(self._agents["unique_id"]) + else: + return ids in self._agents["unique_id"] + + def discard(self, ids: PolarsMaskLike, inplace: bool = True) -> Self: + """Remove an agent from the AgentSetPolars. Does not raise an error if the agent is not found. + + Parameters + ---------- + ids : PolarsMaskLike + The mask of agents to remove. + inplace : bool, optional + Whether to remove the agents in place, by default True. + + Returns + ------- + Self + The updated AgentSetPolars. + """ + return super().discard(ids=ids, inplace=inplace) + + def get( + self, + attr_names: IntoExpr | Iterable[IntoExpr] | None, + mask: PolarsMaskLike = None, + ) -> pl.Series | pl.DataFrame: + masked_df = self._get_masked_df(mask) + attr_names = self.agents.select(attr_names).columns.copy() + if not attr_names: + return masked_df + masked_df = masked_df.select(attr_names) + if masked_df.shape[1] == 1: + return masked_df[masked_df.columns[0]] + return masked_df + + def remove(self, ids: PolarsMaskLike, inplace: bool = True) -> Self: + obj = self._get_obj(inplace=inplace) + initial_len = len(obj._agents) + mask = obj._get_bool_mask(ids) + obj._agents = obj._agents.filter(mask.not_()) + if len(obj._agents) == initial_len: + raise KeyError(f"IDs {ids} not found in agent set.") + return obj + + def set( + self, + attr_names: str | Collection[str] | dict[str, Any] | None = None, + values: Any | None = None, + mask: PolarsMaskLike = None, + inplace: bool = True, + ) -> Self: + + obj = self._get_obj(inplace) + b_mask = obj._get_bool_mask(mask) + masked_df = obj._get_masked_df(mask) + + if not attr_names: + attr_names = masked_df.columns + attr_names.remove("unique_id") + + def process_single_attr( + masked_df: pl.DataFrame, attr_name: str, values: Any + ) -> pl.DataFrame: + if isinstance(values, pl.DataFrame): + return masked_df.with_columns(values.to_series().alias(attr_name)) + elif isinstance(values, pl.Expr): + return masked_df.with_columns(values.alias(attr_name)) + if isinstance(values, pl.Series): + return masked_df.with_columns(values.alias(attr_name)) + else: + if isinstance(values, Collection): + values = pl.Series(values) + else: + values = pl.repeat(values, len(masked_df)) + return masked_df.with_columns(values.alias(attr_name)) + + if isinstance(attr_names, str) and values is not None: + masked_df = process_single_attr(masked_df, attr_names, values) + elif isinstance(attr_names, Collection) and values is not None: + if isinstance(values, Collection) and len(attr_names) == len(values): + for attribute, val in zip(attr_names, values): + masked_df = process_single_attr(masked_df, attribute, val) + else: + for attribute in attr_names: + masked_df = process_single_attr(masked_df, attribute, values) + elif isinstance(attr_names, dict): + for key, val in attr_names.items(): + masked_df = process_single_attr(masked_df, key, val) + else: + raise ValueError( + "attr_names must be a string, a collection of string or a dictionary with columns as keys and values." + ) + non_masked_df = obj._agents.filter(b_mask.not_()) + original_index = obj._agents.select("unique_id") + obj._agents = pl.concat([non_masked_df, masked_df], how="diagonal_relaxed") + obj._agents = original_index.join(obj._agents, on="unique_id", how="left") + return obj + + def select( + self, + mask: PolarsMaskLike = None, + filter_func: Callable[[Self], pl.Series] | None = None, + n: int | None = None, + negate: bool = False, + inplace: bool = True, + ) -> Self: + obj = self._get_obj(inplace) + mask = obj._get_bool_mask(mask) + if filter_func: + mask = mask & filter_func(obj) + if n != None: + mask = (obj._agents["unique_id"]).is_in( + obj._agents.filter(mask).sample(n)["unique_id"] + ) + if negate: + mask = mask.not_() + obj._mask = mask + return obj + + def shuffle(self, inplace: bool = True) -> Self: + obj = self._get_obj(inplace) + obj._agents = obj._agents.sample(fraction=1, shuffle=True) + return obj + + def sort( + self, + by: str | Sequence[str], + ascending: bool | Sequence[bool] = True, + inplace: bool = True, + **kwargs, + ) -> Self: + obj = self._get_obj(inplace) + if isinstance(ascending, bool): + descending = not ascending + else: + descending = [not a for a in ascending] + obj._agents = obj._agents.sort(by=by, descending=descending, **kwargs) + return obj + + def to_pandas(self) -> "AgentSetPandas": + from mesa_frames.concrete.agentset_pandas import AgentSetPandas + + new_obj = AgentSetPandas(self._model) + new_obj._agents = self._agents.to_pandas() + if isinstance(self._mask, pl.Series): + new_obj._mask = self._mask.to_pandas() + else: # self._mask is Expr + new_obj._mask = ( + self._agents["unique_id"] + .is_in(self._agents.filter(self._mask)["unique_id"]) + .to_pandas() + ) + return new_obj + + def _get_bool_mask( + self, + mask: PolarsMaskLike = None, + ) -> pl.Series | pl.Expr: + + def bool_mask_from_series(mask: pl.Series) -> pl.Series: + if ( + isinstance(mask, pl.Series) + and mask.dtype == pl.Boolean + and len(mask) == len(self._agents) + ): + return mask + else: + if not mask.is_in(self._agents["unique_id"]).all(): + raise KeyError( + "Some 'unique_ids' of mask are not present in DataFrame 'unique_id'." + ) + return self._agents["unique_id"].is_in(mask) + + if isinstance(mask, pl.Expr): + return mask + elif isinstance(mask, pl.Series): + return bool_mask_from_series(mask) + elif isinstance(mask, pl.DataFrame): + if "unique_id" in mask.columns: + return bool_mask_from_series(mask["unique_id"]) + elif len(mask.columns) == 1 and mask.dtypes[0] == pl.Boolean: + return bool_mask_from_series(mask[mask.columns[0]]) + else: + raise KeyError( + "DataFrame must have a 'unique_id' column or a single boolean column." + ) + elif mask is None or mask == "all": + return pl.repeat(True, len(self._agents)) + elif mask == "active": + return self._mask + elif isinstance(mask, Collection): + return bool_mask_from_series(pl.Series(mask)) + else: + return bool_mask_from_series(pl.Series([mask])) + + def _get_masked_df( + self, + mask: PolarsMaskLike = None, + ) -> pl.DataFrame: + if (isinstance(mask, pl.Series) and mask.dtype == bool) or isinstance( + mask, pl.Expr + ): + return self._agents.filter(mask) + elif isinstance(mask, pl.DataFrame): + if not mask["unique_id"].is_in(self._agents["unique_id"]).all(): + raise KeyError( + "Some 'unique_id' of mask are not present in DataFrame 'unique_id'." + ) + return mask.select("unique_id").join( + self._agents, on="unique_id", how="left" + ) + elif isinstance(mask, pl.Series): + if not mask.is_in(self._agents["unique_id"]).all(): + raise KeyError( + "Some 'unique_id' of mask are not present in DataFrame 'unique_id'." + ) + mask_df = mask.to_frame("unique_id") + return mask_df.join(self._agents, on="unique_id", how="left") + elif mask is None or mask == "all": + return self._agents + elif mask == "active": + return self._agents.filter(self._mask) + else: + if isinstance(mask, Collection): + mask_series = pl.Series(mask) + else: + mask_series = pl.Series([mask]) + if not mask_series.is_in(self._agents["unique_id"]).all(): + raise KeyError( + "Some 'unique_id' of mask are not present in DataFrame 'unique_id'." + ) + mask_df = mask_series.to_frame("unique_id") + return mask_df.join(self._agents, on="unique_id", how="left") + + def __getattr__(self, key: str) -> pl.Series: + super().__getattr__(key) + return self._agents[key] + + @overload + def __getitem__( + self, + key: str | tuple[PolarsMaskLike, str], + ) -> pl.Series: ... + + @overload + def __getitem__( + self, + key: ( + PolarsMaskLike + | Collection[str] + | tuple[ + PolarsMaskLike, + Collection[str], + ] + ), + ) -> pl.DataFrame: ... + + def __getitem__( + self, + key: ( + str + | Collection[str] + | PolarsMaskLike + | tuple[PolarsMaskLike, str] + | tuple[ + PolarsMaskLike, + Collection[str], + ] + ), + ) -> pl.Series | pl.DataFrame: + attr = super().__getitem__(key) + assert isinstance(attr, (pl.Series, pl.DataFrame)) + return attr + + def __iter__(self) -> Iterator: + return iter(self._agents.iter_rows(named=True)) + + def __len__(self) -> int: + return len(self._agents) + + def __repr__(self) -> str: + return repr(self._agents) + + def __reversed__(self) -> Iterator: + return reversed(iter(self._agents.iter_rows(named=True))) + + def __str__(self) -> str: + return str(self._agents) + + @property + def agents(self) -> pl.DataFrame: + return self._agents + + @agents.setter + def agents(self, agents: pl.DataFrame) -> None: + if "unique_id" not in agents.columns: + raise KeyError("DataFrame must have a unique_id column.") + self._agents = agents + + @property + def active_agents(self) -> pl.DataFrame: + return self.agents.filter(self._mask) + + @active_agents.setter + def active_agents(self, mask: PolarsMaskLike) -> None: + self.select(mask=mask, inplace=True) + + @property + def inactive_agents(self) -> pl.DataFrame: + return self.agents.filter(~self._mask) diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py new file mode 100644 index 0000000..4525b63 --- /dev/null +++ b/mesa_frames/concrete/model.py @@ -0,0 +1,145 @@ +from typing import Any, Sequence + +import numpy as np + +from mesa_frames.abstract.agents import AgentSetDF +from mesa_frames.concrete.agents import AgentsDF + + +class ModelDF: + """Base class for models in the mesa-frames library. + + This class serves as a foundational structure for creating agent-based models. + It includes the basic attributes and methods necessary for initializing and + running a simulation model. + + Attributes + ---------- + running : bool + A boolean indicating if the model should continue running. + schedule : Any + An object to manage the order and execution of agent steps. + current_id : int + A counter for assigning unique IDs to agents. + _agents : AgentsDF + A mapping of each agent type to a dict of its instances. + + Properties + ---------- + agents : AgentsDF + An AgentSet containing all agents in the model, generated from the _agents attribute. + agent_types : list of type + A list of different agent types present in the model. + + Methods + ------- + __new__(cls, seed: int | Sequence[int] | None = None, *args: Any, **kwargs: Any) -> Any + Create a new model object and instantiate its RNG automatically. + __init__(self, *args: Any, **kwargs: Any) -> None + Create a new model. Overload this method with the actual code to start the model. + get_agents_of_type(self, agent_type: type) -> AgentSetDF + Retrieve the AgentSetDF of a specified type. + initialize_data_collector(self, model_reporters: dict | None = None, agent_reporters: dict | None = None, tables: dict | None = None) -> None + Initialize the data collector for the model (not implemented yet). + next_id(self) -> int + Generate and return the next unique identifier for an agent (not implemented yet). + reset_randomizer(self, seed: int | Sequence[int] | None) -> None + Reset the model random number generator with a new or existing seed. + run_model(self) -> None + Run the model until the end condition is reached. + step(self) -> None + Execute a single step of the model's simulation process (needs to be overridden in a subclass). + """ + + random: np.random.Generator + _seed: int | Sequence[int] + running: bool + _agents: AgentsDF + + def __new__( + cls, seed: int | Sequence[int] | None = None, *args: Any, **kwargs: Any + ) -> Any: + """Create a new model object and instantiate its RNG automatically.""" + obj = object.__new__(cls) + obj.reset_randomizer(seed) + return obj + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Create a new model. Overload this method with the actual code to + start the model. Always start with super().__init__() to initialize the + model object properly. + """ + self.running = True + self.schedule = None + self.current_id = 0 + self._agents = AgentsDF() + + def get_agents_of_type(self, agent_type: type) -> AgentSetDF: + """Retrieve the AgentSetDF of a specified type. + + Parameters + ---------- + agent_type : type + The type of AgentSetDF to retrieve. + + Returns + ------- + AgentSetDF + The AgentSetDF of the specified type. + """ + for agentset in self._agents._agentsets: + if agent_type == type(agentset): + return agentset + raise ValueError(f"No agents of type {agent_type} found in the model.") + + def initialize_data_collector( + self, + model_reporters: dict | None = None, + agent_reporters: dict | None = None, + tables: dict | None = None, + ) -> None: + raise NotImplementedError( + "initialize_data_collector() method not implemented yet for ModelDF" + ) + + def next_id(self) -> int: + raise NotImplementedError("next_id() method not implemented for ModelDF") + + def reset_randomizer(self, seed: int | Sequence[int] | None) -> None: + """Reset the model random number generator. + + Parameters: + ---------- + seed : int | None + A new seed for the RNG; if None, reset using the current seed + """ + if seed is None: + seed = np.random.SeedSequence().entropy + assert seed is not None + self._seed = seed + self.random = np.random.default_rng(seed=self._seed) + + def run_model(self) -> None: + """Run the model until the end condition is reached. Overload as + needed. + """ + while self.running: + self.step() + + def step(self) -> None: + """A single step. Fill in here.""" + raise NotImplementedError("step() method needs to be overridden in a subclass.") + + @property + def agents(self) -> AgentsDF: + return self._agents + + @agents.setter + def agents(self, agents: AgentsDF) -> None: + if not isinstance(agents, AgentsDF): + raise TypeError("agents must be an instance of AgentsDF") + self._agents = agents + + @property + def agent_types(self) -> list[type]: + return [agent.__class__ for agent in self._agents._agentsets] diff --git a/mesa_frames/model.py b/mesa_frames/model.py deleted file mode 100644 index 70545f6..0000000 --- a/mesa_frames/model.py +++ /dev/null @@ -1,397 +0,0 @@ -from logging import warn -from typing import Any - -import numpy as np - -from .agent import AgentsDF, AgentSetDF - - -class ModelDF: - random: np.random.Generator - _seed: int - running: bool - _agents: AgentsDF | None - - def __new__(cls, seed: int = 0, *args: Any, **kwargs: Any) -> Any: - """Create a new model object and instantiate its RNG automatically.""" - obj = object.__new__(cls) - if seed == 0: - seed = np.random.SeedSequence().entropy # type: ignore - obj._seed = seed - obj.random = np.random.default_rng(seed=obj._seed) - return obj - - def __init__(self, *args: Any, **kwargs: Any) -> None: - """Create a new model. Overload this method with the actual code to - start the model. Always start with super().__init__() to initialize the - model object properly. - """ - self.running: bool = True - self._agents: AgentsDF | None = None - - @property - def agents(self) -> AgentsDF | None: - return self._agents - - @agents.setter - def agents(self, agents: AgentsDF) -> None: - if not isinstance(agents, AgentsDF): - raise TypeError("agents must be an instance of AgentsDF") - if type(agents) != type(self._agents): - warn( - f"Changing the backend from {type(self._agents)} to {type(agents)}", - RuntimeWarning, - ) - self._agents = agents - - def get_agents_of_type(self, agent_type: type) -> AgentSetDF: - """Retrieve the AgentSetDF of a specified type. - - Parameters - ---------- - agent_type : type - The type of AgentSetDF to retrieve. - - Returns - ------- - AgentSetDF - The AgentSetDF of the specified type. - """ - return self.agents.get_agents_of_type(agent_type) - - def next_id(self) -> int: - raise NotImplementedError("next_id() method not implemented for ModelDF") - - def reset_randomizer(self, seed: int | None = None) -> None: - """Reset the model random number generator. - - Parameters: - ---------- - seed : int | None - A new seed for the RNG; if None, reset using the current seed - """ - - self._seed = np.random.SeedSequence(seed=seed).entropy # type: ignore - - """'def initialize_data_collector( - self, model_reporters=None, agent_reporters=None, tables=None - ) -> None: - if not self._agents: - raise RuntimeError( - "You must create agents before initializing the data collector." - ) - self.datacollector = DataCollectorDF( - model_reporters=model_reporters, - agent_reporters=agent_reporters, - tables=tables, - ) - self.datacollector.collect(self)""" - - -''' -OLD IMPLEMENTATION (HAS TO BE DELETED) - -class ModelDF: - """The base class for all models - - Attributes - ---------- - unique_id : int - The unique_id of the model. - running : bool - Indicates if the model is running or not. - agents : pd.DataFrame | gpd.GeoDataFrame | None - The dataframe containing the agents of the model. - agent_types : list[tuple[type[AgentDF], float]] | None - The list of agent types and their proportions. - p_agents : dict[type[AgentDF], float] | None - The dictionary of agents to create. The keys are the types of agents, - the values are the percentages of each agent type. The sum of the values should be 1. - space - The space where the agents will be placed. Can be None if model does not have a space. - """ - - def __new__(cls, *args, **kwargs): - """Create a new model object and instantiate its RNG automatically - (adds supports to numpy with respect to base model).""" - obj = object.__new__(cls) - obj._seed = kwargs.get("seed") - if obj._seed is None: - # We explicitly specify the seed here so that we know its value in - # advance. - obj._seed = np.random.SeedSequence().entropy - # Use default_rng to get a new Generator instance - obj.random = np.random.default_rng(seed = obj._seed) - return obj - - def __init__(self, unique_id: int | None = None, space=None): - """Create a new model. Overload this method with the actual code to - start the model. Always start with super().__init__() to initialize the - model object properly. - - Parameters - ---------- - unique_id : int | None - The unique_id of the model. - If None, a random unique_id is assigned using a 64-bit random integer. - space - The space where the agents will be placed. Can be None if model does not have a space. - """ - # Initialize default attributes - self.running: bool = True - self.agents: pd.DataFrame | gpd.GeoDataFrame | None = None - self.agent_types: list[tuple[type[AgentDF], float]] | None = None - self.p_agents: dict[type[AgentDF], float] | None = None - # self.schedule : BaseScheduler = None - - # Initialize optional parameters - if not unique_id: - self.unique_id = np.random.randint( - low=-9223372036854775808, high=9223372036854775807, dtype="int64" - ) - else: - self.unique_id = unique_id - self.space = space - - # Initialize data collection - # self.initialize_data_collector(data_collection) - - def get_agents_of_type(self, agent_type: type[AgentDF]) -> pd.Series: - """Returns a boolean mask of the agents dataframe of the model which corresponds to the agent_type. - - Parameters - ---------- - agent_type : type[AgentDF] - The agent_type to get the mask for. - """ - if self.agents is None: - raise RuntimeError( - "You must create agents before getting their masks. Use create_agents() method." - ) - return self.agents["type"].str.contains(agent_type.__name__) # type: ignore - - def run_model(self, n_steps: int | None = None, merged_mro: bool = False) -> None: - """If n_steps are specified, executes model.step() until n_steps are reached. - Otherwise, until self.running is false (as the default mesa.Model.run_model). - - Parameters - ---------- - n_steps : int | None - The number of steps which the model will execute. - Can be None if a running condition turning false is used. - merged_mro: bool - If False, the model will execute one step for each class in p_agent. This is the default behaviour. - If True, the model will execute one step for each inherited agent type in the order of a "merged" MRO. - This may increase performance if there are multiple and complex inheritance as each agent_type (even if parents of different classes), - will be executed only once. Not a viable option if the behavior of a class depends on another. - """ - if n_steps: - if not (isinstance(n_steps, int) and n_steps > 0): - raise TypeError( - "n_steps should be an integer greater than 0 or None if a running condition is used" - ) - for _ in range(n_steps): - self.step(merged_mro) - else: - while self.running: - self.step(merged_mro) - - def step(self, merged_mro: bool = False) -> None: - """Executes one step of the model. - - Parameters - ---------- - merged_mro : bool - If False, the model will execute one step for each class in p_agent. This is the default behaviour. - If True, the model will execute one step for each inherited agent type in the order of a "merged" MRO. - This may increase performance if there are multiple and complex inheritance as each agent_type (even if parents of different classes), - will be executed only once. Not a viable option if the behavior of a class depends on another. - """ - if self.agent_types is None or self.p_agents is None: - raise RuntimeError( - "You must create agents before running the model. Use create_agents() method." - ) - if merged_mro: - for agent_type in self.agent_types: - agent_type[0].step() - else: - for agent in self.p_agents: - agent.step() - - def reset_randomizer(self, seed: int | None = None) -> None: - """Reset the model random number generator. - - Parameters - ---------- - seed : int | None - A new seed for the RNG; if None, reset using the current seed - """ - if seed is None: - seed = self._seed - self.random = np.random.default_rng(seed) - self._seed = seed - - def create_agents( - self, n_agents: int, p_agents: dict[type[AgentDF], float] - ) -> None: - """Populate the self.agents dataframe. - - Parameters - ---------- - n_agents : int | None - The number of agents which the model will create. - p_agents : dict[type[AgentDF], float] - The dictionary of agents to create. The keys are the types of agents, - the values are the percentages of each agent type. The sum of the values should be 1. - """ - - # Verify parameters - if not (isinstance(n_agents, int) and n_agents > 0): - raise TypeError("n_agents should be an integer greater than 0") - if sum(p_agents.values()) != 1: - raise ValueError("Sum of proportions of agents should be 1") - if any(p < 0 or p > 1 for p in p_agents.values()): - raise ValueError("Proportions of agents should be between 0 and 1") - - self.p_agents = p_agents - - start_time = time() - print("Creating agents: ...") - - mros = [[agent.__mro__[:-1], p] for agent, p in p_agents.items()] - mros_copy = deepcopy(mros) - agent_types = [] - - # Create a "merged MRO" (inspired by C3 linearization algorithm) - while True: - candunique_idate_added = False - # if all mros are empty, the merged mro is done - if not any(mro[0] for mro in mros): - break - for mro in mros: - # If mro is empty, continue - if not mro[0]: - continue - # candunique_idate = head - candunique_idate = mro[0][0] - # If candunique_idate appears in the tail of another MRO, skip it for now (because other agent_types depend on it, will be added later) - if any( - candunique_idate in other_mro[0][1:] - for other_mro in mros - if other_mro is not mro - ): - continue - else: - p = 0 - for i, other_mro in enumerate(mros): - if other_mro[0][0] == candunique_idate: - p += other_mro[1] - mros[i][0] = other_mro[0][1:] - else: - continue - agent_types.append((candunique_idate, p)) # Safe to add it - candunique_idate_added = True - # If there wasn't any good head, there is an inconsistent hierarchy - if not candunique_idate_added: - raise ValueError("Inconsistent hierarchy") - self.agent_types = list(agent_types) - - # Create a single DF using vars and values for every class - columns: set[str] = set() - dtypes: dict[str, str] = {} - for agent_type in self.agent_types: - for key, val in agent_type[0].dtypes.items(): - if key not in columns: - columns.add(key) - dtypes[key] = val - - if "geometry" in columns: - if not (self.space and hasattr(self.space, "crs")): - raise ValueError( - "You must specify a space with a crs attribute if you want to create GeoAgents" - ) - self.agents = gpd.GeoDataFrame( - index=pd.RangeIndex(0, n_agents), - columns=list(columns), - crs=self.space.crs, - ) - else: - self.agents = pd.DataFrame( - index=pd.RangeIndex(0, n_agents), columns=list(columns) - ) - - # Populate agents type - start_index = 0 - for i, (_, p) in enumerate(p_agents.items()): - self.agents.loc[ - start_index : start_index + int(n_agents * p) - 1, "type" - ] = str(mros_copy[i][0]) - start_index += int(n_agents * p) - - # Initialize agents - AgentDF.model = self - self.update_agents_masks() - for agent in p_agents: - agent.__init__() - - # Set dtypes - for col, dtype in dtypes.items(): - if "int" in dtype and self.agents[col].isna().sum() > 0: # type: ignore - warn( - f"Pandas does not support NaN values for int{dtype[-2:]} dtypes. Changing dtype to float{dtype[-2:]} for {col}", - RuntimeWarning, - ) - dtypes[col] = "float" + dtype[-2:] - self.agents = self.agents.astype(dtypes) - - # Set agents' unique_id as index (Have to reassign masks because index changed) - self.agents.set_index("id", inplace=True) - self.update_agents_masks() - - print("Created agents: " + "--- %s seconds ---" % (time() - start_time)) - - def _initialize_data_collection(self, how="2d") -> None: - """Initializes the data collection of the model. - - Parameters - ---------- - how : str - The frequency of the data collection. It can be 'xd', 'xd', 'xh', 'weekly', 'daily', 'hourly'. - """ - # TODO: finish implementation of different data collections - if how == "2d": - return - - def update_agents_masks(self) -> None: - """Updates the masks attributes of each agent in self.agent_types. - Useful after agents are created/deleted or index changes. - """ - if self.agent_types is None: - raise RuntimeError( - "You must create agents before updating their masks. Use create_agents() method." - ) - for agent_type in self.agent_types: - agent_type[0].mask = self.get_agents_of_type(agent_type[0]) - - # TODO: implement different data collection frequencies (xw, xd, xh, weekly, daily, hourly, per every step): - """def initialize_data_collector( - self, - model_reporters=None, - agent_reporters=None, - tables=None, - ) -> None: - if not hasattr(self, "schedule") or self.schedule is None: - raise RuntimeError( - "You must initialize the scheduler (self.schedule) before initializing the data collector." - ) - if self.schedule.get_agent_count() == 0: - raise RuntimeError( - "You must add agents to the scheduler before initializing the data collector." - ) - self.datacollector = DataCollector( - model_reporters=model_reporters, - agent_reporters=agent_reporters, - tables=tables, - ) - # Collect data for the first time during initialization. - self.datacollector.collect(self)''' diff --git a/mesa_frames/pandas/agent.py b/mesa_frames/pandas/agent.py deleted file mode 100644 index c309eb2..0000000 --- a/mesa_frames/pandas/agent.py +++ /dev/null @@ -1,337 +0,0 @@ -import pandas as pd - -from .base import AgentSetDF - - -class AgentSetPandas(AgentSetDF): - _agents: pd.DataFrame - _mask: pd.Series[bool] - """A pandas-based implementation of the AgentSet. - - Attributes - ---------- - model : ModelDF - The model to which the AgentSet belongs. - _mask : pd.Series[bool] - A boolean mask indicating which agents are active. - _agents : pd.DataFrame - The agents in the AgentSet. - _skip_copy : list[str] - A list of attributes to skip during the copy process. - """ - - def __new__(cls, model: ModelDF) -> Self: - obj = super().__new__(cls, model) - obj._agents = pd.DataFrame(columns=["unique_id"]).set_index("unique_id") - obj._mask = pd.Series(True, index=obj._agents.index) - return obj - - def __contains__(self, id: Hashable) -> bool: - return id in self._agents.index - - def __deepcopy__(self, memo: dict) -> Self: - obj = super().__deepcopy__(memo) - obj._agents = self._agents.copy(deep=True) - return obj - - def __iter__(self): - return self._agents.iterrows() - - def __reversed__(self) -> Iterable: - return self._agents[::-1].iterrows() - - @property - def agents(self) -> pd.DataFrame: - return self._agents - - @property - def active_agents(self) -> pd.DataFrame: - return self._agents.loc[self._mask] - - @active_agents.setter # When a property is overriden, so it is the getter - def active_agents(self, mask: PandasMaskLike) -> None: - return AgentContainer.active_agents.fset(self, mask) # type: ignore - - @property - def inactive_agents(self) -> pd.DataFrame: - return self._agents.loc[~self._mask] - - def _get_bool_mask( - self, - mask: PandasMaskLike | None = None, - ) -> pd.Series: - if isinstance(mask, pd.Series) and mask.dtype == bool: - return mask - elif isinstance(mask, self.__class__): - return pd.Series( - self._agents.index.isin(mask.agents.index), index=self._agents.index - ) - elif isinstance(mask, pd.DataFrame): - return pd.Series( - self._agents.index.isin(mask.index), index=self._agents.index - ) - elif isinstance(mask, list): - return pd.Series(self._agents.index.isin(mask), index=self._agents.index) - elif mask is None or mask == "all": - return pd.Series(True, index=self._agents.index) - elif mask == "active": - return self._mask - else: - return pd.Series(self._agents.index.isin([mask]), index=self._agents.index) - - def copy( - self, - deep: bool = False, - skip: list[str] | str | None = None, - memo: dict | None = None, - ) -> Self: - obj = super().copy(deep, skip, memo) - obj._agents = self._agents.copy(deep=deep) - obj._mask = self._mask.copy(deep=deep) - return obj - - def select( - self, - mask: PandasMaskLike | None = None, - filter_func: Callable[[Self], PandasMaskLike] | None = None, - n: int | None = None, - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) - bool_mask = obj._get_bool_mask(mask) - if n != None: - bool_mask = pd.Series( - obj._agents.index.isin(obj._agents[bool_mask].sample(n).index), - index=obj._agents.index, - ) - if filter_func: - bool_mask = bool_mask & obj._get_bool_mask(filter_func(obj)) - obj._mask = bool_mask - return obj - - def shuffle(self, inplace: bool = True) -> Self: - obj = self._get_obj(inplace) - obj._agents = obj._agents.sample(frac=1) - return obj - - def sort( - self, - by: str | Sequence[str], - key: ValueKeyFunc | None = None, - ascending: bool | Sequence[bool] = True, - inplace: bool = True, - ) -> Self: - """ - Sort the agents in the agent set based on the given criteria. - - Parameters - ---------- - by : str | Sequence[str] - The attribute(s) to sort by. - key : ValueKeyFunc | None - A function to use for sorting. - ascending : bool | Sequence[bool] - Whether to sort in ascending order. - - Returns - ---------- - AgentSetDF: The sorted agent set. - """ - obj = self._get_obj(inplace) - obj._agents.sort_values(by=by, key=key, ascending=ascending, inplace=True) - return obj - - @overload - def set_attribute( - self, - attr_names: None = None, - value: Any = Any, - mask: PandasMaskLike = PandasMaskLike, - inplace: bool = True, - ) -> Self: ... - - @overload - def set_attribute( - self, - attr_names: dict[str, Any], - value: None, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - @overload - def set_attribute( - self, - attr_names: str | list[str], - value: Any, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> Self: ... - - def set_attribute( - self, - attr_names: str | list[str] | dict[str, Any] | None = None, - value: Any | None = None, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) - mask = obj._get_bool_mask(mask) - if attr_names is None: - attr_names = obj._agents.columns.values.tolist() - if isinstance(attr_names, (str, list)) and value is not None: - obj._agents.loc[mask, attr_names] = value - elif isinstance(attr_names, dict): - for key, value in attr_names.items(): - obj._agents.loc[mask, key] = value - else: - raise ValueError( - "attr_names must be a string or a dictionary with columns as keys and values." - ) - return obj - - @overload - def get_attribute( - self, - attr_names: list[str] | None = None, - mask: PandasMaskLike | None = None, - ) -> pd.DataFrame: ... - - @overload - def get_attribute( - self, - attr_names: str, - mask: PandasMaskLike | None = None, - ) -> pd.Series: ... - - def get_attribute( - self, - attr_names: str | list[str] | None = None, - mask: PandasMaskLike | None = None, - inplace: bool = True, - ) -> pd.Series | pd.DataFrame: - obj = self._get_obj(inplace) - mask = obj._get_bool_mask(mask) - if attr_names is None: - return obj._agents.loc[mask] - else: - return obj._agents.loc[mask, attr_names] - - def add( - self, - other: Self | pd.DataFrame | ListLike | dict[str, Any], - inplace: bool = True, - ) -> Self: - obj = self._get_obj(inplace) - if isinstance(other, obj.__class__): - new_agents = other.agents - elif isinstance(other, pd.DataFrame): - new_agents = other - if "unique_id" != other.index.name: - try: - new_agents.set_index("unique_id", inplace=True, drop=True) - except KeyError: - new_agents["unique_id"] = obj.random.random(len(other)) * 10**8 - elif isinstance(other, dict): - if "unique_id" not in other: - index = obj.random.random(len(other)) * 10**8 - if not isinstance(other["unique_id"], ListLike): - index = [other["unique_id"]] - else: - index = other["unique_id"] - new_agents = ( - pd.DataFrame(other, index=pd.Index(index)) - .reset_index(drop=True) - .set_index("unique_id") - ) - else: # ListLike - if len(other) == len(obj._agents.columns): - # data missing unique_id - new_agents = pd.DataFrame([other], columns=obj._agents.columns) - new_agents["unique_id"] = obj.random.random(1) * 10**8 - elif len(other) == len(obj._agents.columns) + 1: - new_agents = pd.DataFrame( - [other], columns=["unique_id"] + obj._agents.columns.values.tolist() - ) - else: - raise ValueError( - "Length of data must match the number of columns in the AgentSet if being added as a ListLike." - ) - new_agents.set_index("unique_id", inplace=True, drop=True) - obj._agents = pd.concat([obj._agents, new_agents]) - return obj - - def remove(self, id: PandasMaskLike, inplace: bool = True) -> Self: - initial_len = len(self._agents) - obj = self._get_obj(inplace) - mask = obj._get_bool_mask(id) - remove_ids = obj._agents[mask].index - obj._agents.drop(remove_ids, inplace=True) - if len(obj._agents) == initial_len: - raise KeyError(f"IDs {id} not found in agent set.") - return obj - - -class AgentsPandas(AgentsDF): - agentsets: list[AgentSetPandas] - """A pandas implementation of a collection of AgentSetDF. All agents of the model are stored here.""" - - def __init__(self, model: ModelDF): - """Create a new AgentsDF object. - - Parameters - ---------- - model : ModelDF - The model to which the AgentsDF object belongs. - - Attributes - ---------- - agentsets : list[AgentSetDF] - The AgentSetDFs that make up the AgentsDF object. - model : ModelDF - The model to which the AgentSetDF belongs. - """ - super().__init__(model) - - @property - def active_agents(self) -> pd.DataFrame: - return pd.concat([agentset.active_agents for agentset in self.agentsets]) - - @property - def inactive_agents(self) -> pd.DataFrame: - return pd.concat([agentset.inactive_agents for agentset in self.agentsets]) - - def select( - self, - mask: pd.Series[bool] | pd.DataFrame | None = None, - filter_func: Callable[[AgentSetDF], pd.Series[bool]] | None = None, - n: int = 0, - ) -> Self: - n, r = int(n / len(self.agentsets)), n % len(self.agentsets) - new_agentsets: list[AgentSetPandas] = [] - for agentset in self.agentsets: - if mask is None: - agentset_mask = mask - elif isinstance(mask, pd.DataFrame): - agentset_mask = pd.Series( - agentset.agents.index.isin(mask), index=agentset.agents.index - ) - else: - agentset_mask = pd.Series( - agentset.agents.index.isin(mask[mask].index), - index=agentset.agents.index, - ) - agentset.select(mask=agentset_mask, filter_func=filter_func, n=n + r) - if len(agentset.active_agents) > n: - r = len(agentset.active_agents) - n - new_agentsets.append(agentset) - self.agentsets = new_agentsets - return self - - def get_attribute(self, attr_names: str) -> pd.Series[Any]: - return pd.concat( - [agentset.get_attribute(attr_names) for agentset in self.agentsets] - ) - - def add(self, agentsets: AgentSetPandas | list[AgentSetPandas]) -> Self: - return super().add(agentsets) # type: ignore diff --git a/mesa_frames/polars/agent.py b/mesa_frames/polars/agent.py deleted file mode 100644 index 4379f11..0000000 --- a/mesa_frames/polars/agent.py +++ /dev/null @@ -1,236 +0,0 @@ -import polars as pl - -from .base.agent import AgentSetDF -from .base.model import ModelDF - - -class AgentSetPolars(AgentSetDF): - _agents: pl.DataFrame - _mask: pl.Expr | pl.Series - """A polars-based implementation of the AgentSet.""" - - def __init__(self, model: ModelDF): - """Create a new AgentSetDF. - - Parameters - ---------- - model : ModelDF - The model to which the AgentSetDF belongs. - - Attributes - ---------- - agents : pl.DataFrame - The agents in the AgentSetDF. - model : ModelDF - The model to which the AgentSetDF belongs. - """ - super().__init__(model) - self._agents = pl.DataFrame(schema={"unique_id": pl.String}) - self._mask = pl.repeat(True, len(self.agents)) - - @property - def agents(self) -> pl.DataFrame: - if self._agents is None: - self._agents = pl.DataFrame(schema={"unique_id": pl.String}) - return self._agents - - @property - def active_agents(self) -> pl.DataFrame: - return self.agents.filter(self._mask) - - @property - def inactive_agents(self) -> pl.DataFrame: - return self.agents.filter(~self._mask) - - def select( - self, - mask: pl.Expr | pl.Series | pl.DataFrame | None = None, - filter_func: Callable[[Self], pl.Series] | None = None, - n: int = 0, - ) -> Self: - if mask is None: # if not mask doesn't work - mask = pl.repeat(True, len(self.agents)) - elif isinstance(mask, pl.DataFrame): - mask = self.agents["unique_id"].is_in(mask["unique_id"]) - if filter_func: - mask = mask & filter_func(self) - if n != 0: - mask = ( - self.agents.filter(mask) - .sample(n)["unique_id"] - .is_in(self.agents["unique_id"]) - ) - self._mask = mask - return self - - def shuffle(self) -> Self: - self.agents = self.agents.sample(fraction=1) - return self - - def sort( - self, - by: IntoExpr | Iterable[IntoExpr], - *more_by: IntoExpr, - descending: bool | Sequence[bool] = False, - nulls_last: bool = False, - ) -> Self: - """Sort the agents in the agent set based on the given criteria. - - Parameters - ---------- - by (IntoExpr | Iterable[IntoExpr]): The attribute(s) to sort by. - more_by (IntoExpr): Additional attributes to sort by. - descending (bool | Sequence[bool]): Whether to sort in descending order. - nulls_last (bool): Whether to place null values last. - - Returns - ---------- - AgentSetDF: The sorted agent set. - """ - self.agents = self.agents.sort( - by=by, *more_by, descending=descending, nulls_last=nulls_last - ) - return self - - def get_attribute(self, attr_names: str) -> pl.Series: - return self.agents.filter(self._mask)[attr_names] - - def set_attribute(self, attr_names: str, value: Any) -> Self: - if type(value) == pl.Series: - self.agents.filter(self._mask).with_columns(**{attr_names: value}) - else: - self.agents.filter(self._mask).with_columns(**{attr_names: pl.lit(value)}) - return self - - def add( - self, - n: int, - data: FrameInitTypes | None = None, - schema: SchemaDefinition | None = None, - schema_overrides: SchemaDict | None = None, - orient: Orientation | None = None, - infer_schema_length: int | None = N_INFER_DEFAULT, - nan_to_null: bool = False, - ) -> Self: - """Adds new agents to the agent set. - - Parameters - ---------- - n : int - The number of agents to add. - data : dict, Sequence, ndarray, Series, or pandas.DataFrame - Two-dimensional data in various forms; dict input must contain Sequences, Generators, or a range. Sequence may contain Series or other Sequences. - schema : Sequence of str, (str,DataType) pairs, or a {str:DataType,} dict - The DataFrame schema may be declared in several ways: - - As a dict of {name:type} pairs; if type is None, it will be auto-inferred. - - As a list of column names; in this case types are automatically inferred. - - As a list of (name,type) pairs; this is equivalent to the dictionary form. - If you supply a list of column names that does not match the names in the underlying data, the names given here will overwrite them. The number of names given in the schema should match the underlying data dimensions. - schema_overrides : dict, default None - Support type specification or override of one or more columns; note that any dtypes inferred from the schema param will be overridden. - The number of entries in the schema should match the underlying data dimensions, unless a sequence of dictionaries is being passed, in which case a *partial* schema can be declared to prevent specific fields from being loaded. - orient : {'col', 'row'}, default None - Whether to interpret two-dimensional data as columns or as rows. - If None, the orientation is inferred by matching the columns and data dimensions. - If this does not yield conclusive results, column orientation is used. - infer_schema_length : int or None - The maximum number of rows to scan for schema inference. If set to None, the full data may be scanned *(this is slow)*. - This parameter only applies if the input data is a sequence or generator of rows; other input is read as-is. - nan_to_null : bool, default False - If the data comes from one or more numpy arrays, can optionally convert input data np.nan values to null instead. This is a no-op for all other input data. - - Returns - ---------- - AgentSetPolars: The updated agent set. - """ - new_df = pl.DataFrame( - data=data, - schema=schema, - schema_overrides=schema_overrides, - orient=orient, - infer_schema_length=infer_schema_length, - nan_to_null=nan_to_null, - ) - - if "unique_id" not in new_df.columns: - new_df = new_df.with_columns( - unique_id=pl.Series( - values=self.random.random(n) * 10**8, dtype=pl.Int64 - ) - ) - - old_active_agents = self.agents.filter(self._mask)["unique_id"] - self.agents = pl.concat([self.agents, new_df]) - self._mask = self.agents["unique_id"].is_in(old_active_agents) | self.agents[ - "unique_id" - ].is_in(new_df["unique_id"]) - return self - - def discard(self, id: int) -> Self: - with suppress(KeyError): - self.agents = self.agents.filter(self.agents["unique_id"] != id) - return self - - def remove(self, id: int) -> Self: - self.agents = self.agents.filter(self.agents["unique_id"] != id) - return self - - -class AgentsPolars(AgentsDF): - agentsets: list[AgentSetPolars] - """A polars implementation of a collection of AgentSetDF. All agents of the model are stored here.""" - - def __init__(self, model: ModelDF): - """Create a new AgentsDF object. - - Parameters - ---------- - model : ModelDF - The model to which the AgentsDF object belongs. - - Attributes - ---------- - agentsets : list[AgentSetDF] - The AgentSetDFs that make up the AgentsDF object. - model : ModelDF - The model to which the AgentSetDF belongs. - """ - super().__init__(model) - - @property - def active_agents(self) -> pl.DataFrame: - return pl.concat([agentset.active_agents for agentset in self.agentsets]) - - @property - def inactive_agents(self) -> pl.DataFrame: - return pl.concat([agentset.inactive_agents for agentset in self.agentsets]) - - def select( - self, - mask: pl.Expr | pl.Series | pl.DataFrame | None = None, - filter_func: Callable[[AgentSetDF], pl.Series] | None = None, - n: int = 0, - ) -> Self: - n, r = int(n / len(self.agentsets)), n % len(self.agentsets) - new_agentsets: list[AgentSetPolars] = [] - for agentset in self.agentsets: - if mask is None: - agentset_mask = mask - elif isinstance(mask, pl.DataFrame): - agentset_mask = agentset.agents["unique_id"].is_in(mask["unique_id"]) - elif isinstance(mask, pl.Series): - agentset_mask = agentset.agents["unique_id"].is_in(mask) - agentset.select(mask=agentset_mask, filter_func=filter_func, n=n + r) - if len(agentset.active_agents) > n: - r = len(agentset.active_agents) - n - new_agentsets.append(agentset) - self.agentsets = new_agentsets - return self - - def get_attribute(self, attr_names: str) -> pl.Series: - return pl.concat( - [agentset.get_attribute(attr_names) for agentset in self.agentsets] - ) - - def add(self, agentsets: AgentSetPolars | list[AgentSetPolars]) -> Self: - return super().add(agentsets) # type: ignore #child classes are not checked? diff --git a/mesa_frames/types.py b/mesa_frames/types.py new file mode 100644 index 0000000..8ca8726 --- /dev/null +++ b/mesa_frames/types.py @@ -0,0 +1,27 @@ +from typing import Collection, Hashable, Literal + +####----- Agnostic Types -----#### +AgnosticMask = Literal["all", "active"] | Hashable | None + + +###----- Pandas Types -----### +import pandas as pd +from numpy import ndarray + +ArrayLike = pd.api.extensions.ExtensionArray | ndarray +AnyArrayLike = ArrayLike | pd.Index | pd.Series +PandasMaskLike = AgnosticMask | pd.Series | pd.DataFrame | AnyArrayLike + + +###----- Polars Types -----### +import polars as pl + +PolarsMaskLike = AgnosticMask | pl.Expr | pl.Series | pl.DataFrame | Collection[int] + + +###----- Generic -----### + +DataFrame = pd.DataFrame | pl.DataFrame +Series = pd.Series | pl.Series +BoolSeries = pd.Series | pl.Series | pl.Expr +MaskLike = AgnosticMask | PandasMaskLike | PolarsMaskLike diff --git a/tests/test_agent.py b/tests/test_agentset_pandas.py similarity index 60% rename from tests/test_agent.py rename to tests/test_agentset_pandas.py index b3375ae..cd15948 100644 --- a/tests/test_agent.py +++ b/tests/test_agentset_pandas.py @@ -3,45 +3,39 @@ import pandas as pd import pytest import typeguard as tg -from mesa import Model from numpy.random import Generator -from mesa_frames import ( - AgentSetPandas, - AgentSetPolars, - AgentsPandas, - AgentsPolars, - ModelDF, - agent, -) +from mesa_frames import AgentSetPandas, ModelDF @tg.typechecked class ExampleAgentSet(AgentSetPandas): - def __init__(self, model: ModelDF): - self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth") + def __init__(self, model: ModelDF, index: pd.Index): + super().__init__(model) + self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth", index=index) def add_wealth(self, amount: int) -> None: self.agents["wealth"] += amount @pytest.fixture -def fix1_AgentSetPandas(): +def fix1_AgentSetPandas() -> ExampleAgentSet: model = ModelDF() - agents = ExampleAgentSet(model) + agents = ExampleAgentSet(model, pd.Index([0, 1, 2, 3], name="unique_id")) agents.add({"unique_id": [0, 1, 2, 3]}) - agents.agents["wealth"] = agents.starting_wealth - agents.agents["age"] = [10, 20, 30, 40] + agents["wealth"] = agents.starting_wealth + agents["age"] = [10, 20, 30, 40] + return agents @pytest.fixture -def fix2_AgentSetPandas(): +def fix2_AgentSetPandas() -> ExampleAgentSet: model = ModelDF() - agents = ExampleAgentSet(model) - agents.add({"unique_id": [10, 11, 12, 13]}) - agents.agents["wealth"] = agents.starting_wealth + 10 - agents.agents["age"] = [100, 200, 300, 400] + agents = ExampleAgentSet(model, pd.Index([4, 5, 6, 7], name="unique_id")) + agents.add({"unique_id": [4, 5, 6, 7]}) + agents["wealth"] = agents.starting_wealth + 10 + agents["age"] = [100, 200, 300, 400] return agents @@ -49,211 +43,40 @@ def fix2_AgentSetPandas(): class Test_AgentSetPandas: def test__init__(self): model = ModelDF() - forbidden_model = Model() - try: - agents = ExampleAgentSet(forbidden_model) - except Exception as e: - assert type(e) == tg.TypeCheckError - agents = ExampleAgentSet(model) + agents = ExampleAgentSet(model, pd.Index([0, 1, 2, 3])) assert agents.model == model assert isinstance(agents.agents, pd.DataFrame) + assert agents.agents.index.name == "unique_id" assert isinstance(agents._mask, pd.Series) assert isinstance(agents.random, Generator) assert agents.starting_wealth.tolist() == [1, 2, 3, 4] - def test__add__(self, fix1_AgentSetPandas, fix2_AgentSetPandas): - agents = fix1_AgentSetPandas - agents2 = fix2_AgentSetPandas - - # Test with two AgentSetPandas - agents3 = agents + agents2 - assert agents3.agents.index.tolist() == [0, 1, 2, 3, 10, 11, 12, 13] - - # Test with an AgentSetPandas and a DataFrame - agents3 = agents + agents2.agents - assert agents3.agents.index.tolist() == [0, 1, 2, 3, 10, 11, 12, 13] - - # Test with an AgentSetPandas and a list - agents3 = agents + [10, 5] # 10 should be unique id and 5 should be wealth - assert agents3.agents.index.tolist()[:-1] == [0, 1, 2, 3] - assert len(agents3.agents) == 5 - assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 10] - - # Test with an AgentSetPandas and a dict - agents3 = agents + {"unique_id": 10, "wealth": 5} - assert agents3.agents.index.tolist() == [0, 1, 2, 3, 10] - assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 5] - - def test__contains__(self, fix1_AgentSetPandas): - # Test with a single value - agents = fix1_AgentSetPandas - assert 0 in agents - assert 4 not in agents - - def test__copy__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - agents2 = copy(agents) - agents[0, "wealth"] = 5 - assert agents is not agents2 - assert agents[0, "wealth"].values == agents2[0, "wealth"].values - assert agents2.model is agents.model - - def test__deepcopy__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - agents2 = deepcopy(agents) - agents[0, "wealth"] = 5 - assert agents is not agents2 - assert agents[0, "wealth"].values != agents2[0, "wealth"].values - assert agents2.model is agents.model - - def test__getattr__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - assert isinstance(agents.model, ModelDF) - assert agents.wealth.tolist() == [1, 2, 3, 4] - - def test__getitem__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - - # Testing with a string - assert agents["wealth"].tolist() == [1, 2, 3, 4] - - # Test with a tuple[MaskLike, str] - assert agents[0, "wealth"].values == 1 - - # Test with a list[str] - assert agents[["wealth", "age"]].columns.tolist() == ["wealth", "age"] - - # Testing with a tuple[MaskLike, list[str]] - result = agents[0, ["wealth", "age"]] - assert result["wealth"].values.tolist() == [1] - assert result["age"].values.tolist() == [10] - - def test__iadd__(self, fix1_AgentSetPandas, fix2_AgentSetPandas): - agents = deepcopy(fix1_AgentSetPandas) - agents2 = fix2_AgentSetPandas - - # Test with two AgentSetPandas - agents += agents2 - assert agents.agents.index.tolist() == [0, 1, 2, 3, 10, 11, 12, 13] - - # Test with an AgentSetPandas and a DataFrame - agents = deepcopy(fix1_AgentSetPandas) - agents += agents2.agents - assert agents.agents.index.tolist() == [0, 1, 2, 3, 10, 11, 12, 13] - - # Test with an AgentSetPandas and a list - agents = deepcopy(fix1_AgentSetPandas) - agents += [10, 5] - assert agents.agents.index.tolist()[:-1] == [0, 1, 2, 3] - assert len(agents.agents) == 5 - assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 10] - - # Test with an AgentSetPandas and a dict - agents = deepcopy(fix1_AgentSetPandas) - agents += {"unique_id": 10, "wealth": 5} - assert agents.agents.index.tolist() == [0, 1, 2, 3, 10] - - def test__iter__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - for i, agent in agents: - assert isinstance(agent, pd.Series) - assert agent["wealth"] == i + 1 - - def test__isub__(self, fix1_AgentSetPandas): - agents = deepcopy(fix1_AgentSetPandas) - - # Test with two AgentSetPandas - agents -= agents - assert agents.agents.empty - - # Test with an AgentSetPandas and a DataFrame - agents = deepcopy(fix1_AgentSetPandas) - agents -= agents.agents - assert agents.agents.empty - - def test__len__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - assert len(agents) == 4 - - def test__repr__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - repr(agents) - - def test__reversed__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - reversed_wealth = [] - for i, agent in reversed(agents): - reversed_wealth.append(agent["wealth"]) - assert reversed_wealth == [4, 3, 2, 1] - - def test__setitem__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - - agents = deepcopy(agents) # To test passing through a df later - - # Test with key=str, value=Any - agents["wealth"] = 0 - assert agents.agents.wealth.tolist() == [0, 0, 0, 0] - - # Test with key=list[str], value=Any - agents[["wealth", "age"]] = 1 - assert agents.agents.wealth.tolist() == [1, 1, 1, 1] - assert agents.agents.age.tolist() == [1, 1, 1, 1] - - # Test with key=tuple, value=Any - agents[0, "wealth"] = 5 - assert agents.agents.wealth.tolist() == [5, 1, 1, 1] - - # Test with key=MaskLike, value=Any - agents[0] = [9, 99] - assert agents.agents.loc[0, "wealth"] == 9 - assert agents.agents.loc[0, "age"] == 99 - - def test__str__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - str(agents) - - def test__sub__(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - agents2 = agents - agents - assert agents2.agents.empty - assert agents.agents.wealth.tolist() == [1, 2, 3, 4] - - def test_get_object(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - assert agents._get_obj(inplace=True) is agents - assert agents._get_obj(inplace=False) is not agents - - def test_active_agents(self, fix1_AgentSetPandas): - agents = fix1_AgentSetPandas - agents.active_agents = [0, 1] - assert agents.active_agents.index.to_list() == [0, 1] - - def test_add(self, fix1_AgentSetPandas, fix2_AgentSetPandas): + def test_add( + self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet + ): agents = fix1_AgentSetPandas agents2 = fix2_AgentSetPandas - # Test with self - result = agents.add(agents2, inplace=False) - assert result.agents.index.to_list() == [0, 1, 2, 3, 10, 11, 12, 13] - # Test with a DataFrame result = agents.add(agents2.agents, inplace=False) - assert result.agents.index.to_list() == [0, 1, 2, 3, 10, 11, 12, 13] + assert result.agents.index.to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + assert agents.agents.index.name == "unique_id" - # Test with a list + # Test with a list (Sequence[Any]) result = agents.add([10, 5, 10], inplace=False) assert result.agents.index.to_list() == [0, 1, 2, 3, 10] assert result.agents.wealth.to_list() == [1, 2, 3, 4, 5] assert result.agents.age.to_list() == [10, 20, 30, 40, 10] + assert agents.agents.index.name == "unique_id" - # Test with a dict + # Test with a dict[str, Any] agents.add({"unique_id": [4, 5], "wealth": [5, 6], "age": [50, 60]}) assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5, 6] assert agents.agents.index.tolist() == [0, 1, 2, 3, 4, 5] assert agents.agents.age.tolist() == [10, 20, 30, 40, 50, 60] + assert agents.agents.index.name == "unique_id" - def test_contains(self, fix1_AgentSetPandas): + def test_contains(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas # Test with a single value @@ -263,31 +86,22 @@ def test_contains(self, fix1_AgentSetPandas): # Test with a list assert agents.contains([0, 1]).values.tolist() == [True, True] - # Test with a pd.DataFrame - assert agents.contains(pd.DataFrame({"unique_id": [0, 4]})).values.tolist() == [True, False] - - def test_copy(self, fix1_AgentSetPandas): + def test_copy(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas + agents.test_list = [[1, 2, 3]] + # Since pandas have Copy-on-Write, we can't test the deep method on DFs # Test with deep=False agents2 = agents.copy(deep=False) - agents2[0, "wealth"] = 5 - assert agents[0, "wealth"].values == agents2[0, "wealth"].values - assert agents2.model is agents.model + agents2.test_list[0].append(4) + assert agents.test_list[0][-1] == agents2.test_list[0][-1] # Test with deep=True agents2 = fix1_AgentSetPandas.copy(deep=True) - agents2[0, "wealth"] = 3 - assert agents[0, "wealth"].values != agents2[0, "wealth"].values - assert agents2.model is agents.model - - # Test by skipping starting_wealth - agents2 = agents.copy(skip=["starting_wealth"]) - with pytest.raises(KeyError) as e: - agents2.starting_wealth - assert "starting_wealth" in str(e) + agents2.test_list[0].append(4) + assert agents.test_list[-1] != agents2.test_list[-1] - def test_discard(self, fix1_AgentSetPandas): + def test_discard(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas # Test with a single value @@ -296,7 +110,7 @@ def test_discard(self, fix1_AgentSetPandas): # Test with a list result = agents.discard([0, 1], inplace=False) - assert agents.agents.index.tolist() == [2, 3] + assert result.agents.index.tolist() == [2, 3] # Test with a pd.DataFrame result = agents.discard(pd.DataFrame({"unique_id": [0, 1]}), inplace=False) @@ -307,29 +121,29 @@ def test_discard(self, fix1_AgentSetPandas): result = agents.discard("active", inplace=False) assert result.agents.index.to_list() == [2, 3] - def test_do(self, fix1_AgentSetPandas): + def test_do(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas agents.do("add_wealth", 1) assert agents.agents.wealth.tolist() == [2, 3, 4, 5] assert agents.do("add_wealth", 1, return_results=True) == None - def test_get_attribute(self, fix1_AgentSetPandas): + def test_get(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas # Test with a single attribute - assert agents.get_attribute("wealth").tolist() == [1, 2, 3, 4] + assert agents.get("wealth").tolist() == [1, 2, 3, 4] # Test with a list of attributes - assert agents.get_attribute(["wealth", "age"]).columns.tolist() == [ - "wealth", - "age", - ] + result = agents.get(["wealth", "age"]) + assert isinstance(result, pd.DataFrame) + assert result.columns.tolist() == ["wealth", "age"] + assert (result.wealth == agents.agents.wealth).all() # Test with a single attribute and a mask selected = agents.select(agents["wealth"] > 1, inplace=False) - assert selected.get_attribute("wealth", mask="active").tolist() == [2, 3, 4] + assert selected.get("wealth", mask="active").tolist() == [2, 3, 4] - def test_remove(self, fix1_AgentSetPandas): + def test_remove(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas agents.remove([0, 1]) assert agents.agents.index.tolist() == [2, 3] @@ -337,8 +151,8 @@ def test_remove(self, fix1_AgentSetPandas): agents.remove([1]) assert "1" in str(e) - def test_select(self, fix1_AgentSetPandas): - agents: ExampleAgentSet = fix1_AgentSetPandas + def test_select(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas # Test with default arguments. Should select all agents selected = agents.select(inplace=False) @@ -375,27 +189,29 @@ def filter_func(agentset: AgentSetPandas) -> pd.Series: selected = agents.select(mask, filter_func=filter_func, n=1, inplace=False) assert any(el in selected.active_agents.index.tolist() for el in [2, 3]) - def test_set_attribute(self, fix1_AgentSetPandas): + def test_set(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas # Test with a single attribute - agents.set_attribute("wealth", 0) - assert agents.agents.wealth.tolist() == [0, 0, 0, 0] + result = agents.set("wealth", 0, inplace=False) + assert result.agents.wealth.tolist() == [0, 0, 0, 0] # Test with a list of attributes - agents.set_attribute(["wealth", "age"], 1) + result = agents.set(["wealth", "age"], 1, inplace=False) + assert result.agents.wealth.tolist() == [1, 1, 1, 1] + assert result.agents.age.tolist() == [1, 1, 1, 1] # Test with a single attribute and a mask selected = agents.select(agents["wealth"] > 1, inplace=False) - selected.set_attribute("wealth", 0, mask="active") + selected.set("wealth", 0, mask="active") assert selected.agents.wealth.tolist() == [1, 0, 0, 0] # Test with a dictionary - agents.set_attribute({"wealth": 10, "age": 20}) + agents.set({"wealth": 10, "age": 20}) assert agents.agents.wealth.tolist() == [10, 10, 10, 10] assert agents.agents.age.tolist() == [20, 20, 20, 20] - def test_shuffle(self, fix1_AgentSetPandas): + def test_shuffle(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas for _ in range(10): original_order = agents.agents.index.tolist() @@ -404,7 +220,194 @@ def test_shuffle(self, fix1_AgentSetPandas): return assert False - def test_sort(self, fix1_AgentSetPandas): + def test_sort(self, fix1_AgentSetPandas: ExampleAgentSet): agents = fix1_AgentSetPandas agents.sort("wealth", ascending=False) assert agents.agents.wealth.tolist() == [4, 3, 2, 1] + + def test__add__( + self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet + ): + agents = fix1_AgentSetPandas + agents2 = fix2_AgentSetPandas + + # Test with an AgentSetPandas and a DataFrame + agents3 = agents + agents2.agents + assert agents3.agents.index.tolist() == [0, 1, 2, 3, 4, 5, 6, 7] + + # Test with an AgentSetPandas and a list (Sequence[Any]) + agents3 = agents + [10, 5, 5] # unique_id, wealth, age + assert agents3.agents.index.tolist()[:-1] == [0, 1, 2, 3] + assert len(agents3.agents) == 5 + assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 5] + assert agents3.agents.age.tolist() == [10, 20, 30, 40, 5] + + # Test with an AgentSetPandas and a dict + agents3 = agents + {"unique_id": 10, "wealth": 5} + assert agents3.agents.index.tolist() == [0, 1, 2, 3, 10] + assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 5] + + def test__contains__(self, fix1_AgentSetPandas: ExampleAgentSet): + # Test with a single value + agents = fix1_AgentSetPandas + assert 0 in agents + assert 4 not in agents + + def test__copy__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + agents.test_list = [[1, 2, 3]] + + # Since pandas have Copy-on-Write, we can't test the deep method on DFs + # Test with deep=False + agents2 = copy(agents) + agents2.test_list[0].append(4) + assert agents.test_list[0][-1] == agents2.test_list[0][-1] + + def test__deepcopy__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + agents.test_list = [[1, 2, 3]] + + agents2 = deepcopy(agents) + agents2.test_list[0].append(4) + assert agents.test_list[-1] != agents2.test_list[-1] + + def test__getattr__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + assert isinstance(agents.model, ModelDF) + assert agents.wealth.tolist() == [1, 2, 3, 4] + + def test__getitem__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + + # Testing with a string + assert agents["wealth"].tolist() == [1, 2, 3, 4] + + # Test with a tuple[MaskLike, str] + assert agents[0, "wealth"].values == 1 + + # Test with a list[str] + assert agents[["wealth", "age"]].columns.tolist() == ["wealth", "age"] + + # Testing with a tuple[MaskLike, list[str]] + result = agents[0, ["wealth", "age"]] + assert result["wealth"].values.tolist() == [1] + assert result["age"].values.tolist() == [10] + + def test__iadd__( + self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet + ): + agents = deepcopy(fix1_AgentSetPandas) + agents2 = fix2_AgentSetPandas + + # Test with an AgentSetPandas and a DataFrame + agents = deepcopy(fix1_AgentSetPandas) + agents += agents2.agents + assert agents.agents.index.tolist() == [0, 1, 2, 3, 4, 5, 6, 7] + + # Test with an AgentSetPandas and a list + agents = deepcopy(fix1_AgentSetPandas) + agents += [10, 5, 5] # unique_id, wealth, age + assert agents.agents.index.tolist()[:-1] == [0, 1, 2, 3] + assert len(agents.agents) == 5 + assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5] + assert agents.agents.age.tolist() == [10, 20, 30, 40, 5] + + # Test with an AgentSetPandas and a dict + agents = deepcopy(fix1_AgentSetPandas) + agents += {"unique_id": 10, "wealth": 5} + assert agents.agents.index.tolist() == [0, 1, 2, 3, 10] + assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5] + + def test__iter__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + for i, agent in agents: + assert isinstance(agent, pd.Series) + assert agent["wealth"] == i + 1 + + def test__isub__(self, fix1_AgentSetPandas: ExampleAgentSet): + + # Test with an AgentSetPandas and a DataFrame + agents = deepcopy(fix1_AgentSetPandas) + agents -= agents.agents + assert agents.agents.empty + + def test__len__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + assert len(agents) == 4 + + def test__repr__(self, fix1_AgentSetPandas): + agents: ExampleAgentSet = fix1_AgentSetPandas + repr(agents) + + def test__reversed__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + reversed_wealth = [] + for i, agent in reversed(agents): + reversed_wealth.append(agent["wealth"]) + assert reversed_wealth == [4, 3, 2, 1] + + def test__setitem__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + + agents = deepcopy(agents) # To test passing through a df later + + # Test with key=str, value=Any + agents["wealth"] = 0 + assert agents.agents.wealth.tolist() == [0, 0, 0, 0] + + # Test with key=list[str], value=Any + agents[["wealth", "age"]] = 1 + assert agents.agents.wealth.tolist() == [1, 1, 1, 1] + assert agents.agents.age.tolist() == [1, 1, 1, 1] + + # Test with key=tuple, value=Any + agents[0, "wealth"] = 5 + assert agents.agents.wealth.tolist() == [5, 1, 1, 1] + + # Test with key=MaskLike, value=Any + agents[0] = [9, 99] + assert agents.agents.loc[0, "wealth"] == 9 + assert agents.agents.loc[0, "age"] == 99 + + def test__str__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents: ExampleAgentSet = fix1_AgentSetPandas + str(agents) + + def test__sub__(self, fix1_AgentSetPandas: ExampleAgentSet): + agents: ExampleAgentSet = fix1_AgentSetPandas + agents2: ExampleAgentSet = agents - agents.agents + assert agents2.agents.empty + assert agents.agents.wealth.tolist() == [1, 2, 3, 4] + + def test_get_obj(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + assert agents._get_obj(inplace=True) is agents + assert agents._get_obj(inplace=False) is not agents + + def test_agents( + self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet + ): + agents = fix1_AgentSetPandas + agents2 = fix2_AgentSetPandas + assert isinstance(agents.agents, pd.DataFrame) + + # Test agents.setter + agents.agents = agents2.agents + assert agents.agents.index.tolist() == [4, 5, 6, 7] + + def test_active_agents(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + + # Test with select + agents.select(agents["wealth"] > 2, inplace=True) + assert agents.active_agents.index.tolist() == [2, 3] + + # Test with active_agents.setter + agents.active_agents = agents.agents.wealth > 2 + assert agents.active_agents.index.to_list() == [2, 3] + + def test_inactive_agents(self, fix1_AgentSetPandas: ExampleAgentSet): + agents = fix1_AgentSetPandas + + agents.select(agents["wealth"] > 2, inplace=True) + assert agents.inactive_agents.index.to_list() == [0, 1] diff --git a/tests/test_agentset_polars.py b/tests/test_agentset_polars.py new file mode 100644 index 0000000..8ed073e --- /dev/null +++ b/tests/test_agentset_polars.py @@ -0,0 +1,408 @@ +from copy import copy, deepcopy + +import polars as pl +import pytest +import typeguard as tg +from numpy.random import Generator + +from mesa_frames import AgentSetPolars, ModelDF + + +@tg.typechecked +class ExampleAgentSet(AgentSetPolars): + def __init__(self, model: ModelDF): + super().__init__(model) + self.starting_wealth = pl.Series("wealth", [1, 2, 3, 4]) + + def add_wealth(self, amount: int) -> None: + self.agents += amount + + +@pytest.fixture +def fix1_AgentSetPolars() -> ExampleAgentSet: + model = ModelDF() + agents = ExampleAgentSet(model) + agents.add({"unique_id": [0, 1, 2, 3]}) + agents["wealth"] = agents.starting_wealth + agents["age"] = [10, 20, 30, 40] + return agents + + +@pytest.fixture +def fix2_AgentSetPolars() -> ExampleAgentSet: + model = ModelDF() + agents = ExampleAgentSet(model) + agents.add({"unique_id": [4, 5, 6, 7]}) + agents["wealth"] = agents.starting_wealth + 10 + agents["age"] = [100, 200, 300, 400] + return agents + + +class Test_AgentSetPolars: + def test__init__(self): + model = ModelDF() + agents = ExampleAgentSet(model) + agents.add({"unique_id": [0, 1, 2, 3]}) + assert agents.model == model + assert isinstance(agents.agents, pl.DataFrame) + assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3] + assert isinstance(agents._mask, pl.Expr) + assert isinstance(agents.random, Generator) + assert agents.starting_wealth.to_list() == [1, 2, 3, 4] + + def test_add( + self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet + ): + agents = fix1_AgentSetPolars + agents2 = fix2_AgentSetPolars + + # Test with a DataFrame + result = agents.add(agents2.agents, inplace=False) + assert result.agents["unique_id"].to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + + # Test with a list (Sequence[Any]) + result = agents.add([10, 5, 10], inplace=False) + assert result.agents["unique_id"].to_list() == [0, 1, 2, 3, 10] + assert result.agents["wealth"].to_list() == [1, 2, 3, 4, 5] + assert result.agents["age"].to_list() == [10, 20, 30, 40, 10] + + # Test with a dict[str, Any] + agents.add({"unique_id": [4, 5], "wealth": [5, 6], "age": [50, 60]}) + assert agents.agents["wealth"].to_list() == [1, 2, 3, 4, 5, 6] + assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 4, 5] + assert agents.agents["age"].to_list() == [10, 20, 30, 40, 50, 60] + + def test_contains(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Test with a single value + assert agents.contains(0) + assert not agents.contains(4) + + # Test with a list + assert agents.contains([0, 1]).to_list() == [True, True] + + def test_copy(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + agents.test_list = [[1, 2, 3]] + + # Test with deep=False + agents2 = agents.copy(deep=False) + agents2.test_list[0].append(4) + assert agents.test_list[0][-1] == agents2.test_list[0][-1] + + # Test with deep=True + agents2 = fix1_AgentSetPolars.copy(deep=True) + agents2.test_list[0].append(4) + assert agents.test_list[-1] != agents2.test_list[-1] + + def test_discard(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Test with a single value + result = agents.discard(0, inplace=False) + assert result.agents["unique_id"].to_list() == [1, 2, 3] + + # Test with a list + result = agents.discard([0, 1], inplace=False) + assert result.agents["unique_id"].to_list() == [2, 3] + + # Test with a pl.DataFrame + result = agents.discard(pl.DataFrame({"unique_id": [0, 1]}), inplace=False) + assert result.agents["unique_id"].to_list() == [2, 3] + + # Test with active_agents + agents.active_agents = [0, 1] + result = agents.discard("active", inplace=False) + assert result.agents["unique_id"].to_list() == [2, 3] + + def test_do(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + agents.do("add_wealth", 1) + assert agents.agents["wealth"].to_list() == [2, 3, 4, 5] + assert agents.do("add_wealth", 1, return_results=True) is None + + def test_get(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Test with a single attribute + assert agents.get("wealth").to_list() == [1, 2, 3, 4] + + # Test with a list of attributes + result = agents.get(["wealth", "age"]) + assert isinstance(result, pl.DataFrame) + assert result.columns == ["wealth", "age"] + assert result["wealth"].to_list() == agents.agents["wealth"].to_list() + + # Test with a single attribute and a mask + selected = agents.select(agents.agents["wealth"] > 1, inplace=False) + assert selected.get("wealth", mask="active").to_list() == [2, 3, 4] + + def test_remove(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + agents.remove([0, 1]) + assert agents.agents["unique_id"].to_list() == [2, 3] + with pytest.raises(KeyError) as e: + agents.remove([1]) + + def test_select(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Test with default arguments. Should select all agents + selected = agents.select(inplace=False) + assert ( + selected.active_agents["wealth"].to_list() + == agents.agents["wealth"].to_list() + ) + + # Test with a pl.Series[bool] + mask = pl.Series("mask", [True, False, True, True], dtype=pl.Boolean) + selected = agents.select(mask, inplace=False) + assert selected.active_agents["unique_id"].to_list() == [0, 2, 3] + + # Test with a ListLike + mask = [0, 2] + selected = agents.select(mask, inplace=False) + assert selected.active_agents["unique_id"].to_list() == [0, 2] + + # Test with a pl.DataFrame + mask = pl.DataFrame({"unique_id": [0, 1]}) + selected = agents.select(mask, inplace=False) + assert selected.active_agents["unique_id"].to_list() == [0, 1] + + # Test with filter_func + def filter_func(agentset: AgentSetPolars) -> pl.Series: + return agentset.agents["wealth"] > 1 + + selected = agents.select(filter_func=filter_func, inplace=False) + assert selected.active_agents["unique_id"].to_list() == [1, 2, 3] + + # Test with n + selected = agents.select(n=3, inplace=False) + assert len(selected.active_agents) == 3 + + # Test with n, filter_func and mask + mask = pl.Series("mask", [True, False, True, True], dtype=pl.Boolean) + selected = agents.select(mask, filter_func=filter_func, n=1, inplace=False) + assert any(el in selected.active_agents["unique_id"].to_list() for el in [2, 3]) + + def test_set(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Test with a single attribute + result = agents.set("wealth", 0, inplace=False) + assert result.agents["wealth"].to_list() == [0, 0, 0, 0] + + # Test with a list of attributes + result = agents.set(["wealth", "age"], 1, inplace=False) + assert result.agents["wealth"].to_list() == [1, 1, 1, 1] + assert result.agents["age"].to_list() == [1, 1, 1, 1] + + # Test with a single attribute and a mask + selected = agents.select(agents.agents["wealth"] > 1, inplace=False) + selected.set("wealth", 0, mask="active") + assert selected.agents["wealth"].to_list() == [1, 0, 0, 0] + + # Test with a dictionary + agents.set({"wealth": 10, "age": 20}) + assert agents.agents["wealth"].to_list() == [10, 10, 10, 10] + assert agents.agents["age"].to_list() == [20, 20, 20, 20] + + def test_shuffle(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + for _ in range(10): + original_order = agents.agents["unique_id"].to_list() + agents.shuffle() + if original_order != agents.agents["unique_id"].to_list(): + return + assert False + + def test_sort(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + agents.sort("wealth", ascending=False) + assert agents.agents["wealth"].to_list() == [4, 3, 2, 1] + + def test__add__( + self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet + ): + agents = fix1_AgentSetPolars + agents2 = fix2_AgentSetPolars + + # Test with an AgentSetPolars and a DataFrame + agents3 = agents + agents2.agents + assert agents3.agents["unique_id"].to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + + # Test with an AgentSetPolars and a list (Sequence[Any]) + agents3 = agents + [10, 5, 5] # unique_id, wealth, age + assert agents3.agents["unique_id"].to_list()[:-1] == [0, 1, 2, 3] + assert len(agents3.agents) == 5 + assert agents3.agents["wealth"].to_list() == [1, 2, 3, 4, 5] + assert agents3.agents["age"].to_list() == [10, 20, 30, 40, 5] + + # Test with an AgentSetPolars and a dict + agents3 = agents + {"unique_id": 10, "wealth": 5} + assert agents3.agents["unique_id"].to_list() == [0, 1, 2, 3, 10] + assert agents3.agents["wealth"].to_list() == [1, 2, 3, 4, 5] + + def test__contains__(self, fix1_AgentSetPolars: ExampleAgentSet): + # Test with a single value + agents = fix1_AgentSetPolars + assert 0 in agents + assert 4 not in agents + + def test__copy__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + agents.test_list = [[1, 2, 3]] + + # Test with deep=False + agents2 = copy(agents) + agents2.test_list[0].append(4) + assert agents.test_list[0][-1] == agents2.test_list[0][-1] + + def test__deepcopy__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + agents.test_list = [[1, 2, 3]] + + agents2 = deepcopy(agents) + agents2.test_list[0].append(4) + assert agents.test_list[-1] != agents2.test_list[-1] + + def test__getattr__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + assert isinstance(agents.model, ModelDF) + assert agents.wealth.to_list() == [1, 2, 3, 4] + + def test__getitem__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Testing with a string + assert agents["wealth"].to_list() == [1, 2, 3, 4] + + # Test with a tuple[MaskLike, str] + assert agents[0, "wealth"].item() == 1 + + # Test with a list[str] + assert agents[["wealth", "age"]].columns == ["wealth", "age"] + + # Testing with a tuple[MaskLike, list[str]] + result = agents[0, ["wealth", "age"]] + assert result["wealth"].to_list() == [1] + assert result["age"].to_list() == [10] + + def test__iadd__( + self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet + ): + agents = deepcopy(fix1_AgentSetPolars) + agents2 = fix2_AgentSetPolars + + # Test with an AgentSetPolars and a DataFrame + agents = deepcopy(fix1_AgentSetPolars) + agents += agents2.agents + assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 4, 5, 6, 7] + + # Test with an AgentSetPolars and a list + agents = deepcopy(fix1_AgentSetPolars) + agents += [10, 5, 5] # unique_id, wealth, age + assert agents.agents["unique_id"].to_list()[:-1] == [0, 1, 2, 3] + assert len(agents.agents) == 5 + assert agents.agents["wealth"].to_list() == [1, 2, 3, 4, 5] + assert agents.agents["age"].to_list() == [10, 20, 30, 40, 5] + + # Test with an AgentSetPolars and a dict + agents = deepcopy(fix1_AgentSetPolars) + agents += {"unique_id": 10, "wealth": 5} + assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 10] + assert agents.agents["wealth"].to_list() == [1, 2, 3, 4, 5] + + def test__iter__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + for i, agent in enumerate(agents): + assert isinstance(agent, dict) + assert agent["wealth"] == i + 1 + + def test__isub__(self, fix1_AgentSetPolars: ExampleAgentSet): + # Test with an AgentSetPolars and a DataFrame + agents = deepcopy(fix1_AgentSetPolars) + agents -= agents.agents + assert agents.agents.is_empty() + + def test__len__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + assert len(agents) == 4 + + def test__repr__(self, fix1_AgentSetPolars): + agents: ExampleAgentSet = fix1_AgentSetPolars + repr(agents) + + def test__reversed__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + reversed_wealth = [] + for i, agent in reversed(list(enumerate(agents))): + reversed_wealth.append(agent["wealth"]) + assert reversed_wealth == [4, 3, 2, 1] + + def test__setitem__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + agents = deepcopy(agents) # To test passing through a df later + + # Test with key=str, value=Any + agents["wealth"] = 0 + assert agents.agents["wealth"].to_list() == [0, 0, 0, 0] + + # Test with key=list[str], value=Any + agents[["wealth", "age"]] = 1 + assert agents.agents["wealth"].to_list() == [1, 1, 1, 1] + assert agents.agents["age"].to_list() == [1, 1, 1, 1] + + # Test with key=tuple, value=Any + agents[0, "wealth"] = 5 + assert agents.agents["wealth"].to_list() == [5, 1, 1, 1] + + # Test with key=MaskLike, value=Any + agents[0] = [9, 99] + assert agents.agents.item(0, "wealth") == 9 + assert agents.agents.item(0, "age") == 99 + + def test__str__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents: ExampleAgentSet = fix1_AgentSetPolars + str(agents) + + def test__sub__(self, fix1_AgentSetPolars: ExampleAgentSet): + agents: ExampleAgentSet = fix1_AgentSetPolars + agents2: ExampleAgentSet = agents - agents.agents + assert agents2.agents.is_empty() + assert agents.agents["wealth"].to_list() == [1, 2, 3, 4] + + def test_get_obj(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + assert agents._get_obj(inplace=True) is agents + assert agents._get_obj(inplace=False) is not agents + + def test_agents( + self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet + ): + agents = fix1_AgentSetPolars + agents2 = fix2_AgentSetPolars + assert isinstance(agents.agents, pl.DataFrame) + + # Test agents.setter + agents.agents = agents2.agents + assert agents.agents["unique_id"].to_list() == [4, 5, 6, 7] + + def test_active_agents(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + # Test with select + agents.select(agents.agents["wealth"] > 2, inplace=True) + assert agents.active_agents["unique_id"].to_list() == [2, 3] + + # Test with active_agents.setter + agents.active_agents = agents.agents["wealth"] > 2 + assert agents.active_agents["unique_id"].to_list() == [2, 3] + + def test_inactive_agents(self, fix1_AgentSetPolars: ExampleAgentSet): + agents = fix1_AgentSetPolars + + agents.select(agents.agents["wealth"] > 2, inplace=True) + assert agents.inactive_agents["unique_id"].to_list() == [0, 1]