From 12a2f636d41adc4246ec2e20f5525d29a79fe063 Mon Sep 17 00:00:00 2001 From: Adam Amer <136176500+adamamer20@users.noreply.github.com> Date: Sat, 3 Feb 2024 10:11:12 +0100 Subject: [PATCH] API implementation of AgentSet and Agents 1) Agents of the same type are stored in a pd.DataFrame of the AgentSetDF class. This avoids having many missing values (which would occupy memory) if Agents or different type were stored in the same Dataframe. All agents of the model are stored in the AgentsDF class as a list of the AgentSetDF, which virtually supports the same operations as AgentSet. 2) Went with encapsulation, avoiding extensions and subclassing because they didn't always work well with storing additional attributes and it wasn't easy to extend the created classes by subclassing (as it's often done with base mesa.Agent). 3) All operations are inplace (no immutability): if we wanted to keep immutability, there couldn't be method chaining with encapsulation (since methods would have to return a pd.Dataframe) If we were to return an AgentSet. I think it also aligns well with base mesa API. 4) The "select" operations was changed to selecting "active_agents" (by a latent mask) which are the ones effectively used for the "do" operations (since usually you don't want to remove the other agents from the df but simply use a method on a part of agents of the DF). If you wanted to remove the agents you could simply use the "discard" or "remove" methods. --- docs/scripts/readme_plot.py | 63 ++-- mesa_frames/__init__.py | 3 +- mesa_frames/agent.py | 723 +++++++++++++++++++++++++++++------- mesa_frames/model.py | 81 +++- 4 files changed, 694 insertions(+), 176 deletions(-) diff --git a/docs/scripts/readme_plot.py b/docs/scripts/readme_plot.py index 938b311..2cca434 100644 --- a/docs/scripts/readme_plot.py +++ b/docs/scripts/readme_plot.py @@ -1,9 +1,13 @@ +from hmac import new +from random import seed +from typing import TYPE_CHECKING + import mesa import numpy as np +import pandas as pd import perfplot -from mesa_frames.agent import AgentDF -from mesa_frames.model import ModelDF +from mesa_frames import AgentSetDF, ModelDF # Mesa implementation @@ -57,6 +61,14 @@ def run_model(self, n_steps) -> None: self.step() +"""def compute_gini(model): + agent_wealths = model.agents.get("wealth") + x = sorted(agent_wealths) + N = model.num_agents + B = sum(xi * (N - i) for i, xi in enumerate(x)) / (N * sum(x)) + return 1 + (1 / N) - 2 * B""" + + # Mesa Frames implementation def mesa_frames_implementation(n_agents: int) -> None: model = MoneyModelDF(n_agents) @@ -67,45 +79,48 @@ class MoneyModelDF(ModelDF): def __init__(self, N): super().__init__() self.num_agents = N - self.create_agents(N, {MoneyAgentDF: 1}) + self.agents = self.agents.add(MoneyAgentsDF(N, model=self)) + + def step(self): + self.agents = self.agents.do("step") - def step(self, merged_mro=True): - self.agents = self.agents.sample(frac=1) - self.update_agents_masks() - super().step(merged_mro) + def run_model(self, n): + for _ in range(n): + self.step() -class MoneyAgentDF(AgentDF): - dtypes: dict[str, str] = {"wealth": "int64"} +class MoneyAgentsDF(AgentSetDF): + def __init__(self, n: int, model: MoneyModelDF): + super().__init__(model=model) + self.add(n, data={"wealth": np.ones(n)}) - @classmethod - def __init__(cls): - super().__init__() - cls.model.agents.loc[cls.mask, "wealth"] = 1 + def step(self): + wealthy_agents = self.agents["wealth"] > 0 + self.select(wealthy_agents).do("give_money") - @classmethod - def step(cls): - wealthy_agents = cls.model.agents.loc[cls.mask, "wealth"] > 0 - if wealthy_agents.any(): - other_agents = cls.model.agents.index.isin( - cls.model.agents.sample(n=wealthy_agents.sum()).index - ) - cls.model.agents.loc[wealthy_agents, "wealth"] -= 1 - cls.model.agents.loc[other_agents, "wealth"] += 1 + 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 main(): + mesa_frames_implementation(100) out = perfplot.bench( setup=lambda n: n, kernels=[mesa_implementation, mesa_frames_implementation], labels=["mesa", "mesa-frames"], - n_range=[k for k in range(10, 10000, 100)], + n_range=[k for k in range(100, 1000, 100)], xlabel="Number of agents", equality_check=None, title="100 steps of the Boltzmann Wealth model", ) out.show() - out.save("docs/images/readme_plot.png") + # out.save("docs/images/readme_plot.png") if __name__ == "__main__": diff --git a/mesa_frames/__init__.py b/mesa_frames/__init__.py index e98fa69..a80f01a 100644 --- a/mesa_frames/__init__.py +++ b/mesa_frames/__init__.py @@ -1,3 +1,2 @@ -from .agent import AgentSetPandas, AgentDF -from .datacollection import DataCollectorDF +from .agent import AgentsDF, AgentSetDF from .model import ModelDF diff --git a/mesa_frames/agent.py b/mesa_frames/agent.py index 6144ea3..443a4ec 100644 --- a/mesa_frames/agent.py +++ b/mesa_frames/agent.py @@ -1,207 +1,393 @@ from __future__ import annotations from contextlib import suppress -from typing import TYPE_CHECKING, Callable, Generator, Sequence, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Hashable, + Iterable, + Literal, + Sequence, + overload, +) -import numpy as np import pandas as pd +from numpy import ndarray # 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.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 + from .model import ModelDF + ArrayLike = ExtensionArray | ndarray + AnyArrayLike = ArrayLike | pd.Index | pd.Series + ValueKeyFunc = Callable[[pd.Series], pd.Series | AnyArrayLike] | None -class AgentSetPandas(DataFrame): - """ - Attributes - ---------- - agent_type : Agent - The type of the Agent. - model : model - model: The ABM model instance to which this AgentSet belongs.""" + from pandas._typing import Axes, Dtype, ListLikeU - model: ModelDF - agent_type: type[Agent] + +class AgentsDF: + """A collection of AgentSetDFs. 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. + """ + self.agentsets: list[AgentSetDF] = [] + self.model: ModelDF = model @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() - ) - else: - return result + def active_agents(self) -> pd.DataFrame: + """The active agents in the AgentsDF (those that are used for the do, set_attribute, get_attribute operations). + + Returns + ------- + pd.DataFrame + """ + return pd.concat([agentset.active_agents for agentset in self.agentsets]) + + @property + def inactive_agents(self) -> pd.DataFrame: + """The inactive agents in the AgentsDF (those that are not used for the do, set_attribute, get_attribute operations). + + Returns + ------- + pd.DataFrame + """ + return pd.concat([agentset.inactive_agents for agentset in self.agentsets]) + + @property + def random(self) -> Generator | None: + """ + Provide access to the model's random number generator. + + Returns: + ---------- + np.Generator + """ + return self.agentsets[0].model.random def select( self, - filter_func: Callable[[AgentSetPandas], pd.Series[bool]] | None = None, + mask: pd.Series[bool] | pd.DataFrame | None = None, + filter_func: Callable[[AgentSetDF], pd.Series[bool]] | None = None, n: int = 0, - inplace: bool = False, - ) -> AgentSetPandas | None: + ) -> AgentsDF: """ - Select a subset of agents from the AgentSet based on a filter function and/or quantity limit. + Change active_agents to a subset of agents from the AgentSet based on a mask, filter function and/or quantity limit. Attributes: ---------- - filter_func : Callable[[AgentSetPandas], pd.Series[bool]], optional + mask : pd.Series[bool] | pd.DataFrame | None, optional + A boolean mask indicating which agents should be included in the result. + If it's a DataFrame, it uses the indexes present in that dataframe. + If None, no filtering is applied. Defaults to None. + filter_func : Callable[[AgentSetDF], 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. + AgentsDF + The same AgentsDF with the updated active_agents property for each AgentSetDF. """ - 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]) + n, r = int(n / len(self.agentsets)), n % len(self.agentsets) + new_agentsets: list[AgentSetDF] = [] + 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 shuffle(self) -> AgentsDF: + """Randomly shuffles the agents in each AgentSetDF. - 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)) + Returns: + ---------- + AgentsDF + The same AgentsDF with the agents shuffled in each AgentSetDF. + """ + self.agentsets = [agentset.shuffle() for agentset in self.agentsets] + return self def sort( self, by: str | Sequence[str], key: ValueKeyFunc | None, ascending: bool | Sequence[bool] = True, - inplace: bool = False, - ) -> AgentSetPandas | None: + ) -> AgentsDF: + """ + Sort the agents in each AgentSetDF based on a specified attribute or custom function. + + Parameters: + ---------- + by : str | Sequence[str]) + A single attribute name or a list of attribute names based on which the agents are sorted. + key : ValueKeyFunc | None + 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 True. + + Returns: + ---------- + AgentsDF + The same AgentsDF with the agents sorted in each AgentSetDF. + """ + 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, + ) -> AgentsDF: ... + + @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, + ) -> AgentsDF | list[Any]: + """Invoke a method on each AgentSetDF. + + + Parameters + ---------- + method_name : str + The name of the method to call on each agent. + return_results : bool, optional + If True, the results of the method calls are returned. Defaults to False. + *args + Variable length argument list passed to the method being called. + **kwargs + Arbitrary keyword arguments passed to the method being called. + + Returns + ------- + AgentsDF | list[Any] + The same AgentsDF with each AgentSetDF updated based on the method call or the results of the method calls. + + """ + 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 get_attribute(self, attr_name: str) -> pd.Series[Any]: """ - Sort the agents in the AgentSetPandas based on a specified attribute or custom function. + Retrieve a specified attribute for active agents in AgentsDF. - 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. + Parameters: + ---------- + attr_name : str + The name of the attribute to retrieve. Returns: - AgentSetPandas: A sorted AgentSetPandas. Returns the current AgentSetPandas if inplace is True. + ---------- + pd.Series[Any] + A list of attribute values from each active agent in AgentsDF. """ - return cast( - "AgentSetPandas", - self.sort_values(by=by, key=key, ascending=ascending, inplace=inplace), + return pd.concat( + [agentset.get_attribute(attr_name) for agentset in self.agentsets] ) - def do(self, method_name: str, *args, sequential=False, **kwargs) -> AgentSetPandas: + def set_attribute(self, attr_name: str, value: Any) -> AgentsDF: """ - Invoke a method on each agent in the AgentSet. + Set a specified attribute for active agents in AgentsDF. 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. + attr_name : str + The name of the attribute to set for each agent. + value : Any + The value assigned to the attribute. If the value is a scalar, it is assigned to all active agents. + If the value is array-like, it must be the same length as the number of active agents. Returns: ---------- - AgentSetPandas: The results of the method calls + AgentsDF + The updated Agents """ - method = getattr(self, method_name) - if sequential: - return self.apply(method, axis=0, args=args, **kwargs) + self.agentsets = [ + agentset.set_attribute(attr_name, value) for agentset in self.agentsets + ] + return self + + def add(self, agentsets: AgentSetDF | list[AgentSetDF]) -> AgentsDF: + """Add an AgentSetDF or a list of AgentSetDFs to the AgentsDF. + + Parameters + ---------- + agentsets : AgentSetDF | list[AgentSetDF] + + Returns + ------- + AgentsDF + The updated AgentsDF. + """ + if isinstance(agentsets, AgentSetDF): + self.agentsets.append(agentsets) else: - return self.apply(method, axis=1, args=args, **kwargs) + self.agentsets = self.agentsets + agentsets + return self + + def discard(self, id: int) -> AgentsDF: + """Remove a specified agent. If the agent is not found, does not raise an error. + + Parameters + ---------- + id : int + The ID of the agent to remove. + + Returns + ---------- + Agents + The updated Agents.""" + self.agentsets = [agentset.discard(id) for agentset in self.agentsets] + return self - def get_attribute(self, attr_name: str) -> AgentSetPandas: + def remove(self, id: int) -> AgentsDF: + """Remove an agent from the AgentsDF. If the agent is not found, raises a KeyError. + + Parameters + ---------- + id : int + The ID of the agent to remove. + + Returns + ---------- + AgentsDF + The updated AgentsDF. + """ + 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.") + + def to_frame(self) -> pd.DataFrame: + """Convert the AgentsDF to a single DataFrame. + + Returns + ------- + pd.DataFrame + A DataFrame containing all agents from all AgentSetDFs. """ - Retrieve a specified attribute from each agent in the AgentSet. + return pd.concat([agentset.agents for agentset in self.agentsets]) - Args: - attr_name (str): The name of the attribute to retrieve from each agent. + def get_agents_of_type(self, agent_type: type) -> AgentSetDF: + """Retrieve the AgentSetDF of a specified type. - Returns: - list[Any]: A list of attribute values from each agent in the set. + Parameters + ---------- + agent_type : type + The type of AgentSetDF to retrieve. + + Returns + ------- + AgentSetDF + The AgentSetDF of the specified type. """ - return self[attr_name] + for agentset in self.agentsets: + if isinstance(agentset, agent_type): + return agentset + raise ValueError(f"No AgentSetDF of type {agent_type} found.") - def add_agents(self, n: int): - """Add n agents to the AgentSet. + +class AgentSetDF: + """A DataFrame-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 ---------- - n : int - The number of agents to add. + agents : pd.DataFrame + The agents in the AgentSetDF. + model : ModelDF + The model to which the AgentSetDF belongs. """ - # 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) + self.agents: pd.DataFrame = pd.DataFrame() + self.model: ModelDF = model + self._mask: pd.Series[bool] = pd.Series(True, index=self.agents.index) - def remove(self, agent: Agent): - """Remove an agent from the agentset.""" - self.drop(agent.unique_id, inplace=True) + @property + def active_agents(self) -> pd.DataFrame: + """The active agents in the AgentSetDF (those that are used for the do, set_attribute, get_attribute operations). + + Returns + ------- + pd.DataFrame + """ + return self.agents.loc[self._mask] + + @property + def inactive_agents(self) -> pd.DataFrame: + """The inactive agents in the AgentSetDF (those that are not used for the do, set_attribute, get_attribute operations). + + Returns + ------- + pd.DataFrame + """ + return self.agents.loc[~self._mask] @property def random(self) -> Generator: @@ -209,15 +395,264 @@ def random(self) -> Generator: Provide access to the model's random number generator. Returns: - Random: The random number generator associated with the model. + ---------- + np.Generator """ return self.model.random + def select( + self, + mask: pd.Series[bool] | pd.DataFrame | None = None, + filter_func: Callable[[AgentSetDF], pd.Series[bool]] | None = None, + n: int = 0, + ) -> AgentSetDF: + """ + Change active_agents to a subset of agents from the AgentSetDF based on a mask, filter function and/or quantity limit. -class AgentDF(Agent): - unique_id: int + Attributes: + ---------- + mask : pd.Series[bool] | pd.DataFrame | None, optional + A boolean mask indicating which agents should be included in the result. + If it's a DataFrame, it uses the indexes present in that dataframe. + If None, no filtering is applied. Defaults to None. + filter_func : Callable[[AgentSetDF], 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. + + Returns: + ---------- + AgentSetDF + The same AgentSetDF with the updated active_agents property. + """ + if mask is None: + mask = pd.Series(True, index=self.agents.index) + elif isinstance(mask, pd.DataFrame): + mask = pd.Series( + self.agents.index.isin(mask.index), index=self.agents.index + ) + if filter_func: + mask = mask & filter_func(self) + if n != 0: + mask = pd.Series(self.agents[mask].sample(n).index.isin(self.agents.index)) + self._mask = mask + return self + + def shuffle(self) -> AgentSetDF: + """Randomly shuffles the agents in the AgentSetDF. + + Returns: + ---------- + AgentSetDF + The same AgentSetDF with the agents shuffled. + """ + self.agents = self.agents.sample(frac=1) + return self + + def sort( + self, + by: str | Sequence[str], + key: ValueKeyFunc | None, + ascending: bool | Sequence[bool] = True, + ) -> AgentSetDF: + """ + Sort the agents in the AgentSetDF based on a specified attribute or custom function. + + Parameters: + ---------- + by : str | Sequence[str]) + A single attribute name or a list of attribute names based on which the agents are sorted. + key : ValueKeyFunc | None + 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 True. + + Returns: + ---------- + AgentSetDF + The same AgentSetDF with the agents sorted. + """ + self.agents.sort_values(by=by, key=key, ascending=ascending, inplace=True) + return self + + @overload + def do( + self, + method_name: str, + return_results: Literal[False] = False, + *args, + **kwargs, + ) -> AgentSetDF: ... + + @overload + def do( + self, + method_name: str, + return_results: Literal[True], + *args, + **kwargs, + ) -> Any: ... + + def do( + self, + method_name: str, + return_results: bool = False, + *args, + **kwargs, + ) -> AgentSetDF | Any: + """ + Invoke a method on the AgentSetDF. + + Parameters: + ---------- + method_name : str + The name of the method to call on each agent. + return_results : bool, optional + If True, the results of the method calls are returned. Defaults to False. + *args + Variable length argument list passed to the method being called. + **kwargs + Arbitrary keyword arguments passed to the method being called. + + Returns: + ---------- + AgentSetDF | Any + The same AgentSetDF with the agents updated based on the method call or the results of the method calls. + """ + method = getattr(self, method_name) + if return_results: + return method(*args, **kwargs) + else: + method(*args, **kwargs) + return self - def __init__(self, unique_id: int, model: ModelDF): - self.unique_id = unique_id - self.model = model - self.pos: Position | None = None + def get_attribute(self, attr_name: str) -> pd.Series[Any]: + """ + Retrieve a specified attribute for active agents in the AgentSetDF. + + Parameters: + ---------- + attr_name : str + The name of the attribute to retrieve. + + Returns: + ---------- + pd.Series[Any] + A list of attribute values from each active agent in AgentSetDF. + """ + return self.agents.loc[ + self.agents.index.isin(self.active_agents.index), attr_name + ] + + def set_attribute(self, attr_name: str, value: Any) -> AgentSetDF: + """ + Set a specified attribute for active agents in the AgentSetDF. + + Parameters: + ---------- + attr_name : str + The name of the attribute to set for each agent. + value : Any + The value assigned to the attribute. If the value is a scalar, it is assigned to all active agents. + If the value is array-like, it must be the same length as the number of active agents. + + Returns: + ---------- + AgentSetDF + The updated AgentSetDF. + """ + self.agents.loc[self.agents.index.isin(self.active_agents.index), attr_name] = ( + value + ) + return self + + def add( + self, + n: int, + data: ( + ListLikeU + | pd.DataFrame + | dict[Any, Any] + | Iterable[ListLikeU | tuple[Hashable, ListLikeU] | dict[Any, Any]] + | None + ) = None, + index: Axes | None = None, + copy: bool = False, + columns: Axes | None = None, + dtype: Dtype | None = None, + ) -> AgentSetDF: + """Add n agents to the AgentSet. + If index is not specified, the agents are assigned 64-bit random unique IDs using the model's generator. + The other arguments are passed to the pandas.DataFrame constructor. + + Attributes + ---------- + n : int + The number of agents to add. + data : array-like, Iterable, dict, or DataFrame, optional + Dict can contain Series, arrays, constants, dataclass or list-like objects. If data is a dict, column order follows insertion-order. If a dict contains Series which have an index defined, it is aligned by its index. This alignment also occurs if data is a Series or a DataFrame itself. Alignment is done on Series/DataFrame inputs. + If data is a list of dicts, column order follows insertion-order. + index : Index or array-like + Index to use for resulting frame. Will default to RangeIndex if no indexing information part of input data and no index provided. + columns : Index or array-like + Column labels to use for resulting frame when data does not have them, defaulting to RangeIndex(0, 1, 2, …, n). If data contains column labels, will perform column selection instead. + dtype : dtype, default None + Data type to force. Only a single dtype is allowed. If None, infer. + + Returns + ---------- + AgentSetDF + The updated self + """ + if not index: + index = pd.Index((self.random.random(n) * 10**8).astype(int)) + + new_df = pd.DataFrame( + data=data, + index=index, + columns=columns, + dtype=dtype, + copy=copy, + ) + + self.agents = pd.concat([self.agents, new_df]) + + if self._mask.empty: + self._mask = pd.Series(True, index=new_df.index) + else: + self._mask = pd.concat([self._mask, pd.Series(True, index=new_df.index)]) + + return self + + def discard(self, id: int) -> AgentSetDF: + """Remove a specified agent. If the agent is not found, does not raise an error. + + Parameters + ---------- + id : int + The ID of the agent to remove. + + Returns + ---------- + AgentSetDF + The updated AgentSetDF.""" + with suppress(KeyError): + self.agents.drop(id, inplace=True) + return self + + def remove(self, id: int) -> AgentSetDF: + """Remove an agent from the AgentSetDF. If the agent is not found, raises a KeyError. + + Parameters + ---------- + id : int + The ID of the agent to remove. + + Returns + ---------- + AgentSetDF + The updated AgentSetDF. + """ + self.agents.drop(id, inplace=True) + return self diff --git a/mesa_frames/model.py b/mesa_frames/model.py index a5a5ffe..8940fd3 100644 --- a/mesa_frames/model.py +++ b/mesa_frames/model.py @@ -1,14 +1,83 @@ -from copy import deepcopy -from time import time -from warnings import warn +from typing import TYPE_CHECKING, Any import geopandas as gpd import numpy as np import pandas as pd -from mesa_frames.agent import AgentDF +from .agent import AgentsDF, AgentSetDF +class ModelDF: + random: np.random.Generator + _seed: int + running: bool + + def __new__(cls, *args: Any, **kwargs: Any) -> Any: + """Create a new model object and instantiate its RNG automatically.""" + obj = object.__new__(cls) + if "seed" in kwargs: + obj._seed = kwargs["seed"] + else: + obj._seed = np.random.SeedSequence().entropy # type: ignore + 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 = AgentsDF([]) + self.agents: AgentsDF = AgentsDF(self) + + 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 @@ -39,7 +108,7 @@ def __new__(cls, *args, **kwargs): # advance. obj._seed = np.random.SeedSequence().entropy # Use default_rng to get a new Generator instance - obj.random = np.random.default_rng(obj._seed) + obj.random = np.random.default_rng(seed = obj._seed) return obj def __init__(self, unique_id: int | None = None, space=None): @@ -312,4 +381,4 @@ def update_agents_masks(self) -> None: tables=tables, ) # Collect data for the first time during initialization. - self.datacollector.collect(self)""" + self.datacollector.collect(self)'''