From 158a121ea862eeac00a06bcabd7d4942265a592c Mon Sep 17 00:00:00 2001 From: Adam Amer <136176500+adamamer20@users.noreply.github.com> Date: Sat, 24 Aug 2024 20:35:13 +0200 Subject: [PATCH] Enforce correct numpy docstring formatting with ruff.pydocstyle (#74) * adding pydocstyle to pyproject.toml * adding mesa_frames __init__ docstring * adding mesa_frames/abstract __init__ docstring * fixing mesa_frames/abstract/agents.py for ruff.pydocstyle * fixing mesa_frames/abstract/mixin.py for ruff.pydocstyle * fixing mesa_frames/abstract/space.py for ruff.pydocstyle * adding __init__ to mesa_frames/concrete * fixing arguments in __init__ of the module * fixing mesa_frames/concrete/model.py for ruff.pydocstyle * fixing mesa_frames/concrete/agents.py for ruff.pydocstyle * fixing mesa_frames/concrete/pandas/agentset.py for ruff.pydocstyle * adding __init__ to mesa_frames/concrete/pandas * fixing mesa_frames/concrete/pandas/mixin.py for ruff.pydocstyle * removing methods section (hard to mantain, and automatically generated by sphinx) * removing properties (sphinx will generate that) * moving parameters to __init__ * fixing mesa_frames/concrete/pandas/space.py * adding __init__ to mesa_frames/concrete/polars * removing docstring (sphinx will autogenerate from AgentSetDF) * removing docstring (sphinx will autogenerate from GridDF) * removing docstring (sphinx will autogenerate from DFMixin) * deactivate D101 and D102 (too many errors due to inheritance) * deactivating d105 (doesn't work well with inheritance) * fixing mesa_frames/concrete/polars/agentset.py for ruff.pydocstyle * fixing mesa_frames/concrete/polars/mixin.py for ruff.pydocstyle * fixing mesa_frames/concrete/polars/space.py * ignoring in tests, examples or docs * test -> tests * adding __init__ to type aliases * removing commented test * test fix * test fix * fix: flaky test --- mesa_frames/__init__.py | 44 ++++ mesa_frames/abstract/__init__.py | 42 +++ mesa_frames/abstract/agents.py | 174 ++++++------ mesa_frames/abstract/mixin.py | 55 +++- mesa_frames/abstract/space.py | 336 +++++++++++++----------- mesa_frames/concrete/__init__.py | 60 +++++ mesa_frames/concrete/agents.py | 119 +++++---- mesa_frames/concrete/model.py | 123 ++++++--- mesa_frames/concrete/pandas/__init__.py | 50 ++++ mesa_frames/concrete/pandas/agentset.py | 184 +++++++------ mesa_frames/concrete/pandas/mixin.py | 47 +++- mesa_frames/concrete/pandas/space.py | 56 ++++ mesa_frames/concrete/polars/__init__.py | 56 ++++ mesa_frames/concrete/polars/agentset.py | 138 +++++----- mesa_frames/concrete/polars/mixin.py | 45 ++++ mesa_frames/concrete/polars/space.py | 46 +++- mesa_frames/types_.py | 2 + pyproject.toml | 12 + 18 files changed, 1076 insertions(+), 513 deletions(-) diff --git a/mesa_frames/__init__.py b/mesa_frames/__init__.py index 95b6906..9914980 100644 --- a/mesa_frames/__init__.py +++ b/mesa_frames/__init__.py @@ -1,3 +1,47 @@ +""" +mesa-frames: High-performance extension for the Mesa agent-based modeling framework. + +mesa-frames extends the Mesa framework to support complex simulations with thousands +of agents by storing agents in DataFrames. This approach significantly enhances +performance and scalability while maintaining a syntax similar to Mesa. + +Key Features: +- Utilizes DataFrame storage for agents, enabling vectorized operations +- Supports both pandas and Polars as backend libraries +- Provides similar syntax to Mesa for ease of transition +- Allows for vectorized functions when simultaneous activation of agents is possible +- Implements SIMD processing for optimized simultaneous operations +- Includes GridDF for efficient grid-based spatial modeling + +Main Components: +- AgentSetPandas: Agent set implementation using pandas backend +- AgentSetPolars: Agent set implementation using Polars backend +- ModelDF: Base model class for mesa-frames +- GridDF: Grid space implementation for spatial modeling + +Usage: +To use mesa-frames, import the necessary components and subclass them as needed: + + from mesa_frames import AgentSetPolars, ModelDF, GridDF + + class MyAgent(AgentSetPolars): + # Define your agent logic here + + class MyModel(ModelDF): + def __init__(self, width, height): + super().__init__() + self.grid = GridDF(width, height, self) + # Define your model logic here + +Note: mesa-frames is in early development. API and usage patterns may change. + +For more detailed information, refer to the full documentation and API reference. + +Developed by: Adam Amer +License: MIT +GitHub: https://github.com/adamamer20/mesa_frames +""" + from mesa_frames.concrete.agents import AgentsDF from mesa_frames.concrete.model import ModelDF from mesa_frames.concrete.pandas.agentset import AgentSetPandas diff --git a/mesa_frames/abstract/__init__.py b/mesa_frames/abstract/__init__.py index e69de29..40bfddb 100644 --- a/mesa_frames/abstract/__init__.py +++ b/mesa_frames/abstract/__init__.py @@ -0,0 +1,42 @@ +""" +mesa-frames abstract components. + +This package contains abstract base classes and mixins that define the core +interfaces and shared functionality for the mesa-frames extension. + +Classes: + agents.py: + - AgentContainer: Abstract base class for agent containers. + - AgentSetDF: Abstract base class for agent sets using DataFrames. + + mixin.py: + - CopyMixin: Mixin class providing fast copy functionality. + - DataFrameMixin: Mixin class defining the interface for DataFrame operations. + + space.py: + - SpaceDF: Abstract base class for all space classes. + - DiscreteSpaceDF: Abstract base class for discrete space classes (Grids and Networks). + - GridDF: Abstract base class for grid classes. + +These abstract classes and mixins provide the foundation for the concrete +implementations in mesa-frames, ensuring consistent interfaces and shared +functionality across different backend implementations (e.g., pandas, Polars). + +Usage: + These classes are not meant to be instantiated directly. Instead, they + should be inherited by concrete implementations in the mesa-frames package. + + For example: + + from mesa_frames.abstract import AgentSetDF, DataFrameMixin + + class ConcreteAgentSet(AgentSetDF): + # Implement abstract methods here + ... + +Note: + The abstract classes use Python's ABC (Abstract Base Class) module to define + abstract methods that must be implemented by concrete subclasses. + +For more detailed information on each class, refer to their individual docstrings. +""" diff --git a/mesa_frames/abstract/agents.py b/mesa_frames/abstract/agents.py index 011c8a5..c9dbd22 100644 --- a/mesa_frames/abstract/agents.py +++ b/mesa_frames/abstract/agents.py @@ -1,3 +1,45 @@ +""" +Abstract base classes for agent containers in mesa-frames. + +This module defines the core abstractions for agent containers in the mesa-frames +extension. It provides the foundation for implementing agent storage and +manipulation using DataFrame-based approaches. + +Classes: + AgentContainer(CopyMixin): + An abstract base class that defines the common interface for all agent + containers in mesa-frames. It inherits from CopyMixin to provide fast + copying functionality. + + AgentSetDF(AgentContainer, DataFrameMixin): + An abstract base class for agent sets that use DataFrames as the underlying + storage mechanism. It inherits from both AgentContainer and DataFrameMixin + to combine agent container functionality with DataFrame operations. + +These abstract classes are designed to be subclassed by concrete implementations +that use specific DataFrame libraries (e.g., pandas, Polars) as their backend. + +Usage: + These classes should not be instantiated directly. Instead, they should be + subclassed to create concrete implementations: + + from mesa_frames.abstract.agents import AgentSetDF + + class AgentSetPandas(AgentSetDF): + def __init__(self, model): + super().__init__(model) + # Implementation using pandas DataFrame + ... + + # Implement other abstract methods + +Note: + The abstract methods in these classes use Python's @abstractmethod decorator, + ensuring that concrete subclasses must implement these methods. + +Attributes and methods of each class are documented in their respective docstrings. +""" + from __future__ import annotations # PEP 563: postponed evaluation of type annotations from abc import abstractmethod @@ -28,31 +70,6 @@ class AgentContainer(CopyMixin): """An abstract class for containing agents. Defines the common interface for AgentSetDF and AgentsDF. - Methods - ------- - copy(deep: bool = False, memo: dict | None = None) -> Self - Create a copy of the AgentContainer. - discard(ids: IdsLike, inplace: bool = True) -> Self - Removes an agent from the AgentContainer. Does not raise an error if the agent is not found. - add(other: Any, inplace: bool = True) -> Self - Add agents to the AgentContainer. - contains(ids: IdsLike) -> bool | BoolSeries - Check if agents with the specified IDs are in the AgentContainer. - do(method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any | dict[str, Any] - Invoke a method on the AgentContainer. - get(attr_names: str | Collection[str] | None = None, mask: AgentMask | 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: IdsLike, inplace: bool = True) -> Self - Removes an agent from the AgentContainer. - select(mask: AgentMask | None = None, filter_func: Callable[[Self], AgentMask] | 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: AgentMask | 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 @@ -80,7 +97,7 @@ def discard( agents: IdsLike | AgentSetDF | Collection[AgentSetDF], inplace: bool = True, ) -> Self: - """Removes agents from the AgentContainer. Does not raise an error if the agent is not found. + """Remove agents from the AgentContainer. Does not raise an error if the agent is not found. Parameters ---------- @@ -216,7 +233,7 @@ def get( attr_names: str | Collection[str] | None = None, mask: AgentMask | None = None, ) -> Series | dict[str, Series] | DataFrame | dict[str, DataFrame]: - """Retrieves the value of a specified attribute for each agent in the AgentContainer. + """Retrieve the value of a specified attribute for each agent in the AgentContainer. Parameters ---------- @@ -238,7 +255,7 @@ def remove( agents: IdsLike | AgentSetDF | Collection[AgentSetDF], inplace: bool = True, ) -> Self: - """Removes the agents from the AgentContainer + """Remove the agents from the AgentContainer. Parameters ---------- @@ -313,7 +330,7 @@ def set( mask: AgentMask | None = None, inplace: bool = True, ) -> Self: - """Sets the value of a specified attribute or attributes for each agent in the mask in AgentContainer. + """Set the value of a specified attribute or attributes for each agent in the mask in AgentContainer. Parameters ---------- @@ -379,7 +396,21 @@ def sort( A new or updated AgentContainer. """ - def __add__(self, other) -> Self: + def __add__( + self, other: DataFrameInput | AgentSetDF | Collection[AgentSetDF] + ) -> Self: + """Add agents to a new AgentContainer through the + operator. + + Parameters + ---------- + other : DataFrameInput | AgentSetDF | Collection[AgentSetDF] + The agents to add. + + Returns + ------- + Self + A new AgentContainer with the added agents. + """ return self.add(agents=other, inplace=False) def __contains__(self, agents: int | AgentSetDF) -> bool: @@ -417,7 +448,7 @@ def __getitem__( | tuple[AgentMask, Collection[str]] ), ) -> Series | DataFrame | dict[str, Series] | dict[str, DataFrame]: - """Implements the [] operator for the AgentContainer. + """Implement 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. @@ -531,7 +562,7 @@ def __setitem__( @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). + """Fallback for retrieving attributes of the AgentContainer. Retrieve an attribute of the underlying DataFrame(s). Parameters ---------- @@ -615,7 +646,8 @@ def random(self) -> Generator: Returns ------- - Generator""" + Generator + """ return self.model.random @property @@ -709,53 +741,6 @@ def pos(self) -> DataFrame | dict[str, DataFrame]: class AgentSetDF(AgentContainer, DataFrameMixin): """The AgentSetDF class is a container for agents of the same type. - 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: AgentMask, 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: AgentMask | None = None) -> Series | DataFrame - Retrieve the value of a specified attribute for each agent in the AgentSetDF. - remove(self, ids: AgentMask, inplace: bool = True) -> Self - Removes an agent from the AgentSetDF. - select(self, mask: AgentMask | None = None, filter_func: Callable[[Self], AgentMask] | 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: AgentMask | 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] | AgentMask | tuple[AgentMask, str] | tuple[AgentMask, 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 @@ -790,7 +775,7 @@ def add( agents: DataFrameInput, inplace: bool = True, ) -> Self: - """Add agents to the AgentSetDF + """Add agents to the AgentSetDF. Agents can be the input to the DataFrame constructor. So, the input can be: - A DataFrame: adds the agents from the DataFrame. @@ -812,6 +797,20 @@ def add( ... def discard(self, agents: IdsLike, inplace: bool = True) -> Self: + """Remove an agent from the AgentSetDF. Does not raise an error if the agent is not found. + + Parameters + ---------- + agents : IdsLike + The ids to remove + inplace : bool, optional + Whether to remove the agent in place, by default True + + Returns + ------- + Self + The updated AgentSetDF. + """ return super().discard(agents, inplace) @overload @@ -896,9 +895,9 @@ def get( ) -> Series | DataFrame: ... @abstractmethod - def step(self) -> None: ... - - """A single step of the AgentSetDF. This method should be overridden by subclasses.""" + def step(self) -> None: + """Run a single step of the AgentSetDF. This method should be overridden by subclasses.""" + ... def remove(self, agents: IdsLike, inplace: bool = True) -> Self: if agents is None or (isinstance(agents, Iterable) and len(agents) == 0): @@ -922,7 +921,7 @@ def _concatenate_agentsets( @abstractmethod def _get_bool_mask(self, mask: AgentMask) -> BoolSeries: - """Get the equivalent boolean mask based on the input mask + """Get the equivalent boolean mask based on the input mask. Parameters ---------- @@ -936,7 +935,7 @@ def _get_bool_mask(self, mask: AgentMask) -> BoolSeries: @abstractmethod def _get_masked_df(self, mask: AgentMask) -> DataFrame: - """Get the df filtered by the input mask + """Get the df filtered by the input mask. Parameters ---------- @@ -966,8 +965,7 @@ def _get_obj_copy( @abstractmethod def _discard(self, ids: IdsLike) -> Self: - """Removes an agent from the DataFrame of the AgentSetDF. Gets called by self.model.agents.remove and self.model.agents.discard. - + """Remove an agent from the DataFrame of the AgentSetDF. Gets called by self.model.agents.remove and self.model.agents.discard. Parameters ---------- diff --git a/mesa_frames/abstract/mixin.py b/mesa_frames/abstract/mixin.py index a2453dd..99ad85d 100644 --- a/mesa_frames/abstract/mixin.py +++ b/mesa_frames/abstract/mixin.py @@ -1,3 +1,47 @@ +""" +Mixin classes for mesa-frames abstract components. + +This module defines mixin classes that provide common functionality and interfaces +for various components in the mesa-frames extension. These mixins are designed to +be used with the abstract base classes to create flexible and extensible +implementations. + +Classes: + CopyMixin(ABC): + A mixin class that provides a fast copy method for classes that inherit it. + This is useful for creating efficient copies of large data structures, such + as DataFrames containing agent data. + + DataFrameMixin(ABC): + A mixin class that defines an interface for DataFrame operations. This mixin + provides a common set of methods that should be implemented by concrete + backend classes (e.g., pandas or Polars implementations) to ensure consistent + DataFrame manipulation across the mesa-frames package. + +These mixin classes are not meant to be instantiated directly. Instead, they should +be inherited alongside other base classes to add specific functionality or to +enforce a common interface. + +Usage: + Mixin classes are typically used in multiple inheritance scenarios: + + from mesa_frames.abstract.mixin import CopyMixin, DataFrameMixin + + class MyDataFrameClass(SomeBaseClass, CopyMixin, DataFrameMixin): + def __init__(self): + super().__init__() + # Implementation + + # Implement abstract methods from DataFrameMixin + +Note: + The DataFrameMixin uses Python's @abstractmethod decorator for its methods, + ensuring that classes inheriting from it must implement these methods. + +Attributes and methods of each mixin class are documented in their respective +docstrings. +""" + from abc import ABC, abstractmethod from collections.abc import Collection, Hashable, Iterator, Sequence from copy import copy, deepcopy @@ -9,14 +53,7 @@ class CopyMixin(ABC): - """A mixin class that provides a fast copy method for the class that inherits it. - - Methods - ------- - copy(deep: bool = False, memo: dict | None = None) -> Self - Create a copy of the object. If deep is True, a deep copy is created. If deep is False, a shallow copy is created. - - """ + """A mixin class that provides a fast copy method for the class that inherits it.""" _copy_with_method: dict[str, tuple[str, list[str]]] = {} _copy_only_reference: list[str] = [ @@ -137,6 +174,8 @@ def __deepcopy__(self, memo: dict) -> Self: class DataFrameMixin(ABC): + """A mixin class which defines an interface for DataFrame operations. Most methods are abstract and should be implemented by the concrete backend.""" + def _df_remove(self, df: DataFrame, mask: Mask, index_cols: str) -> DataFrame: return self._df_get_masked_df(df, index_cols, mask, negate=True) diff --git a/mesa_frames/abstract/space.py b/mesa_frames/abstract/space.py index 7b618c8..c8f7919 100644 --- a/mesa_frames/abstract/space.py +++ b/mesa_frames/abstract/space.py @@ -1,3 +1,52 @@ +""" +Abstract base classes for spatial components in mesa-frames. + +This module defines the core abstractions for spatial structures in the mesa-frames +extension. It provides the foundation for implementing various types of spaces, +including discrete spaces and grids, using DataFrame-based approaches for improved +performance and scalability. + +Classes: + SpaceDF(CopyMixin, DataFrameMixin): + An abstract base class that defines the common interface for all space + classes in mesa-frames. It combines fast copying functionality with + DataFrame operations. + + DiscreteSpaceDF(SpaceDF): + An abstract base class for discrete space implementations, such as grids + and networks. It extends SpaceDF with methods specific to discrete spaces. + + GridDF(DiscreteSpaceDF): + An abstract base class for grid-based spaces. It inherits from + DiscreteSpaceDF and adds grid-specific functionality. + +These abstract classes are designed to be subclassed by concrete implementations +that use specific DataFrame libraries (e.g., pandas, Polars) as their backend. +They provide a common interface and shared functionality across different types +of spatial structures in agent-based models. + +Usage: + These classes should not be instantiated directly. Instead, they should be + subclassed to create concrete implementations: + + from mesa_frames.abstract.space import GridDF + + class GridPandas(GridDF): + def __init__(self, model, dimensions, torus, capacity, neighborhood_type): + super().__init__(model, dimensions, torus, capacity, neighborhood_type) + # Implementation using pandas DataFrame + ... + + # Implement other abstract methods + +Note: + The abstract methods in these classes use Python's @abstractmethod decorator, + ensuring that concrete subclasses must implement these methods. + +Attributes and methods of each class are documented in their respective docstrings. +For more detailed information on each class, refer to their individual docstrings. +""" + from abc import abstractmethod from collections.abc import Callable, Collection, Sequence, Sized from itertools import product @@ -35,63 +84,7 @@ class SpaceDF(CopyMixin, DataFrameMixin): - """The SpaceDF class is an abstract class that defines the interface for all space classes in mesa_frames. - - Methods - ------- - __init__(model: 'ModelDF') - Create a new SpaceDF object. - random_agents(n: int) -> DataFrame - Return a random sample of agents from the space. - get_directions( - pos0: SpaceCoordinate | SpaceCoordinates | None = None, - pos1: SpaceCoordinate | SpaceCoordinates | None = None, - agents0: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, - agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, - normalize: bool = False, - ) -> DataFrame - Returns the directions from pos0 to pos1 or agents0 and agents1. - get_distances( - pos0: SpaceCoordinate | SpaceCoordinates | None = None, - pos1: SpaceCoordinate | SpaceCoordinates | None = None, - agents0: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, - agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, - ) -> DataFrame - Returns the distances from pos0 to pos1 or agents0 and agents1. - get_neighbors( - radius: int | float | Sequence[int] | Sequence[float] | ArrayLike, - pos: Space - ) -> DataFrame - Get the neighboring agents from given positions or agents according to the specified radiuses. - move_agents( - agents: IdsLike | AgentContainer | Collection[AgentContainer], - pos - ) -> Self - Place agents in the space according to the specified coordinates. - move_to_empty( - agents: IdsLike | AgentContainer | Collection[AgentContainer], - inplace: bool = True, - ) -> Self - Move agents to empty cells/positions in the space. - random_pos( - n: int, - ) -> DataFrame - Return a random sample of positions from the space. - remove_agents( - agents: IdsLike | AgentContainer | Collection[AgentContainer], - inplace: bool = True, - ) - Remove agents from the space. - swap_agents( - agents0: IdsLike | AgentContainer | Collection[AgentContainer], - agents1: IdsLike | AgentContainer | Collection[AgentContainer], - ) -> Self - Swap the positions of the agents in the space. - - Parameters - ---------- - model : 'ModelDF' - """ + """The SpaceDF class is an abstract class that defines the interface for all space classes in mesa_frames.""" _model: "ModelDF" _agents: DataFrame # | GeoDataFrame # Stores the agents placed in the space @@ -103,6 +96,12 @@ class SpaceDF(CopyMixin, DataFrameMixin): ] # The column names of the positions in the _agents dataframe (eg. ['dim_0', 'dim_1', ...] in Grids, ['node_id', 'edge_id'] in Networks) def __init__(self, model: "ModelDF") -> None: + """Create a new SpaceDF. + + Parameters + ---------- + model : ModelDF + """ self._model = model def move_agents( @@ -111,8 +110,9 @@ def move_agents( pos: SpaceCoordinate | SpaceCoordinates, inplace: bool = True, ) -> Self: - """Move agents in the Space to the specified coordinates. If some agents are not placed, - raises a RuntimeWarning. + """Move agents in the Space to the specified coordinates. + + If some agents are not placed,raises a RuntimeWarning. Parameters ---------- @@ -196,6 +196,7 @@ def swap_agents( inplace: bool = True, ) -> Self: """Swap the positions of the agents in the space. + agents0 and agents1 must have the same length and all agents must be placed in the space. Parameters @@ -249,7 +250,8 @@ def get_directions( agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, normalize: bool = False, ) -> DataFrame: - """Returns the directions from pos0 to pos1 or agents0 and agents1. + """Return the directions from pos0 to pos1 or agents0 and agents1. + If the space is a Network, the direction is the shortest path between the two nodes. In all other cases, the direction is the direction vector between the two positions. Either positions (pos0, pos1) or agents (agents0, agents1) must be specified, not both and they must have the same length. @@ -282,7 +284,8 @@ def get_distances( agents0: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, agents1: IdsLike | AgentContainer | Collection[AgentContainer] | None = None, ) -> DataFrame: - """Returns the distances from pos0 to pos1 or agents0 and agents1. + """Return the distances from pos0 to pos1 or agents0 and agents1. + If the space is a Network, the distance is the number of nodes of the shortest path between the two nodes. In all other cases, the distance is Euclidean/l2/Frobenius norm. You should specify either positions (pos0, pos1) or agents (agents0, agents1), not both and they must have the same length. @@ -314,6 +317,7 @@ def get_neighbors( include_center: bool = False, ) -> DataFrame: """Get the neighboring agents from given positions or agents according to the specified radiuses. + Either positions (pos0, pos1) or agents (agents0, agents1) must be specified, not both and they must have the same length. Parameters @@ -409,6 +413,8 @@ def remove_agents( ) -> Self: """Remove agents from the space. + Does not remove the agents from the model. + Parameters ---------- agents : IdsLike | AgentContainer | Collection[AgentContainer] @@ -462,7 +468,9 @@ def _place_or_move_agents( pos: SpaceCoordinate | SpaceCoordinates, is_move: bool, ) -> Self: - """A unique method for moving or placing agents (only the RuntimeWarning changes). + """Move or place agents. + + Only the runtime warning change. Parameters ---------- @@ -479,14 +487,28 @@ def _place_or_move_agents( """ @abstractmethod - def __repr__(self) -> str: ... + def __repr__(self) -> str: + """Return a string representation of the SpaceDF. + + Returns + ------- + str + """ + ... @abstractmethod - def __str__(self) -> str: ... + def __str__(self) -> str: + """Return a string representation of the SpaceDF. + + Returns + ------- + str + """ + ... @property def agents(self) -> DataFrame: # | GeoDataFrame: - """Get the ids of the agents placed in the cell set, along with their coordinates or geometries + """Get the ids of the agents placed in the cell set, along with their coordinates or geometries. Returns ------- @@ -516,31 +538,7 @@ def random(self) -> Generator: class DiscreteSpaceDF(SpaceDF): - """The DiscreteSpaceDF class is an abstract class that defines the interface for all discrete space classes (Grids and Networks) in mesa_frames. - - Methods - ------- - __init__(model: 'ModelDF', capacity: int | None = None) - Create a new DiscreteSpaceDF object. - is_available(pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame - Check whether the input positions are available (there exists at least one remaining spot in the cells). - is_empty(pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame - Check whether the input positions are empty (there isn't any single agent in the cells). - is_full(pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame - Check whether the input positions are full (there isn't any spot available in the cells). - move_to_empty(agents: IdsLike | AgentContainer | Collection[AgentContainer], inplace: bool = True) -> Self - Move agents to empty cells in the space (cells where there isn't any single agent). - move_to_available(agents: IdsLike | AgentContainer | Collection[AgentContainer], inplace: bool = True) -> Self - Move agents to available cells in the space (cells where there is at least one spot available). - sample_cells(n: int, cell_type: Literal["any", "empty", "available", "full"] = "any", with_replacement: bool = True) -> DataFrame - Sample cells from the grid according to the specified cell_type. - get_neighborhood(radius: int | float | Sequence[int] | Sequence[float] | ArrayLike, pos: DiscreteCoordinate | DiscreteCoordinates | None = None, agents: IdsLike | AgentContainer | Collection[AgentContainer] = None, include_center: bool = False) -> DataFrame - Get the neighborhood cells from the given positions (pos) or agents according to the specified radiuses. - get_cells(coords: DiscreteCoordinate | DiscreteCoordinates | None = None) -> DataFrame - Retrieve a dataframe of specified cells with their properties and agents. - set_cells(properties: DataFrame, cells: DiscreteCoordinates | None = None, inplace: bool = True) -> Self - Set the properties of the specified cells. - """ + """The DiscreteSpaceDF class is an abstract class that defines the interface for all discrete space classes (Grids and Networks) in mesa_frames.""" _agents: DataFrame _capacity: int | None # The maximum capacity for cells (default is infinite) @@ -554,22 +552,20 @@ def __init__( model: "ModelDF", capacity: int | None = None, ): - """Create a DiscreteSpaceDF object. - NOTE: The capacity specified here is the default capacity, - it can be set also per cell through the set_cells method. + """Create a new DiscreteSpaceDF. Parameters ---------- model : ModelDF The model to which the space belongs capacity : int | None, optional - The maximum capacity for cells, by default None (infinite) + The maximum capacity for cells (default is infinite), by default None """ super().__init__(model) self._capacity = capacity def is_available(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: - """Check whether the input positions are available (there exists at least one remaining spot in the cells) + """Check whether the input positions are available (there exists at least one remaining spot in the cells). Parameters ---------- @@ -584,7 +580,7 @@ def is_available(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFra return self._check_cells(pos, "available") def is_empty(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: - """Check whether the input positions are empty (there isn't any single agent in the cells) + """Check whether the input positions are empty (there isn't any single agent in the cells). Parameters ---------- @@ -599,7 +595,7 @@ def is_empty(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: return self._check_cells(pos, "empty") def is_full(self, pos: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: - """Check whether the input positions are full (there isn't any spot available in the cells) + """Check whether the input positions are full (there isn't any spot available in the cells). Parameters ---------- @@ -718,6 +714,7 @@ def set_cells( inplace: bool = True, ) -> Self: """Set the properties of the specified cells. + This method mirrors the functionality of mesa's PropertyLayer, but allows also to set properties only of specific cells. Either the cells DF must contain both the cells' coordinates and the properties or the cells' coordinates can be specified separately with the cells argument. @@ -778,6 +775,7 @@ def get_neighborhood( include_center: bool = False, ) -> DataFrame: """Get the neighborhood cells from the given positions (pos) or agents according to the specified radiuses. + Either positions (pos) or agents must be specified, not both. Parameters @@ -989,15 +987,48 @@ def _update_capacity_agents( """ ... - def __getitem__(self, cells: DiscreteCoordinate | DiscreteCoordinates): + def __getitem__(self, cells: DiscreteCoordinate | DiscreteCoordinates) -> DataFrame: + """Get the properties and agents of the specified cells. + + Parameters + ---------- + cells : DiscreteCoordinate | DiscreteCoordinates + The cells to get the properties for + + Returns + ------- + DataFrame + A DataFrame with the properties and agents of the cells + """ return self.get_cells(cells) def __getattr__(self, key: str) -> DataFrame: + """Get the properties of the cells. + + Parameters + ---------- + key : str + The property to get + + Returns + ------- + DataFrame + A DataFrame with the properties of the cells + """ # Fallback, if key (property) is not found in the object, # then it must mean that it's in the _cells dataframe return self._cells[key] def __setitem__(self, cells: DiscreteCoordinates, properties: DataFrame): + """Set the properties of the specified cells. + + Parameters + ---------- + cells : DiscreteCoordinates + The cells to set the properties for + properties : DataFrame + The properties to set + """ self.set_cells(cells=cells, properties=properties, inplace=True) def __repr__(self) -> str: @@ -1011,6 +1042,8 @@ def __str__(self) -> str: @property def cells(self) -> DataFrame: """ + Obtain the properties and agents of the cells in the grid. + Returns ------- DataFrame @@ -1024,18 +1057,39 @@ def cells(self, df: DataFrame): @property def empty_cells(self) -> DataFrame: + """Get the empty cells (cells without any agent) in the grid. + + Returns + ------- + DataFrame + A DataFrame with the empty cells + """ return self._sample_cells( None, with_replacement=False, condition=self._empty_cell_condition ) @property def available_cells(self) -> DataFrame: + """Get the available cells (cells with at least one spot available) in the grid. + + Returns + ------- + DataFrame + A DataFrame with the available cells + """ return self._sample_cells( None, with_replacement=False, condition=self._available_cell_condition ) @property def full_cells(self) -> DataFrame: + """Get the full cells (cells without any spot available) in the grid. + + Returns + ------- + DataFrame + A DataFrame with the full cells + """ return self._sample_cells( None, with_replacement=False, condition=self._full_cell_condition ) @@ -1055,6 +1109,7 @@ def remaining_capacity(self) -> int | None: class GridDF(DiscreteSpaceDF): """The GridDF class is an abstract class that defines the interface for all grid classes in mesa-frames. + Inherits from DiscreteSpaceDF. Warning @@ -1075,23 +1130,6 @@ class GridDF(DiscreteSpaceDF): .. _np.genfromtxt: https://numpy.org/doc/stable/reference/generated/numpy.genfromtxt.html .. _mesa-examples Sugarscape model: https://github.com/projectmesa/mesa-examples/blob/e137a60e4e2f2546901bec497e79c4a7b0cc69bb/examples/sugarscape_g1mt/sugarscape_g1mt/model.py#L93-L94 - - - Methods - ------- - __init__(model: 'ModelDF', dimensions: Sequence[int], torus: bool = False, capacity: int | None = None, neighborhood_type: str = 'moore') - Create a new GridDF object. - out_of_bounds(pos: GridCoordinate | GridCoordinates) -> DataFrame - Check whether the input positions are out of bounds in a non-toroidal grid. - - Properties - ---------- - dimensions : Sequence[int] - The dimensions of the grid - neighborhood_type : Literal['moore', 'von_neumann', 'hexagonal'] - The type of neighborhood to consider - torus : bool - If the grid is a torus """ _cells_capacity: ( @@ -1111,42 +1149,20 @@ def __init__( capacity: int | None = None, neighborhood_type: str = "moore", ): - """Create a new GridDF object. - - Warning - ------- - For rectangular grids: - In this implementation, [0, ..., 0] is the bottom-left corner and - [dimensions[0]-1, ..., dimensions[n-1]-1] is the top-right corner, consistent with - Cartesian coordinates and Matplotlib/Seaborn plot outputs. - The convention is different from `np.genfromtxt`_ and its use in the - `mesa-examples Sugarscape model`_, where [0, ..., 0] is the top-left corner - and [dimensions[0]-1, ..., dimensions[n-1]-1] is the bottom-right corner. - - For hexagonal grids: - The coordinates are ordered according to the axial coordinate system. - In this system, the hexagonal grid uses two axes (q and r) at 60 degrees to each other. - The q-axis points to the right, and the r-axis points up and to the right. - The [0, 0] coordinate is at the bottom-left corner of the grid. - - .. _np.genfromtxt: https://numpy.org/doc/stable/reference/generated/numpy.genfromtxt.html - .. _mesa-examples Sugarscape model: https://github.com/projectmesa/mesa-examples/blob/e137a60e4e2f2546901bec497e79c4a7b0cc69bb/examples/sugarscape_g1mt/sugarscape_g1mt/model.py#L93-L94 + """Create a new GridDF. Parameters ---------- - model : 'ModelDF' - The model object to which the grid belongs - dimensions: Sequence[int] - The dimensions of the grid. For hexagonal grids, this should be [q_max, r_max]. + model : ModelDF + The model to which the space belongs + dimensions : Sequence[int] + The dimensions of the grid torus : bool, optional - If the grid should be a torus, by default False + If the grid is a torus, by default False capacity : int | None, optional - The maximum number of agents that can be placed in a cell, by default None - neighborhood_type: str, optional - The type of neighborhood to consider, by default 'moore'. - If 'moore', the neighborhood is the 8 cells around the center cell (up, down, left, right, and diagonals). - If 'von_neumann', the neighborhood is the 4 cells around the center cell (up, down, left, right). - If 'hexagonal', the neighborhood are 6 cells around the center cell distributed in a hexagonal shape. + The maximum capacity for cells (default is infinite), by default None + neighborhood_type : str, optional + The type of neighborhood to consider, by default "moore" """ super().__init__(model, capacity) self._dimensions = dimensions @@ -1758,12 +1774,38 @@ def _generate_empty_grid( @property def dimensions(self) -> Sequence[int]: + """The dimensions of the grid. + + They are set uniquely at the creation of the grid. + + Returns + ------- + Sequence[int] + The dimensions of the grid + """ return self._dimensions @property - def neighborhood_type(self) -> str: + def neighborhood_type(self) -> Literal["moore", "von_neumann", "hexagonal"]: + """The type of neighborhood to consider (moore, von_neumann, hexagonal). + + It is set uniquely at the creation of the grid. + + Returns + ------- + Literal['moore', 'von_neumann', 'hexagonal'] + """ return self._neighborhood_type @property def torus(self) -> bool: + """If the grid is a torus (wraps around at the edges). + + Can be set uniquely at the creation of the grid. + + Returns + ------- + bool + Whether the grid is a torus + """ return self._torus diff --git a/mesa_frames/concrete/__init__.py b/mesa_frames/concrete/__init__.py index e69de29..f2cc5e4 100644 --- a/mesa_frames/concrete/__init__.py +++ b/mesa_frames/concrete/__init__.py @@ -0,0 +1,60 @@ +""" +Concrete implementations of mesa-frames components. + +This package contains the concrete implementations of the abstract base classes +defined in mesa_frames.abstract. It provides ready-to-use classes for building +agent-based models using DataFrame-based storage, with support for both pandas +and Polars backends. + +Subpackages: + pandas: Contains pandas-based implementations of agent sets, mixins, and spatial structures. + polars: Contains Polars-based implementations of agent sets, mixins, and spatial structures. + +Modules: + agents: Defines the AgentsDF class, a collection of AgentSetDFs. + model: Provides the ModelDF class, the base class for models in mesa-frames. + +Classes: + From pandas.agentset: + AgentSetPandas(AgentSetDF, PandasMixin): A pandas-based implementation of the AgentSet. + + From pandas.mixin: + PandasMixin(DataFrameMixin): A pandas-based implementation of DataFrame operations. + + From pandas.space: + GridPandas(GridDF, PandasMixin): A pandas-based implementation of Grid. + + From polars subpackage: + Similar classes as in the pandas subpackage, but using Polars as the backend. + + From agents: + AgentsDF(AgentContainer): A collection of AgentSetDFs. All agents of the model are stored here. + + From model: + ModelDF: Base class for models in the mesa-frames library. + +Usage: + Users can import the concrete implementations directly from this package: + + from mesa_frames.concrete import ModelDF, AgentsDF + from mesa_frames.concrete.pandas import AgentSetPandas, GridPandas + + # For Polars-based implementations + from mesa_frames.concrete.polars import AgentSetPolars, GridPolars + + class MyModel(ModelDF): + def __init__(self): + super().__init__() + self.agents.add(AgentSetPandas(self)) + self.space = GridPandas(self, dimensions=[10, 10]) + # ... other initialization code + +Note: + The choice between pandas and Polars implementations depends on the user's + preference and performance requirements. Both provide similar functionality + but may have different performance characteristics depending on the specific + use case. + +For more detailed information on each class, refer to their respective module +and class docstrings. +""" diff --git a/mesa_frames/concrete/agents.py b/mesa_frames/concrete/agents.py index 08be785..adf4752 100644 --- a/mesa_frames/concrete/agents.py +++ b/mesa_frames/concrete/agents.py @@ -1,3 +1,49 @@ +""" +Concrete implementation of the agents collection for mesa-frames. + +This module provides the concrete implementation of the agents collection class +for the mesa-frames library. It defines the AgentsDF class, which serves as a +container for all agent sets in a model, leveraging DataFrame-based storage for +improved performance. + +Classes: + AgentsDF(AgentContainer): + A collection of AgentSetDFs. This class acts as a container for all + agents in the model, organizing them into separate AgentSetDF instances + based on their types. + +The AgentsDF class is designed to be used within ModelDF instances to manage +all agents in the simulation. It provides methods for adding, removing, and +accessing agents and agent sets, while taking advantage of the performance +benefits of DataFrame-based agent storage. + +Usage: + The AgentsDF class is typically instantiated and used within a ModelDF subclass: + + from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.agents import AgentsDF + from mesa_frames.concrete.pandas import AgentSetPandas + + class MyCustomModel(ModelDF): + def __init__(self): + super().__init__() + # Adding agent sets to the collection + self.agents += AgentSetPandas(self) + self.agents += AnotherAgentSetPandas(self) + + def step(self): + # Step all agent sets + self.agents.do("step") + +Note: + This concrete implementation builds upon the abstract AgentContainer class + defined in the mesa_frames.abstract package, providing a ready-to-use + agents collection that integrates with the DataFrame-based agent storage system. + +For more detailed information on the AgentsDF class and its methods, refer to +the class docstring. +""" + from collections import defaultdict from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import TYPE_CHECKING, Literal, cast @@ -21,17 +67,8 @@ class AgentsDF(AgentContainer): - _agentsets: list[AgentSetDF] - _ids: pl.Series """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[AgentSetDF, tuple[str, list[str]]] - A dictionary of attributes to copy with a specified method and arguments. - Properties ---------- active_agents(self) -> dict[AgentSetDF, pd.DataFrame] @@ -45,53 +82,19 @@ class AgentsDF(AgentContainer): 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 | Iterable[AgentSetDF], inplace: bool = True) -> Self - Add agents to the AgentsDF. - contains(self, ids: IdsLike) -> bool | pl.Series - Check if agents with the specified IDs are in the AgentsDF. - copy(self, deep: bool = False, memo: dict | None = None) -> Self - Create a copy of the AgentsDF. - discard(self, ids: IdsLike, inplace: bool = True) -> Self - Remove an agent from the AgentsDF. Does not raise an error if the agent is not found. - do(self, method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any - Invoke a method on the AgentsDF. - get(self, attr_names: str | Collection[str] | None = None, mask: AgentMask = None) -> dict[AgentSetDF, Series] | dict[AgentSetDF, DataFrame] - Retrieve the value of a specified attribute for each agent in the AgentsDF. - remove(self, ids: IdsLike, inplace: bool = True) -> Self - Remove agents from the AgentsDF. - select(self, mask: AgentMask = None, filter_func: Callable[[Self], AgentMask] | 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[AgentSetDF, Any] | None = None, values: Any | None = None, mask: AgentMask | 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. - _check_ids(self, other: AgentSetDF | Iterable[AgentSetDF]) -> None - Check if the IDs of the agents to be added are unique. - __add__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self - Add AgentSetDFs to a new AgentsDF through the + operator. - __getattr__(self, key: str) -> Any - Retrieve an attribute of the underlying agent sets. - __iadd__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self - Add AgentSetDFs to the AgentsDF through the += operator. - __iter__(self) -> Iterator - Get an iterator for the agents in the AgentsDF. - __len__(self) -> int - 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. """ + _agentsets: list[AgentSetDF] + _ids: pl.Series + def __init__(self, model: "ModelDF") -> None: + """Initialize a new AgentsDF. + + Parameters + ---------- + model : ModelDF + The model associated with the AgentsDF. + """ self._model = model self._agentsets = [] self._ids = pl.Series(name="unique_id", dtype=pl.Int64) @@ -422,7 +425,7 @@ def _get_bool_masks( def _return_agentsets_list( self, agentsets: AgentSetDF | Iterable[AgentSetDF] ) -> list[AgentSetDF]: - """Convert the agentsets to a list of AgentSetDF + """Convert the agentsets to a list of AgentSetDF. Parameters ---------- @@ -588,6 +591,14 @@ def active_agents( @property def agentsets_by_type(self) -> dict[type[AgentSetDF], Self]: + """Get the agent sets in the AgentsDF grouped by type. + + Returns + ------- + dict[type[AgentSetDF], Self] + A dictionary mapping agent set types to the corresponding AgentsDF. + """ + def copy_without_agentsets() -> Self: return self.copy(deep=False, skip=["_agentsets"]) diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py index 959922b..bd09637 100644 --- a/mesa_frames/concrete/model.py +++ b/mesa_frames/concrete/model.py @@ -1,3 +1,45 @@ +""" +Concrete implementation of the model class for mesa-frames. + +This module provides the concrete implementation of the base model class for +the mesa-frames library. It defines the ModelDF class, which serves as the +foundation for creating agent-based models using DataFrame-based agent storage. + +Classes: + ModelDF: + The base class for models in the mesa-frames library. This class + provides the core functionality for initializing and running + agent-based simulations using DataFrame-backed agent sets. + +The ModelDF class is designed to be subclassed by users to create specific +model implementations. It provides the basic structure and methods necessary +for setting up and running simulations, while leveraging the performance +benefits of DataFrame-based agent storage. + +Usage: + To create a custom model, subclass ModelDF and implement the necessary + methods: + + from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.agents import AgentSetPandas + + class MyCustomModel(ModelDF): + def __init__(self, num_agents): + super().__init__() + self.agents += AgentSetPandas(self) + # Initialize your model-specific attributes and agent sets + + def run_model(self): + # Implement the logic for a single step of your model + for _ in range(10): + self.step() + + # Add any other custom methods for your model + +For more detailed information on the ModelDF class and its methods, refer to +the class docstring. +""" + from collections.abc import Sequence from typing import TYPE_CHECKING @@ -17,25 +59,6 @@ class ModelDF: It includes the basic attributes and methods necessary for initializing and running a simulation 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). - Properties ---------- agents : AgentsDF @@ -51,7 +74,9 @@ class ModelDF: _space: SpaceDF | None # This will be a MultiSpaceDF object def __init__(self, seed: int | Sequence[int] | None = None) -> None: - """Create a new model. Overload this method with the actual code to + """Create a new model. + + Overload this method with the actual code to start the model. Always start with super().__init__(seed) to initialize the model object properly. @@ -85,25 +110,12 @@ def get_agents_of_type(self, agent_type: type) -> "AgentSetDF": 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: + Parameters ---------- - seed : int | None + seed : int | Sequence[int] | None A new seed for the RNG; if None, reset using the current seed """ if seed is None: @@ -113,18 +125,34 @@ def reset_randomizer(self, seed: int | Sequence[int] | None) -> None: 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. + """Run the model until the end condition is reached. + + Overload as needed. """ while self.running: self.step() def step(self) -> None: - """A single step. The default method calls the step() method of all agents. Overload as needed.""" + """Run a single step. + + The default method calls the step() method of all agents. Overload as needed. + """ self.agents.step() @property def agents(self) -> AgentsDF: + """Get the AgentsDF object containing all agents in the model. + + Returns + ------- + AgentsDF + The AgentsDF object containing all agents in the model. + + Raises + ------ + ValueError + If the model has not been initialized properly with super().__init__(). + """ try: return self._agents except AttributeError: @@ -140,10 +168,29 @@ def agents(self, agents: AgentsDF) -> None: @property def agent_types(self) -> list[type]: + """Get a list of different agent types present in the model. + + Returns + ------- + list[type] + A list of the different agent types present in the model. + """ return [agent.__class__ for agent in self._agents._agentsets] @property def space(self) -> SpaceDF: + """Get the space object associated with the model. + + Returns + ------- + SpaceDF + The space object associated with the model. + + Raises + ------ + ValueError + If the space has not been set for the model. + """ if not self._space: raise ValueError( "You haven't set the space for the model. Use model.space = your_space" diff --git a/mesa_frames/concrete/pandas/__init__.py b/mesa_frames/concrete/pandas/__init__.py index e69de29..3783581 100644 --- a/mesa_frames/concrete/pandas/__init__.py +++ b/mesa_frames/concrete/pandas/__init__.py @@ -0,0 +1,50 @@ +""" +Pandas-based implementations for mesa-frames. + +This subpackage contains concrete implementations of mesa-frames components +using pandas as the backend for DataFrame operations. It provides high-performance, +pandas-based classes for agent sets, spatial structures, and DataFrame operations. + +Modules: + agentset: Defines the AgentSetPandas class, a pandas-based implementation of AgentSet. + mixin: Provides the PandasMixin class, implementing DataFrame operations using pandas. + space: Contains the GridPandas class, a pandas-based implementation of Grid. + +Classes: + AgentSetPandas(AgentSetDF, PandasMixin): + A pandas-based implementation of the AgentSet, using pandas DataFrames + for efficient agent storage and manipulation. + + PandasMixin(DataFrameMixin): + A mixin class that implements DataFrame operations using pandas, + providing methods for data manipulation and analysis. + + GridPandas(GridDF, PandasMixin): + A pandas-based implementation of Grid, using pandas DataFrames for + efficient spatial operations and agent positioning. + +Usage: + These classes can be imported and used directly in mesa-frames models: + + from mesa_frames.concrete.pandas import AgentSetPandas, GridPandas + from mesa_frames.concrete.model import ModelDF + + class MyAgents(AgentSetPandas): + def __init__(self, model): + super().__init__(model) + # Initialize agents + + class MyModel(ModelDF): + def __init__(self, width, height): + super().__init__() + self.agents.add(MyAgents(self)) + self.space = GridPandas(self, dimensions=[width, height]) + +Note: + Using these pandas-based implementations requires pandas to be installed. + The performance characteristics will depend on the pandas version and the + specific operations used in the model. + +For more detailed information on each class, refer to their respective module +and class docstrings. +""" diff --git a/mesa_frames/concrete/pandas/agentset.py b/mesa_frames/concrete/pandas/agentset.py index 6248384..758a8d2 100644 --- a/mesa_frames/concrete/pandas/agentset.py +++ b/mesa_frames/concrete/pandas/agentset.py @@ -1,103 +1,90 @@ +""" +Pandas-based implementation of AgentSet for mesa-frames. + +This module provides a concrete implementation of the AgentSet class using pandas +as the backend for DataFrame operations. It defines the AgentSetPandas class, +which combines the abstract AgentSetDF functionality with pandas-specific +operations for efficient agent management and manipulation. + +Classes: + AgentSetPandas(AgentSetDF, PandasMixin): + A pandas-based implementation of the AgentSet. This class uses pandas + DataFrames to store and manipulate agent data, providing high-performance + operations for large numbers of agents. + +The AgentSetPandas class is designed to be used within ModelDF instances or as +part of an AgentsDF collection. It leverages the power of pandas for fast and +efficient data operations on agent attributes and behaviors. + +Usage: + The AgentSetPandas class can be used directly in a model or as part of an + AgentsDF collection: + + from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.pandas.agentset import AgentSetPandas + import numpy as np + + class MyAgents(AgentSetPandas): + def __init__(self, model): + super().__init__(model) + # Initialize with some agents + self.add({'unique_id': np.arange(100), 'wealth': 10}) + + def step(self): + # Implement step behavior using pandas operations + self.agents['wealth'] += 1 + + class MyModel(ModelDF): + def __init__(self): + super().__init__() + self.agents += MyAgents(self) + + def step(self): + self.agents.step() + +Note: + This implementation relies on pandas, so users should ensure that pandas + is installed and imported. The performance characteristics of this class + will depend on the pandas version and the specific operations used. + +For more detailed information on the AgentSetPandas class and its methods, +refer to the class docstring. +""" + from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import TYPE_CHECKING import pandas as pd import polars as pl +from typing_extensions import Any, Self, overload + from mesa_frames.abstract.agents import AgentSetDF from mesa_frames.concrete.pandas.mixin import PandasMixin from mesa_frames.concrete.polars.agentset import AgentSetPolars from mesa_frames.types_ import AgentPandasMask, PandasIdsLike -from typing_extensions import Any, Self, overload if TYPE_CHECKING: from mesa_frames.concrete.model import ModelDF -class AgentSetPandas(AgentSetDF, PandasMixin): +class AgentSetPandas(AgentSetDF, PandasMixin): # noqa : D101 _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: PandasIdsLike) -> bool | pd.Series - Check if agents with the specified IDs are in the AgentSetPandas. - copy(self, deep: bool = False, memo: dict | None = None) -> Self - Create a copy of the AgentSetPandas. - discard(self, ids: PandasIdsLike, inplace: bool = True) -> Self - Remove an agent from the AgentSetPandas. Does not raise an error if the agent is not found. - do(self, method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any - Invoke a method on the AgentSetPandas. - get(self, attr_names: str | Collection[str] | None, mask: PandasMaskLike = None) -> pd.Series | pd.DataFrame - Retrieve the value of a specified attribute for each agent in the AgentSetPandas. - remove(self, ids: PandasIdsLike, inplace: bool = True) -> Self - Remove agents from the AgentSetPandas. - select(self, mask: PandasMaskLike = None, filter_func: Callable[[Self], PandasMaskLike] | None = None, n: int | None = None, negate: bool = False, inplace: bool = True) -> Self - Select agents in the AgentSetPandas based on the given criteria. - 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: + """Initialize a new AgentSetPandas. + + Overload this method to add custom initialization logic but make sure to call super().__init__(model). + + Parameters + ---------- + model : ModelDF + The model associated with the AgentSetPandas. + """ self._model = model self._agents = ( pd.DataFrame(columns=["unique_id"]) @@ -106,7 +93,7 @@ def __init__(self, model: "ModelDF") -> None: ) self._mask = pd.Series(True, index=self._agents.index, dtype=pd.BooleanDtype()) - def add( + def add( # noqa : D102 self, agents: pd.DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True, @@ -156,7 +143,7 @@ def contains(self, agents: int) -> bool: ... @overload def contains(self, agents: PandasIdsLike) -> pd.Series: ... - def contains(self, agents: PandasIdsLike) -> bool | pd.Series: + def contains(self, agents: PandasIdsLike) -> bool | pd.Series: # noqa : D102 if isinstance(agents, pd.Series): return agents.isin(self._agents.index) elif isinstance(agents, pd.Index): @@ -168,7 +155,7 @@ def contains(self, agents: PandasIdsLike) -> bool | pd.Series: else: return agents in self._agents.index - def get( + def get( # noqa : D102 self, attr_names: str | Collection[str] | None = None, mask: AgentPandasMask = None, @@ -184,7 +171,7 @@ def get( if isinstance(attr_names, Collection): return self._agents.loc[mask, list(attr_names)] - def set( + def set( # noqa : D102 self, attr_names: str | dict[str, Any] | Collection[str] | None = None, values: Any | None = None, @@ -222,7 +209,7 @@ def set( obj._agents = obj._agents.reindex(original_index) return obj - def select( + def select( # noqa : D102 self, mask: AgentPandasMask = None, filter_func: Callable[[Self], AgentPandasMask] | None = None, @@ -244,12 +231,12 @@ def select( obj._mask = bool_mask return obj - def shuffle(self, inplace: bool = True) -> Self: + def shuffle(self, inplace: bool = True) -> Self: # noqa : D102 obj = self._get_obj(inplace) obj._agents = obj._agents.sample(frac=1) return obj - def sort( + def sort( # noqa : D102 self, by: str | Sequence[str], ascending: bool | Sequence[bool] = True, @@ -261,6 +248,15 @@ def sort( return obj def to_polars(self) -> AgentSetPolars: + """Convert the AgentSetPandas to an AgentSetPolars. + + NOTE: If a methods is not backend-agnostic (i.e., it uses pandas-specific functionality), when the method is called on the Polars version of the object, it will raise an error. + + Returns + ------- + AgentSetPolars + An AgentSetPolars object with the same agents and active agents as the AgentSetPandas. + """ new_obj = AgentSetPolars(self._model) new_obj._agents = pl.DataFrame(self._agents) new_obj._mask = pl.Series(self._mask) @@ -388,24 +384,24 @@ def _update_mask( dtype=pd.BooleanDtype(), ) - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> Any: # noqa : D105 super().__getattr__(name) return getattr(self._agents, name) - def __iter__(self) -> Iterator[dict[str, Any]]: + def __iter__(self) -> Iterator[dict[str, Any]]: # noqa : D105 for index, row in self._agents.iterrows(): row_dict = row.to_dict() row_dict["unique_id"] = index yield row_dict - def __len__(self) -> int: + def __len__(self) -> int: # noqa : D105 return len(self._agents) - def __reversed__(self) -> Iterator: + def __reversed__(self) -> Iterator: # noqa : D105 return iter(self._agents[::-1].iterrows()) @property - def agents(self) -> pd.DataFrame: + def agents(self) -> pd.DataFrame: # noqa : D105 return self._agents @agents.setter @@ -419,7 +415,7 @@ def agents(self, new_agents: pd.DataFrame) -> None: self._agents = new_agents @property - def active_agents(self) -> pd.DataFrame: + def active_agents(self) -> pd.DataFrame: # noqa : D102 return self._agents.loc[self._mask] @active_agents.setter @@ -427,13 +423,13 @@ def active_agents(self, mask: AgentPandasMask) -> None: self.select(mask=mask, inplace=True) @property - def inactive_agents(self) -> pd.DataFrame: + def inactive_agents(self) -> pd.DataFrame: # noqa : D102 return self._agents.loc[~self._mask] @property - def index(self) -> pd.Index: + def index(self) -> pd.Index: # noqa : D102 return self._agents.index @property - def pos(self) -> pd.DataFrame: + def pos(self) -> pd.DataFrame: # noqa : D102 return super().pos diff --git a/mesa_frames/concrete/pandas/mixin.py b/mesa_frames/concrete/pandas/mixin.py index 27579a8..a809325 100644 --- a/mesa_frames/concrete/pandas/mixin.py +++ b/mesa_frames/concrete/pandas/mixin.py @@ -1,13 +1,52 @@ +""" +Pandas-specific mixin for DataFrame operations in mesa-frames. + +This module provides a concrete implementation of the DataFrameMixin using pandas +as the backend for DataFrame operations. It defines the PandasMixin class, which +implements DataFrame operations specific to pandas. + +Classes: + PandasMixin(DataFrameMixin): + A pandas-based implementation of DataFrame operations. This class provides + methods for manipulating data stored in pandas DataFrames, + tailored for use in mesa-frames components like AgentSetPandas and GridPandas. + +The PandasMixin class is designed to be used as a mixin with other mesa-frames +classes, providing them with pandas-specific DataFrame functionality. It implements +the abstract methods defined in the DataFrameMixin, ensuring consistent DataFrame +operations across the mesa-frames package. + +Usage: + The PandasMixin is typically used in combination with other base classes: + + from mesa_frames.abstract import AgentSetDF + from mesa_frames.concrete.pandas.mixin import PandasMixin + + class AgentSetPandas(AgentSetDF, PandasMixin): + def __init__(self, model): + super().__init__(model) + ... + + def _some_private_method(self): + # Use pandas operations provided by the mixin + result = self._df_add(self.agents, 10) + # ... further processing ... + + +For more detailed information on the PandasMixin class and its methods, refer to +the class docstring. +""" + from collections.abc import Callable, Collection, Hashable, Iterator, Sequence from typing import Literal import numpy as np -from mesa_frames.abstract.mixin import DataFrameMixin -from mesa_frames.types_ import DataFrame, PandasMask -from typing_extensions import Any, overload - import pandas as pd import polars as pl +from typing_extensions import Any, overload + +from mesa_frames.abstract.mixin import DataFrameMixin +from mesa_frames.types_ import DataFrame, PandasMask class PandasMixin(DataFrameMixin): diff --git a/mesa_frames/concrete/pandas/space.py b/mesa_frames/concrete/pandas/space.py index f63c9d0..8889906 100644 --- a/mesa_frames/concrete/pandas/space.py +++ b/mesa_frames/concrete/pandas/space.py @@ -1,8 +1,64 @@ +""" +Pandas-based implementation of spatial structures for mesa-frames. + +This module provides concrete implementations of spatial structures using pandas +as the backend for DataFrame operations. It defines the GridPandas class, which +implements a 2D grid structure using pandas DataFrames for efficient spatial +operations and agent positioning. + +Classes: + GridPandas(GridDF, PandasMixin): + A pandas-based implementation of a 2D grid. This class uses pandas + DataFrames to store and manipulate spatial data, providing high-performance + operations for large-scale spatial simulations. + +The GridPandas class is designed to be used within ModelDF instances to represent +the spatial environment of the simulation. It leverages the power of pandas for +fast and efficient data operations on spatial attributes and agent positions. + +Usage: + The GridPandas class can be used directly in a model to represent the + spatial environment: + + from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.pandas.space import GridPandas + from mesa_frames.concrete.pandas.agentset import AgentSetPandas + + class MyAgents(AgentSetPandas): + # ... agent implementation ... + + class MyModel(ModelDF): + def __init__(self, width, height): + super().__init__() + self.space = GridPandas(self, [width, height]) + self.agents += MyAgents(self) + + def step(self): + # Move agents + self.space.move_agents(self.agents, positions) + # ... other model logic ... + +Features: + - Efficient storage and retrieval of agent positions + - Fast operations for moving agents and querying neighborhoods + - Seamless integration with pandas-based agent sets + - Support for various boundary conditions (e.g., wrapped, bounded) + +Note: + This implementation relies on pandas, so users should ensure that pandas + is installed and imported. The performance characteristics of this class + will depend on the pandas version and the specific operations used. + +For more detailed information on the GridPandas class and its methods, +refer to the class docstring. +""" + from collections.abc import Callable, Sequence from typing import Literal import numpy as np import pandas as pd + from mesa_frames.abstract.space import GridDF from mesa_frames.concrete.pandas.mixin import PandasMixin diff --git a/mesa_frames/concrete/polars/__init__.py b/mesa_frames/concrete/polars/__init__.py index e69de29..2faa9e1 100644 --- a/mesa_frames/concrete/polars/__init__.py +++ b/mesa_frames/concrete/polars/__init__.py @@ -0,0 +1,56 @@ +""" +Polars-based implementations for mesa-frames. + +This subpackage contains concrete implementations of mesa-frames components +using Polars as the backend for DataFrame operations. It provides high-performance, +Polars-based classes for agent sets, spatial structures, and DataFrame operations. + +Modules: + agentset: Defines the AgentSetPolars class, a Polars-based implementation of AgentSet. + mixin: Provides the PolarsMixin class, implementing DataFrame operations using Polars. + space: Contains the GridPolars class, a Polars-based implementation of Grid. + +Classes: + AgentSetPolars(AgentSetDF, PolarsMixin): + A Polars-based implementation of the AgentSet, using Polars DataFrames + for efficient agent storage and manipulation. + + PolarsMixin(DataFrameMixin): + A mixin class that implements DataFrame operations using Polars, + providing methods for data manipulation and analysis. + + GridPolars(GridDF, PolarsMixin): + A Polars-based implementation of Grid, using Polars DataFrames for + efficient spatial operations and agent positioning. + +Usage: + These classes can be imported and used directly in mesa-frames models: + + from mesa_frames.concrete.polars import AgentSetPolars, GridPolars + from mesa_frames.concrete.model import ModelDF + + class MyAgents(AgentSetPolars): + def __init__(self, model): + super().__init__(model) + # Initialize agents + + class MyModel(ModelDF): + def __init__(self, width, height): + super().__init__() + self.agents = MyAgents(self) + self.grid = GridPolars(width, height, self) + +Features: + - High-performance DataFrame operations using Polars + - Efficient memory usage and fast computation + - Support for lazy evaluation and query optimization + - Seamless integration with other mesa-frames components + +Note: + Using these Polars-based implementations requires Polars to be installed. + Polars offers excellent performance for large datasets and complex operations, + making it suitable for large-scale agent-based models. + +For more detailed information on each class, refer to their respective module +and class docstrings. +""" diff --git a/mesa_frames/concrete/polars/agentset.py b/mesa_frames/concrete/polars/agentset.py index 2a1a59f..e3d0fc3 100644 --- a/mesa_frames/concrete/polars/agentset.py +++ b/mesa_frames/concrete/polars/agentset.py @@ -1,3 +1,62 @@ +""" +Polars-based implementation of AgentSet for mesa-frames. + +This module provides a concrete implementation of the AgentSet class using Polars +as the backend for DataFrame operations. It defines the AgentSetPolars class, +which combines the abstract AgentSetDF functionality with Polars-specific +operations for efficient agent management and manipulation. + +Classes: + AgentSetPolars(AgentSetDF, PolarsMixin): + A Polars-based implementation of the AgentSet. This class uses Polars + DataFrames to store and manipulate agent data, providing high-performance + operations for large numbers of agents. + +The AgentSetPolars class is designed to be used within ModelDF instances or as +part of an AgentsDF collection. It leverages the power of Polars for fast and +efficient data operations on agent attributes and behaviors. + +Usage: + The AgentSetPolars class can be used directly in a model or as part of an + AgentsDF collection: + + from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.polars.agentset import AgentSetPolars + import polars as pl + + class MyAgents(AgentSetPolars): + def __init__(self, model): + super().__init__(model) + # Initialize with some agents + self.add(pl.DataFrame({'id': range(100), 'wealth': 10})) + + def step(self): + # Implement step behavior using Polars operations + self.agents = self.agents.with_columns(new_wealth = pl.col('wealth') + 1) + + class MyModel(ModelDF): + def __init__(self): + super().__init__() + self.agents += MyAgents(self) + + def step(self): + self.agents.step() + +Features: + - Efficient storage and manipulation of large agent populations + - Fast vectorized operations on agent attributes + - Support for lazy evaluation and query optimization + - Seamless integration with other mesa-frames components + +Note: + This implementation relies on Polars, so users should ensure that Polars + is installed and imported. The performance characteristics of this class + will depend on the Polars version and the specific operations used. + +For more detailed information on the AgentSetPolars class and its methods, +refer to the class docstring. +""" + from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import TYPE_CHECKING @@ -22,85 +81,6 @@ class AgentSetPolars(AgentSetDF, PolarsMixin): _copy_only_reference: list[str] = ["_model", "_mask"] _mask: pl.Expr | pl.Series - """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: PolarsIdsLike) -> bool | pl.Series - Check if agents with the specified IDs are in the AgentSetPolars. - copy(self, deep: bool = False, memo: dict | None = None) -> Self - Create a copy of the AgentSetPolars. - discard(self, ids: PolarsIdsLike, inplace: bool = True) -> Self - Remove an agent from the AgentSetPolars. Does not raise an error if the agent is not found. - do(self, method_name: str, *args, return_results: bool = False, inplace: bool = True, **kwargs) -> Self | Any - Invoke a method on the AgentSetPolars. - get(self, attr_names: IntoExpr | Iterable[IntoExpr] | None, mask: PolarsMaskLike = None) -> pl.Series | pl.DataFrame - Retrieve the value of a specified attribute for each agent in the AgentSetPolars. - remove(self, ids: PolarsIdsLike, inplace: bool = True) -> Self - Remove agents from the AgentSetPolars. - select(self, mask: PolarsMaskLike = None, filter_func: Callable[[Self], PolarsMaskLike] | None = None, n: int | None = None, negate: bool = False, inplace: bool = True) -> Self - Select agents in the AgentSetPolars based on the given criteria. - 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. diff --git a/mesa_frames/concrete/polars/mixin.py b/mesa_frames/concrete/polars/mixin.py index d0a69a9..4736723 100644 --- a/mesa_frames/concrete/polars/mixin.py +++ b/mesa_frames/concrete/polars/mixin.py @@ -1,3 +1,48 @@ +""" +Polars-specific mixin for DataFrame operations in mesa-frames. + +This module provides a concrete implementation of the DataFrameMixin using Polars +as the backend for DataFrame operations. It defines the PolarsMixin class, which +implements DataFrame operations specific to Polars, offering efficient data +manipulation and analysis capabilities for mesa-frames components. + +Classes: + PolarsMixin(DataFrameMixin): + A Polars-based implementation of DataFrame operations. This class provides + methods for manipulating and analyzing data stored in Polars DataFrames, + tailored for use in mesa-frames components like AgentSetPolars and GridPolars. + +The PolarsMixin class is designed to be used as a mixin with other mesa-frames +classes, providing them with Polars-specific DataFrame functionality. It implements +the abstract methods defined in the DataFrameMixin, ensuring consistent DataFrame +operations across the mesa-frames package. + +Usage: + The PolarsMixin is typically used in combination with other base classes: + + from mesa_frames.abstract import AgentSetDF + from mesa_frames.concrete.polars.mixin import PolarsMixin + + class AgentSetPolars(AgentSetDF, PolarsMixin): + def __init__(self, model): + super().__init__(model) + self.agents = pl.DataFrame() # Initialize empty DataFrame + + def some_method(self): + # Use Polars operations provided by the mixin + result = self._df_groupby(self.agents, 'some_column') + # ... further processing ... + +Features: + - High-performance DataFrame operations using Polars + - Support for both eager and lazy evaluation + - Efficient memory usage and fast computation + - Integration with Polars' query optimization capabilities + +For more detailed information on the PolarsMixin class and its methods, refer to +the class docstring. +""" + from collections.abc import Callable, Collection, Hashable, Iterator, Sequence from typing import Literal diff --git a/mesa_frames/concrete/polars/space.py b/mesa_frames/concrete/polars/space.py index 3a8389a..27c9b90 100644 --- a/mesa_frames/concrete/polars/space.py +++ b/mesa_frames/concrete/polars/space.py @@ -1,8 +1,52 @@ +""" +Polars-based implementation of spatial structures for mesa-frames. + +This module provides concrete implementations of spatial structures using Polars +as the backend for DataFrame operations. It defines the GridPolars class, which +implements a 2D grid structure using Polars DataFrames for efficient spatial +operations and agent positioning. + +Classes: + GridPolars(GridDF, PolarsMixin): + A Polars-based implementation of a 2D grid. This class uses Polars + DataFrames to store and manipulate spatial data, providing high-performance + operations for large-scale spatial simulations. + +The GridPolars class is designed to be used within ModelDF instances to represent +the spatial environment of the simulation. It leverages the power of Polars for +fast and efficient data operations on spatial attributes and agent positions. + +Usage: + The GridPolars class can be used directly in a model to represent the + spatial environment: + + from mesa_frames.concrete.model import ModelDF + from mesa_frames.concrete.polars.space import GridPolars + from mesa_frames.concrete.polars.agentset import AgentSetPolars + + class MyAgents(AgentSetPolars): + # ... agent implementation ... + + class MyModel(ModelDF): + def __init__(self, width, height): + super().__init__() + self.space = GridPolars(self, [width, height]) + self.agents += MyAgents(self) + + def step(self): + # Move agents + self.space.move_agents(self.agents) + # ... other model logic ... + +For more detailed information on the GridPolars class and its methods, +refer to the class docstring. +""" + from collections.abc import Callable, Sequence +from typing import Literal import numpy as np import polars as pl -from typing import Literal from mesa_frames.abstract.space import GridDF from mesa_frames.concrete.polars.mixin import PolarsMixin diff --git a/mesa_frames/types_.py b/mesa_frames/types_.py index 4c20cc3..06d13f0 100644 --- a/mesa_frames/types_.py +++ b/mesa_frames/types_.py @@ -1,3 +1,5 @@ +"""Type aliases for the mesa_frames package.""" + from collections.abc import Collection, Sequence from typing import Literal diff --git a/pyproject.toml b/pyproject.toml index 1e68eb0..26db853 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,3 +48,15 @@ features = ["dev"] [tool.hatch.build.targets.wheel] packages = ["mesa_frames"] + +[tool.ruff.lint] +select = ["D"] +ignore = ["D101", "D102", "D105"] + +[tool.ruff.lint.pydocstyle] +convention = "numpy" + +[tool.ruff.lint.per-file-ignores] +"tests/*" = ["D"] +"examples/*" = ["D"] +"docs/*" = ["D"] \ No newline at end of file