diff --git a/mesa_frames/__init__.py b/mesa_frames/__init__.py index e69de29..e98fa69 100644 --- a/mesa_frames/__init__.py +++ b/mesa_frames/__init__.py @@ -0,0 +1,3 @@ +from .agent import AgentSetPandas, AgentDF +from .datacollection import DataCollectorDF +from .model import ModelDF diff --git a/mesa_frames/agent.py b/mesa_frames/agent.py index 178acd5..6144ea3 100644 --- a/mesa_frames/agent.py +++ b/mesa_frames/agent.py @@ -1,80 +1,223 @@ -from typing import TYPE_CHECKING, Optional +from __future__ import annotations + +from contextlib import suppress +from typing import TYPE_CHECKING, Callable, Generator, Sequence, cast import numpy as np import pandas as pd -from numpy.random import randint + +# import polars as pl +from mesa import Agent +from model import ModelDF +from numpy.random import Generator +from pandas import DataFrame if TYPE_CHECKING: - from mesa_frames.model import ModelDF + from mesa.model import Model + from mesa.space import Position + from pandas import Index, Series + from pandas.core.arrays.base import ExtensionArray + ArrayLike = ExtensionArray | np.ndarray + AnyArrayLike = ArrayLike | Index | Series + ValueKeyFunc = Callable[[Series], Series | AnyArrayLike] | None + from mesa.space import Position -class AgentDF: - """The AgentDF class is the base class for other agents. - It should be used as inherited class for new agents classes. +class AgentSetPandas(DataFrame): + """ Attributes ---------- - dtypes : dict[str, str] - The attributes of the Agent as a dictionary of columns and data types. It contains: - - id : int64 - -- The unique id of the Agent. - - type : str - -- The type of the Agent. - model : Optional['ModelDF'] - The model of the simulation where the Agent is used. See src/model.py. Default: None - mask : pd.Series | None - The mask of the agents dataframe in the model which corresponds to the Agent class. - Initialized when model is created. Default: None - """ + agent_type : Agent + The type of the Agent. + model : model + model: The ABM model instance to which this AgentSet belongs.""" - dtypes: dict[str, str] = { - "id": "int64", - "type": "str", - } - model: Optional["ModelDF"] = None - mask: pd.Series | None = None - - @classmethod - def __init__(cls): - """Initializes the Agent class. - Assigns a 64-bit random id to each agent in the model. - """ - if cls.mask is None or cls.model is None or cls.model.agents is None: - raise ValueError( - "The Agent classes have not been initialized. Please use the model.create_agents() method to initialize the mask." + model: ModelDF + agent_type: type[Agent] + + @property + def _constructor(self): + return AgentSetPandas + + def __new__(cls, n: int, agent_type: type[Agent], model: ModelDF, *args, **kwargs): + return super().__new__(cls, *args, **kwargs) + + def __init__( + self, n: int, agent_type: type[Agent], model: ModelDF, *args, **kwargs + ): + super(DataFrame, self).__init__(*args, **kwargs) + self.model = model + self.agent_type = agent_type + self.add_agents(n) + + def __getitem__(self, key) -> AgentSetPandas: + result = super().__getitem__(key) + if isinstance(result, DataFrame): + # Create AgentSetPandas with DataFrame-specific data + return AgentSetPandas(0, self.agent_type, self.model, data=result) + elif isinstance(result, pd.Series): + # Convert Series to DataFrame and then create AgentSetPandas + return AgentSetPandas( + 0, self.agent_type, self.model, data=result.to_frame() ) - cls.model.agents.loc[cls.mask, "id"] = np.random.randint( - low=-9223372036854775808, - high=9223372036854775807, - size=cls.mask.sum(), - dtype="int64", + else: + return result + + def select( + self, + filter_func: Callable[[AgentSetPandas], pd.Series[bool]] | None = None, + n: int = 0, + inplace: bool = False, + ) -> AgentSetPandas | None: + """ + Select a subset of agents from the AgentSet based on a filter function and/or quantity limit. + + Attributes: + ---------- + filter_func : Callable[[AgentSetPandas], pd.Series[bool]], optional + A function that takes the AgentSet and returns a boolean mask over the agents indicating which agents + should be included in the result. Defaults to None, meaning no filtering is applied. + n : int, optional + The number of agents to select. If 0, all matching agents are selected. Defaults to 0. + inplace : bool, optional + If True, modifies the current AgentSet; otherwise, returns a new AgentSet. Defaults to False. + + Returns: + ---------- + AgentSet: A new AgentSet containing the selected agents, unless inplace is True, in which case the current AgentSet is updated. + """ + mask = pd.Series(True, index=self.index) + if filter_func: + mask = filter_func(self) + mask = mask & self.sample(n).index.isin(mask.index) + if inplace: + # Apply the mask in-place + self.loc[:, :] = self[mask] + else: + # Return a new instance + return AgentSetPandas(0, self.agent_type, self.model, self[mask]) + + def shuffle(self, inplace: bool = False) -> AgentSetPandas | None: + """Randomly shuffle the agents in the AgentSet.""" + if inplace: + self.loc[:, :] = self.sample(frac=1) + else: + return AgentSetPandas(0, self.agent_type, self.model, self.sample(frac=1)) + + def sort( + self, + by: str | Sequence[str], + key: ValueKeyFunc | None, + ascending: bool | Sequence[bool] = True, + inplace: bool = False, + ) -> AgentSetPandas | None: + """ + Sort the agents in the AgentSetPandas based on a specified attribute or custom function. + + Args: + key (Callable[[Agent], Any] | str): A function or attribute name based on which the agents are sorted. + ascending (bool, optional): If True, the agents are sorted in ascending order. Defaults to False. + inplace (bool, optional): If True, sorts the agents in the current AgentSetPandas; otherwise, returns a new sorted AgentSet. Defaults to False. + + Returns: + AgentSetPandas: A sorted AgentSetPandas. Returns the current AgentSetPandas if inplace is True. + """ + return cast( + "AgentSetPandas", + self.sort_values(by=by, key=key, ascending=ascending, inplace=inplace), ) - @classmethod - def step(cls): - """The step method of the Agent class. - It should be decorated as @classmethod and should act on the mask portion of the agents dataframe of the model. + def do(self, method_name: str, *args, sequential=False, **kwargs) -> AgentSetPandas: """ - pass + Invoke a method on each agent in the AgentSet. + Parameters: + ---------- + method_name (str): The name of the method to call on each agent. + *args: Variable length argument list passed to the method being called. + sequential = False + **kwargs: Arbitrary keyword arguments passed to the method being called. -class GeoAgentDF(AgentDF): - """The GeoAgentDF extends the AgentDF class to include a geometry attribute. - The agents will be stored in a GeoDataFrame. + Returns: + ---------- + AgentSetPandas: The results of the method calls + """ + method = getattr(self, method_name) + if sequential: + return self.apply(method, axis=0, args=args, **kwargs) + else: + return self.apply(method, axis=1, args=args, **kwargs) - Attributes - ---------- - dtypes : dict[str, str] - The attributes of the Agent as a dictionary of columns and data types. It contains: - - geometry : geometry - -- The geometry of the Agent. - """ + def get_attribute(self, attr_name: str) -> AgentSetPandas: + """ + Retrieve a specified attribute from each agent in the AgentSet. + + Args: + attr_name (str): The name of the attribute to retrieve from each agent. + + Returns: + list[Any]: A list of attribute values from each agent in the set. + """ + return self[attr_name] + + def add_agents(self, n: int): + """Add n agents to the AgentSet. + + Attributes + ---------- + n : int + The number of agents to add. + """ + # First, let's collect attributes from each agent_type. + callables = [] + values = [] + attributes = [] + for agent_type in reversed(self.agent_type.__mro__): + for attribute in agent_type.__dict__.keys(): + if attribute[:2] != "__": + attributes.append(attribute) + value = getattr(agent_type, attribute) + if callable(value): + callables.append((attribute, value)) + else: + values.append((attribute, value)) + # Now, let's create the agents. + self.index = pd.Index(self.model.random.random(n) % 1) + self.columns = list(attributes) + + # Finally, let's assign the values to the attribtutes. + + for attribute, value in values: + self[attribute] = value + + for attribute, value in callables: + self[attribute] = value(self) + + def discard(self, agent: Agent) -> AgentSetPandas | None: + """Remove an agent from the agentset.""" + with suppress(KeyError): + self.drop(agent.unique_id, inplace=True) - dtypes: dict[str, str] = {"geometry": "geometry"} + def remove(self, agent: Agent): + """Remove an agent from the agentset.""" + self.drop(agent.unique_id, inplace=True) - @classmethod - def step(cls): - """The step method of the GeoAgentDF class. - It should act on the mask portion of the agents dataframe of the model. + @property + def random(self) -> Generator: """ - pass + Provide access to the model's random number generator. + + Returns: + Random: The random number generator associated with the model. + """ + return self.model.random + + +class AgentDF(Agent): + unique_id: int + + def __init__(self, unique_id: int, model: ModelDF): + self.unique_id = unique_id + self.model = model + self.pos: Position | None = None