+"""
+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
+
+import numpy as np
+
+from mesa_frames.abstract.space import SpaceDF
+from mesa_frames.concrete.agents import AgentsDF
+
+if TYPE_CHECKING:
+ from mesa_frames.abstract.agents import AgentSetDF
+
+
+
+[docs]
+class ModelDF:
+ """Base class for models in the mesa-frames library.
+
+ This class serves as a foundational structure for creating agent-based models.
+ It includes the basic attributes and methods necessary for initializing and
+ running a simulation model.
+
+ """
+
+ random: np.random.Generator
+ running: bool
+ _seed: int | Sequence[int]
+ _agents: AgentsDF # Where the agents are stored
+ _space: SpaceDF | None # This will be a MultiSpaceDF object
+
+
+[docs]
+ def __init__(self, seed: int | Sequence[int] | None = None) -> None:
+ """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.
+
+ Parameters
+ ----------
+ seed : int | Sequence[int] | None, optional
+ The seed for the model's generator
+ """
+ self.random = None
+ self.reset_randomizer(seed)
+ self.running = True
+ self.current_id = 0
+ self._agents = AgentsDF(self)
+ self._space = None
+
+
+
+[docs]
+ def get_agents_of_type(self, agent_type: type) -> "AgentSetDF":
+ """Retrieve the AgentSetDF of a specified type.
+
+ Parameters
+ ----------
+ agent_type : type
+ The type of AgentSetDF to retrieve.
+
+ Returns
+ -------
+ AgentSetDF
+ The AgentSetDF of the specified type.
+ """
+ for agentset in self._agents._agentsets:
+ if isinstance(agentset, agent_type):
+ return agentset
+ raise ValueError(f"No agents of type {agent_type} found in the model.")
+
+
+
+[docs]
+ def reset_randomizer(self, seed: int | Sequence[int] | None) -> None:
+ """Reset the model random number generator.
+
+ Parameters
+ ----------
+ seed : int | Sequence[int] | None
+ A new seed for the RNG; if None, reset using the current seed
+ """
+ if seed is None:
+ seed = np.random.SeedSequence().entropy
+ assert seed is not None
+ self._seed = seed
+ self.random = np.random.default_rng(seed=self._seed)
+
+
+
+[docs]
+ def run_model(self) -> None:
+ """Run the model until the end condition is reached.
+
+ Overload as needed.
+ """
+ while self.running:
+ self.step()
+
+
+
+[docs]
+ def step(self) -> None:
+ """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:
+ raise ValueError(
+ "You haven't called super().__init__() in your model. Make sure to call it in your __init__ method."
+ )
+
+ @agents.setter
+ def agents(self, agents: AgentsDF) -> None:
+ if not isinstance(agents, AgentsDF):
+ raise TypeError("agents must be an instance of AgentsDF")
+ self._agents = agents
+
+ @property
+ def agent_types(self) -> list[type]:
+ """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"
+ )
+ return self._space
+
+ @space.setter
+ def space(self, space: SpaceDF) -> None:
+ self._space = space
+
+
+"""
+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 mesa_frames.utils import copydoc
+
+if TYPE_CHECKING:
+ from mesa_frames.concrete.model import ModelDF
+
+
+
+[docs]
+@copydoc(AgentSetDF)
+class AgentSetPandas(AgentSetDF, PandasMixin):
+ """pandas-based implementation of AgentSetDF."""
+
+ _agents: pd.DataFrame
+ _mask: pd.Series
+ _copy_with_method: dict[str, tuple[str, list[str]]] = {
+ "_agents": ("copy", ["deep"]),
+ "_mask": ("copy", ["deep"]),
+ }
+
+
+[docs]
+ 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"])
+ .astype({"unique_id": "int64"})
+ .set_index("unique_id")
+ )
+ self._mask = pd.Series(True, index=self._agents.index, dtype=pd.BooleanDtype())
+
+
+
+[docs]
+ def add( # noqa : D102
+ self,
+ agents: pd.DataFrame | Sequence[Any] | dict[str, Any],
+ inplace: bool = True,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ if isinstance(agents, pd.DataFrame):
+ new_agents = agents
+ if "unique_id" != agents.index.name:
+ try:
+ new_agents.set_index("unique_id", inplace=True, drop=True)
+ except KeyError:
+ raise KeyError("DataFrame must have a unique_id column/index.")
+ elif isinstance(agents, dict):
+ if "unique_id" not in agents:
+ raise KeyError("Dictionary must have a unique_id key.")
+ index = agents.pop("unique_id")
+ if not isinstance(index, list):
+ index = [index]
+ new_agents = pd.DataFrame(agents, index=pd.Index(index, name="unique_id"))
+ else:
+ if len(agents) != len(obj._agents.columns) + 1:
+ raise ValueError(
+ "Length of data must match the number of columns in the AgentSet if being added as a Collection."
+ )
+ columns = pd.Index(["unique_id"]).append(obj._agents.columns.copy())
+ new_agents = pd.DataFrame([agents], columns=columns).set_index(
+ "unique_id", drop=True
+ )
+
+ if new_agents.index.dtype != "int64":
+ new_agents.index = new_agents.index.astype("int64")
+
+ if not obj._agents.index.intersection(new_agents.index).empty:
+ raise KeyError("Some IDs already exist in the agent set.")
+
+ original_active_indices = obj._mask.index[obj._mask].copy()
+
+ obj._agents = pd.concat([obj._agents, new_agents])
+
+ obj._update_mask(original_active_indices, new_agents.index)
+
+ return obj
+
+
+ @overload
+ def contains(self, agents: int) -> bool: ...
+
+ @overload
+ def contains(self, agents: PandasIdsLike) -> pd.Series: ...
+
+
+[docs]
+ 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):
+ return pd.Series(
+ agents.isin(self._agents.index), index=agents, dtype=pd.BooleanDtype()
+ )
+ elif isinstance(agents, Collection):
+ return pd.Series(list(agents), index=list(agents)).isin(self._agents.index)
+ else:
+ return agents in self._agents.index
+
+
+
+[docs]
+ def get( # noqa : D102
+ self,
+ attr_names: str | Collection[str] | None = None,
+ mask: AgentPandasMask = None,
+ ) -> pd.Index | pd.Series | pd.DataFrame:
+ mask = self._get_bool_mask(mask)
+ if attr_names is None:
+ return self._agents.loc[mask]
+ else:
+ if isinstance(attr_names, str) and attr_names == "unique_id":
+ return self._agents.loc[mask].index
+ if isinstance(attr_names, str):
+ return self._agents.loc[mask, attr_names]
+ if isinstance(attr_names, Collection):
+ return self._agents.loc[mask, list(attr_names)]
+
+
+
+[docs]
+ def set( # noqa : D102
+ self,
+ attr_names: str | dict[str, Any] | Collection[str] | None = None,
+ values: Any | None = None,
+ mask: AgentPandasMask = None,
+ inplace: bool = True,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ b_mask = obj._get_bool_mask(mask)
+ masked_df = obj._get_masked_df(mask)
+
+ if not attr_names:
+ attr_names = masked_df.columns
+
+ if isinstance(attr_names, dict):
+ for key, val in attr_names.items():
+ masked_df.loc[:, key] = val
+ elif (
+ isinstance(attr_names, str)
+ or (
+ isinstance(attr_names, Collection)
+ and all(isinstance(n, str) for n in attr_names)
+ )
+ ) and values is not None:
+ if not isinstance(attr_names, str): # isinstance(attr_names, Collection)
+ attr_names = list(attr_names)
+ masked_df.loc[:, attr_names] = values
+ else:
+ raise ValueError(
+ "Either attr_names must be a dictionary with columns as keys and values or values must be provided."
+ )
+
+ non_masked_df = obj._agents[~b_mask]
+ original_index = obj._agents.index
+ obj._agents = pd.concat([non_masked_df, masked_df])
+ obj._agents = obj._agents.reindex(original_index)
+ return obj
+
+
+
+[docs]
+ def select( # noqa : D102
+ self,
+ mask: AgentPandasMask = None,
+ filter_func: Callable[[Self], AgentPandasMask] | None = None,
+ n: int | None = None,
+ negate: bool = False,
+ inplace: bool = True,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ bool_mask = obj._get_bool_mask(mask)
+ if filter_func:
+ bool_mask = bool_mask & obj._get_bool_mask(filter_func(obj))
+ if negate:
+ bool_mask = ~bool_mask
+ if n is not None:
+ bool_mask = pd.Series(
+ obj._agents.index.isin(obj._agents[bool_mask].sample(n).index),
+ index=obj._agents.index,
+ )
+ obj._mask = bool_mask
+ return obj
+
+
+
+[docs]
+ def shuffle(self, inplace: bool = True) -> Self: # noqa : D102
+ obj = self._get_obj(inplace)
+ obj._agents = obj._agents.sample(frac=1)
+ return obj
+
+
+
+[docs]
+ def sort( # noqa : D102
+ self,
+ by: str | Sequence[str],
+ ascending: bool | Sequence[bool] = True,
+ inplace: bool = True,
+ **kwargs,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ obj._agents.sort_values(by=by, ascending=ascending, **kwargs, inplace=True)
+ return obj
+
+
+
+[docs]
+ 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)
+ return new_obj
+
+
+ def _concatenate_agentsets(
+ self,
+ agentsets: Iterable[Self],
+ duplicates_allowed: bool = True,
+ keep_first_only: bool = True,
+ original_masked_index: pd.Index | None = None,
+ ) -> Self:
+ if not duplicates_allowed:
+ indices = [self._agents.index.to_series()] + [
+ agentset._agents.index.to_series() for agentset in agentsets
+ ]
+ pd.concat(indices, verify_integrity=True)
+ if duplicates_allowed & keep_first_only:
+ final_df = self._agents.copy()
+ final_mask = self._mask.copy()
+ for obj in iter(agentsets):
+ final_df = final_df.combine_first(obj._agents)
+ final_mask = final_mask.combine_first(obj._mask)
+ else:
+ final_df = pd.concat([obj._agents for obj in agentsets])
+ final_mask = pd.concat([obj._mask for obj in agentsets])
+ self._agents = final_df
+ self._mask = final_mask
+ if not isinstance(original_masked_index, type(None)):
+ ids_to_remove = original_masked_index.difference(self._agents.index)
+ if not ids_to_remove.empty:
+ self.remove(ids_to_remove, inplace=True)
+ return self
+
+ def _get_bool_mask(
+ self,
+ mask: AgentPandasMask = None,
+ ) -> pd.Series:
+ if isinstance(mask, pd.Series) and mask.dtype == bool:
+ return mask
+ elif isinstance(mask, pd.DataFrame):
+ return pd.Series(
+ self._agents.index.isin(mask.index), index=self._agents.index
+ )
+ elif isinstance(mask, list):
+ return pd.Series(self._agents.index.isin(mask), index=self._agents.index)
+ elif mask is None or isinstance(mask, str) and mask == "all":
+ return pd.Series(True, index=self._agents.index)
+ elif isinstance(mask, str) and mask == "active":
+ return self._mask
+ elif isinstance(mask, Collection):
+ return pd.Series(self._agents.index.isin(mask), index=self._agents.index)
+ else:
+ return pd.Series(self._agents.index.isin([mask]), index=self._agents.index)
+
+ def _get_masked_df(
+ self,
+ mask: AgentPandasMask = None,
+ ) -> pd.DataFrame:
+ if isinstance(mask, pd.Series) and mask.dtype == bool:
+ return self._agents.loc[mask]
+ elif isinstance(mask, pd.DataFrame):
+ if mask.index.name != "unique_id":
+ if "unique_id" in mask.columns:
+ mask.set_index("unique_id", inplace=True, drop=True)
+ else:
+ raise KeyError("DataFrame must have a unique_id column/index.")
+ return pd.DataFrame(index=mask.index).join(
+ self._agents, on="unique_id", how="left"
+ )
+ elif isinstance(mask, pd.Series):
+ mask_df = mask.to_frame("unique_id").set_index("unique_id")
+ return mask_df.join(self._agents, on="unique_id", how="left")
+ elif mask is None or mask == "all":
+ return self._agents
+ elif mask == "active":
+ return self._agents.loc[self._mask]
+ else:
+ mask_series = pd.Series(mask)
+ mask_df = mask_series.to_frame("unique_id").set_index("unique_id")
+ return mask_df.join(self._agents, on="unique_id", how="left")
+
+ @overload
+ def _get_obj_copy(self, obj: pd.Series) -> pd.Series: ...
+
+ @overload
+ def _get_obj_copy(self, obj: pd.DataFrame) -> pd.DataFrame: ...
+
+ @overload
+ def _get_obj_copy(self, obj: pd.Index) -> pd.Index: ...
+
+ def _get_obj_copy(
+ self, obj: pd.Series | pd.DataFrame | pd.Index
+ ) -> pd.Series | pd.DataFrame | pd.Index:
+ return obj.copy()
+
+ def _discard(
+ self,
+ ids: PandasIdsLike,
+ ) -> Self:
+ mask = self._get_bool_mask(ids)
+ remove_ids = self._agents[mask].index
+ original_active_indices = self._mask.index[self._mask].copy()
+ self._agents.drop(remove_ids, inplace=True)
+ self._update_mask(original_active_indices)
+ return self
+
+ def _update_mask(
+ self,
+ original_active_indices: pd.Index,
+ new_active_indices: pd.Index | None = None,
+ ) -> None:
+ # Update the mask with the old active agents and the new agents
+ if new_active_indices is None:
+ self._mask = pd.Series(
+ self._agents.index.isin(original_active_indices),
+ index=self._agents.index,
+ dtype=pd.BooleanDtype(),
+ )
+ else:
+ self._mask = pd.Series(
+ self._agents.index.isin(original_active_indices)
+ | self._agents.index.isin(new_active_indices),
+ index=self._agents.index,
+ dtype=pd.BooleanDtype(),
+ )
+
+
+[docs]
+ def __getattr__(self, name: str) -> Any: # noqa : D105
+ super().__getattr__(name)
+ return getattr(self._agents, name)
+
+
+
+[docs]
+ 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
+
+
+
+
+
+
+[docs]
+ def __reversed__(self) -> Iterator: # noqa : D105
+ return iter(self._agents[::-1].iterrows())
+
+
+ @property
+ def agents(self) -> pd.DataFrame: # noqa : D105
+ return self._agents
+
+ @agents.setter
+ def agents(self, new_agents: pd.DataFrame) -> None:
+ if new_agents.index.name == "unique_id":
+ pass
+ elif "unique_id" in new_agents.columns:
+ new_agents.set_index("unique_id", inplace=True, drop=True)
+ else:
+ raise KeyError("The DataFrame should have a 'unique_id' index/column")
+ self._agents = new_agents
+
+ @property
+ def active_agents(self) -> pd.DataFrame: # noqa : D102
+ return self._agents.loc[self._mask]
+
+ @active_agents.setter
+ def active_agents(self, mask: AgentPandasMask) -> None:
+ self.select(mask=mask, inplace=True)
+
+ @property
+ def inactive_agents(self) -> pd.DataFrame: # noqa : D102
+ return self._agents.loc[~self._mask]
+
+ @property
+ def index(self) -> pd.Index: # noqa : D102
+ return self._agents.index
+
+ @property
+ def pos(self) -> pd.DataFrame: # noqa : D102
+ return super().pos
+
+
+"""
+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
+from mesa_frames.utils import copydoc
+
+
+
+[docs]
+@copydoc(GridDF)
+class GridPandas(GridDF, PandasMixin):
+ """pandas-based implementation of GridDF."""
+
+ _agents: pd.DataFrame
+ _copy_with_method: dict[str, tuple[str, list[str]]] = {
+ "_agents": ("copy", ["deep"]),
+ "_cells": ("copy", ["deep"]),
+ "_cells_capacity": ("copy", []),
+ "_offsets": ("copy", ["deep"]),
+ }
+ _cells: pd.DataFrame
+ _cells_capacity: np.ndarray
+ _offsets: pd.DataFrame
+
+ def _empty_cell_condition(self, cap: np.ndarray) -> np.ndarray:
+ # Create a boolean mask of the same shape as cap
+ empty_mask = np.ones_like(cap, dtype=bool)
+
+ if not self._agents.empty:
+ # Get the coordinates of all agents
+ agent_coords = self._agents[self._pos_col_names].to_numpy(int)
+
+ # Mark cells containing agents as not empty
+ empty_mask[tuple(agent_coords.T)] = False
+
+ return empty_mask
+
+ def _generate_empty_grid(
+ self, dimensions: Sequence[int], capacity: int
+ ) -> np.ndarray:
+ if not capacity:
+ capacity = np.inf
+ return np.full(dimensions, capacity)
+
+ def _sample_cells(
+ self,
+ n: int | None,
+ with_replacement: bool,
+ condition: Callable[[np.ndarray], np.ndarray],
+ respect_capacity: bool = True,
+ ) -> pd.DataFrame:
+ # Get the coordinates of cells that meet the condition
+ coords = np.array(np.where(condition(self._cells_capacity))).T
+
+ # If the grid has infinite capacity, there is no need to respect capacity
+ if np.any(self._cells_capacity == np.inf):
+ respect_capacity = False
+
+ if respect_capacity and condition != self._full_cell_condition:
+ capacities = self._cells_capacity[tuple(coords.T)]
+ else:
+ # If not respecting capacity or for full cells, set capacities to 1
+ capacities = np.ones(len(coords), dtype=int)
+
+ if n is not None:
+ if with_replacement:
+ if respect_capacity and condition != self._full_cell_condition:
+ assert (
+ n <= capacities.sum()
+ ), "Requested sample size exceeds the total available capacity."
+
+ sampled_coords = np.empty((0, coords.shape[1]), dtype=coords.dtype)
+ while len(sampled_coords) < n:
+ remaining_samples = n - len(sampled_coords)
+ sampled_indices = self.random.choice(
+ len(coords),
+ size=remaining_samples,
+ replace=True,
+ )
+ unique_indices, counts = np.unique(
+ sampled_indices, return_counts=True
+ )
+
+ if respect_capacity and condition != self._full_cell_condition:
+ # Calculate valid counts for each unique index
+ valid_counts = np.minimum(counts, capacities[unique_indices])
+ # Update capacities
+ capacities[unique_indices] -= valid_counts
+ else:
+ valid_counts = counts
+
+ # Create array of repeated coordinates
+ new_coords = np.repeat(coords[unique_indices], valid_counts, axis=0)
+ # Extend sampled_coords
+ sampled_coords = np.vstack((sampled_coords, new_coords))
+
+ if respect_capacity and condition != self._full_cell_condition:
+ # Update coords and capacities
+ mask = capacities > 0
+ coords = coords[mask]
+ capacities = capacities[mask]
+
+ sampled_coords = sampled_coords[:n]
+ self.random.shuffle(sampled_coords)
+ else:
+ assert n <= len(
+ coords
+ ), "Requested sample size exceeds the number of available cells."
+ sampled_indices = self.random.choice(len(coords), size=n, replace=False)
+ sampled_coords = coords[sampled_indices]
+ else:
+ sampled_coords = coords
+
+ # Convert the coordinates to a DataFrame
+ sampled_cells = pd.DataFrame(sampled_coords, columns=self._pos_col_names)
+ return sampled_cells
+
+ def _update_capacity_agents(
+ self,
+ agents: pd.DataFrame,
+ operation: Literal["movement", "removal"],
+ ) -> np.ndarray:
+ # Update capacity for agents that were already on the grid
+ masked_df = self._df_get_masked_df(
+ self._agents, index_cols="agent_id", mask=agents
+ )
+
+ if operation == "movement":
+ # Increase capacity at old positions
+ old_positions = tuple(masked_df[self._pos_col_names].to_numpy(int).T)
+ np.add.at(self._cells_capacity, old_positions, 1)
+
+ # Decrease capacity at new positions
+ new_positions = tuple(agents[self._pos_col_names].to_numpy(int).T)
+ np.add.at(self._cells_capacity, new_positions, -1)
+ elif operation == "removal":
+ # Increase capacity at the positions of removed agents
+ positions = tuple(masked_df[self._pos_col_names].to_numpy(int).T)
+ np.add.at(self._cells_capacity, positions, 1)
+ return self._cells_capacity
+
+ def _update_capacity_cells(self, cells: pd.DataFrame) -> np.ndarray:
+ # Get the coordinates of the cells to update
+ coords = cells.index
+
+ # Get the current capacity of updatable cells
+ current_capacity = self._cells.reindex(coords, fill_value=self._capacity)[
+ "capacity"
+ ].to_numpy()
+
+ # Calculate the number of agents currently in each cell
+ agents_in_cells = current_capacity - self._cells_capacity[tuple(zip(*coords))]
+
+ # Update the capacity in self._cells_capacity
+ new_capacity = cells["capacity"].to_numpy() - agents_in_cells
+
+ # Assert that no new capacity is negative
+ assert np.all(
+ new_capacity >= 0
+ ), "New capacity of a cell cannot be less than the number of agents in it."
+
+ self._cells_capacity[tuple(zip(*coords))] = new_capacity
+
+ return self._cells_capacity
+
+ @property
+ def remaining_capacity(self) -> int:
+ if not self._capacity:
+ return np.inf
+ return self._cells_capacity.sum()
+
+
+"""
+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
+
+import polars as pl
+from polars._typing import IntoExpr
+from typing_extensions import Any, Self, overload
+
+from mesa_frames.concrete.agents import AgentSetDF
+from mesa_frames.concrete.polars.mixin import PolarsMixin
+from mesa_frames.types_ import AgentPolarsMask, PolarsIdsLike
+from mesa_frames.utils import copydoc
+
+if TYPE_CHECKING:
+ from mesa_frames.concrete.model import ModelDF
+ from mesa_frames.concrete.pandas.agentset import AgentSetPandas
+
+
+
+[docs]
+@copydoc(AgentSetDF)
+class AgentSetPolars(AgentSetDF, PolarsMixin):
+ """Polars-based implementation of AgentSetDF."""
+
+ _agents: pl.DataFrame
+ _copy_with_method: dict[str, tuple[str, list[str]]] = {
+ "_agents": ("clone", []),
+ }
+ _copy_only_reference: list[str] = ["_model", "_mask"]
+ _mask: pl.Expr | pl.Series
+
+
+[docs]
+ def __init__(self, model: "ModelDF") -> None:
+ """Initialize a new AgentSetPolars.
+
+ Parameters
+ ----------
+ model : ModelDF
+ The model that the agent set belongs to.
+ """
+ self._model = model
+ self._agents = pl.DataFrame(schema={"unique_id": pl.Int64})
+ self._mask = pl.repeat(True, len(self._agents), dtype=pl.Boolean, eager=True)
+
+
+
+[docs]
+ def add(
+ self,
+ agents: pl.DataFrame | Sequence[Any] | dict[str, Any],
+ inplace: bool = True,
+ ) -> Self:
+ """Add agents to the AgentSetPolars.
+
+ Parameters
+ ----------
+ agents : pl.DataFrame | Sequence[Any] | dict[str, Any]
+ The agents to add.
+ inplace : bool, optional
+ Whether to add the agents in place, by default True.
+
+ Returns
+ -------
+ Self
+ The updated AgentSetPolars.
+ """
+ obj = self._get_obj(inplace)
+ if isinstance(agents, pl.DataFrame):
+ if "unique_id" not in agents.columns:
+ raise KeyError("DataFrame must have a unique_id column.")
+ new_agents = agents
+ elif isinstance(agents, dict):
+ if "unique_id" not in agents:
+ raise KeyError("Dictionary must have a unique_id key.")
+ new_agents = pl.DataFrame(agents)
+ else:
+ if len(agents) != len(obj._agents.columns):
+ raise ValueError(
+ "Length of data must match the number of columns in the AgentSet if being added as a Collection."
+ )
+ new_agents = pl.DataFrame([agents], schema=obj._agents.schema)
+
+ if new_agents["unique_id"].dtype != pl.Int64:
+ raise TypeError("unique_id column must be of type int64.")
+
+ # If self._mask is pl.Expr, then new mask is the same.
+ # If self._mask is pl.Series[bool], then new mask has to be updated.
+
+ if isinstance(obj._mask, pl.Series):
+ original_active_indices = obj._agents.filter(obj._mask)["unique_id"]
+
+ obj._agents = pl.concat([obj._agents, new_agents], how="diagonal_relaxed")
+
+ if isinstance(obj._mask, pl.Series):
+ obj._update_mask(original_active_indices, new_agents["unique_id"])
+
+ return obj
+
+
+ @overload
+ def contains(self, agents: int) -> bool: ...
+
+ @overload
+ def contains(self, agents: PolarsIdsLike) -> pl.Series: ...
+
+
+[docs]
+ def contains(
+ self,
+ agents: PolarsIdsLike,
+ ) -> bool | pl.Series:
+ if isinstance(agents, pl.Series):
+ return agents.is_in(self._agents["unique_id"])
+ elif isinstance(agents, Collection):
+ return pl.Series(agents).is_in(self._agents["unique_id"])
+ else:
+ return agents in self._agents["unique_id"]
+
+
+
+[docs]
+ def get(
+ self,
+ attr_names: IntoExpr | Iterable[IntoExpr] | None,
+ mask: AgentPolarsMask = None,
+ ) -> pl.Series | pl.DataFrame:
+ masked_df = self._get_masked_df(mask)
+ attr_names = self.agents.select(attr_names).columns.copy()
+ if not attr_names:
+ return masked_df
+ masked_df = masked_df.select(attr_names)
+ if masked_df.shape[1] == 1:
+ return masked_df[masked_df.columns[0]]
+ return masked_df
+
+
+
+[docs]
+ def set(
+ self,
+ attr_names: str | Collection[str] | dict[str, Any] | None = None,
+ values: Any | None = None,
+ mask: AgentPolarsMask = None,
+ inplace: bool = True,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ b_mask = obj._get_bool_mask(mask)
+ masked_df = obj._get_masked_df(mask)
+
+ if not attr_names:
+ attr_names = masked_df.columns
+ attr_names.remove("unique_id")
+
+ def process_single_attr(
+ masked_df: pl.DataFrame, attr_name: str, values: Any
+ ) -> pl.DataFrame:
+ if isinstance(values, pl.DataFrame):
+ return masked_df.with_columns(values.to_series().alias(attr_name))
+ elif isinstance(values, pl.Expr):
+ return masked_df.with_columns(values.alias(attr_name))
+ if isinstance(values, pl.Series):
+ return masked_df.with_columns(values.alias(attr_name))
+ else:
+ if isinstance(values, Collection):
+ values = pl.Series(values)
+ else:
+ values = pl.repeat(values, len(masked_df))
+ return masked_df.with_columns(values.alias(attr_name))
+
+ if isinstance(attr_names, str) and values is not None:
+ masked_df = process_single_attr(masked_df, attr_names, values)
+ elif isinstance(attr_names, Collection) and values is not None:
+ if isinstance(values, Collection) and len(attr_names) == len(values):
+ for attribute, val in zip(attr_names, values):
+ masked_df = process_single_attr(masked_df, attribute, val)
+ else:
+ for attribute in attr_names:
+ masked_df = process_single_attr(masked_df, attribute, values)
+ elif isinstance(attr_names, dict):
+ for key, val in attr_names.items():
+ masked_df = process_single_attr(masked_df, key, val)
+ else:
+ raise ValueError(
+ "attr_names must be a string, a collection of string or a dictionary with columns as keys and values."
+ )
+ non_masked_df = obj._agents.filter(b_mask.not_())
+ original_index = obj._agents.select("unique_id")
+ obj._agents = pl.concat([non_masked_df, masked_df], how="diagonal_relaxed")
+ obj._agents = original_index.join(obj._agents, on="unique_id", how="left")
+ return obj
+
+
+
+[docs]
+ def select(
+ self,
+ mask: AgentPolarsMask = None,
+ filter_func: Callable[[Self], pl.Series] | None = None,
+ n: int | None = None,
+ negate: bool = False,
+ inplace: bool = True,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ mask = obj._get_bool_mask(mask)
+ if filter_func:
+ mask = mask & filter_func(obj)
+ if n is not None:
+ mask = (obj._agents["unique_id"]).is_in(
+ obj._agents.filter(mask).sample(n)["unique_id"]
+ )
+ if negate:
+ mask = mask.not_()
+ obj._mask = mask
+ return obj
+
+
+
+[docs]
+ def shuffle(self, inplace: bool = True) -> Self:
+ obj = self._get_obj(inplace)
+ obj._agents = obj._agents.sample(fraction=1, shuffle=True)
+ return obj
+
+
+
+[docs]
+ def sort(
+ self,
+ by: str | Sequence[str],
+ ascending: bool | Sequence[bool] = True,
+ inplace: bool = True,
+ **kwargs,
+ ) -> Self:
+ obj = self._get_obj(inplace)
+ if isinstance(ascending, bool):
+ descending = not ascending
+ else:
+ descending = [not a for a in ascending]
+ obj._agents = obj._agents.sort(by=by, descending=descending, **kwargs)
+ return obj
+
+
+ def to_pandas(self) -> "AgentSetPandas":
+ from mesa_frames.concrete.pandas.agentset import AgentSetPandas
+
+ new_obj = AgentSetPandas(self._model)
+ new_obj._agents = self._agents.to_pandas()
+ if isinstance(self._mask, pl.Series):
+ new_obj._mask = self._mask.to_pandas()
+ else: # self._mask is Expr
+ new_obj._mask = (
+ self._agents["unique_id"]
+ .is_in(self._agents.filter(self._mask)["unique_id"])
+ .to_pandas()
+ )
+ return new_obj
+
+ def _concatenate_agentsets(
+ self,
+ agentsets: Iterable[Self],
+ duplicates_allowed: bool = True,
+ keep_first_only: bool = True,
+ original_masked_index: pl.Series | None = None,
+ ) -> Self:
+ if not duplicates_allowed:
+ indices_list = [self._agents["unique_id"]] + [
+ agentset._agents["unique_id"] for agentset in agentsets
+ ]
+ all_indices = pl.concat(indices_list)
+ if all_indices.is_duplicated().any():
+ raise ValueError(
+ "Some ids are duplicated in the AgentSetDFs that are trying to be concatenated"
+ )
+ if duplicates_allowed & keep_first_only:
+ # Find the original_index list (ie longest index list), to sort correctly the rows after concatenation
+ max_length = max(len(agentset) for agentset in agentsets)
+ for agentset in agentsets:
+ if len(agentset) == max_length:
+ original_index = agentset._agents["unique_id"]
+ final_dfs = [self._agents]
+ final_active_indices = [self._agents["unique_id"]]
+ final_indices = self._agents["unique_id"].clone()
+ for obj in iter(agentsets):
+ # Remove agents that are already in the final DataFrame
+ final_dfs.append(
+ obj._agents.filter(pl.col("unique_id").is_in(final_indices).not_())
+ )
+ # Add the indices of the active agents of current AgentSet
+ final_active_indices.append(obj._agents.filter(obj._mask)["unique_id"])
+ # Update the indices of the agents in the final DataFrame
+ final_indices = pl.concat(
+ [final_indices, final_dfs[-1]["unique_id"]], how="vertical"
+ )
+ # Left-join original index with concatenated dfs to keep original ids order
+ final_df = original_index.to_frame().join(
+ pl.concat(final_dfs, how="diagonal_relaxed"), on="unique_id", how="left"
+ )
+ #
+ final_active_index = pl.concat(final_active_indices, how="vertical")
+
+ else:
+ final_df = pl.concat(
+ [obj._agents for obj in agentsets], how="diagonal_relaxed"
+ )
+ final_active_index = pl.concat(
+ [obj._agents.filter(obj._mask)["unique_id"] for obj in agentsets]
+ )
+ final_mask = final_df["unique_id"].is_in(final_active_index)
+ self._agents = final_df
+ self._mask = final_mask
+ # If some ids were removed in the do-method, we need to remove them also from final_df
+ if not isinstance(original_masked_index, type(None)):
+ ids_to_remove = original_masked_index.filter(
+ original_masked_index.is_in(self._agents["unique_id"]).not_()
+ )
+ if not ids_to_remove.is_empty():
+ self.remove(ids_to_remove, inplace=True)
+ return self
+
+ def _get_bool_mask(
+ self,
+ mask: AgentPolarsMask = None,
+ ) -> pl.Series | pl.Expr:
+ def bool_mask_from_series(mask: pl.Series) -> pl.Series:
+ if (
+ isinstance(mask, pl.Series)
+ and mask.dtype == pl.Boolean
+ and len(mask) == len(self._agents)
+ ):
+ return mask
+ return self._agents["unique_id"].is_in(mask)
+
+ if isinstance(mask, pl.Expr):
+ return mask
+ elif isinstance(mask, pl.Series):
+ return bool_mask_from_series(mask)
+ elif isinstance(mask, pl.DataFrame):
+ if "unique_id" in mask.columns:
+ return bool_mask_from_series(mask["unique_id"])
+ elif len(mask.columns) == 1 and mask.dtypes[0] == pl.Boolean:
+ return bool_mask_from_series(mask[mask.columns[0]])
+ else:
+ raise KeyError(
+ "DataFrame must have a 'unique_id' column or a single boolean column."
+ )
+ elif mask is None or mask == "all":
+ return pl.repeat(True, len(self._agents))
+ elif mask == "active":
+ return self._mask
+ elif isinstance(mask, Collection):
+ return bool_mask_from_series(pl.Series(mask))
+ else:
+ return bool_mask_from_series(pl.Series([mask]))
+
+ def _get_masked_df(
+ self,
+ mask: AgentPolarsMask = None,
+ ) -> pl.DataFrame:
+ if (isinstance(mask, pl.Series) and mask.dtype == pl.Boolean) or isinstance(
+ mask, pl.Expr
+ ):
+ return self._agents.filter(mask)
+ elif isinstance(mask, pl.DataFrame):
+ if not mask["unique_id"].is_in(self._agents["unique_id"]).all():
+ raise KeyError(
+ "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
+ )
+ return mask.select("unique_id").join(
+ self._agents, on="unique_id", how="left"
+ )
+ elif isinstance(mask, pl.Series):
+ if not mask.is_in(self._agents["unique_id"]).all():
+ raise KeyError(
+ "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
+ )
+ mask_df = mask.to_frame("unique_id")
+ return mask_df.join(self._agents, on="unique_id", how="left")
+ elif mask is None or mask == "all":
+ return self._agents
+ elif mask == "active":
+ return self._agents.filter(self._mask)
+ else:
+ if isinstance(mask, Collection):
+ mask_series = pl.Series(mask)
+ else:
+ mask_series = pl.Series([mask])
+ if not mask_series.is_in(self._agents["unique_id"]).all():
+ raise KeyError(
+ "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
+ )
+ mask_df = mask_series.to_frame("unique_id")
+ return mask_df.join(self._agents, on="unique_id", how="left")
+
+ @overload
+ def _get_obj_copy(self, obj: pl.Series) -> pl.Series: ...
+
+ @overload
+ def _get_obj_copy(self, obj: pl.DataFrame) -> pl.DataFrame: ...
+
+ def _get_obj_copy(self, obj: pl.Series | pl.DataFrame) -> pl.Series | pl.DataFrame:
+ return obj.clone()
+
+ def _discard(self, ids: PolarsIdsLike) -> Self:
+ mask = self._get_bool_mask(ids)
+
+ if isinstance(self._mask, pl.Series):
+ original_active_indices = self._agents.filter(self._mask)["unique_id"]
+
+ self._agents = self._agents.filter(mask.not_())
+
+ if isinstance(self._mask, pl.Series):
+ self._update_mask(original_active_indices)
+
+ return self
+
+ def _update_mask(
+ self, original_active_indices: pl.Series, new_indices: pl.Series | None = None
+ ) -> None:
+ if new_indices is not None:
+ self._mask = self._agents["unique_id"].is_in(
+ original_active_indices
+ ) | self._agents["unique_id"].is_in(new_indices)
+ else:
+ self._mask = self._agents["unique_id"].is_in(original_active_indices)
+
+
+[docs]
+ def __getattr__(self, key: str) -> pl.Series:
+ super().__getattr__(key)
+ return self._agents[key]
+
+
+ @overload
+ def __getitem__(
+ self,
+ key: str | tuple[AgentPolarsMask, str],
+ ) -> pl.Series: ...
+
+ @overload
+ def __getitem__(
+ self,
+ key: (
+ AgentPolarsMask
+ | Collection[str]
+ | tuple[
+ AgentPolarsMask,
+ Collection[str],
+ ]
+ ),
+ ) -> pl.DataFrame: ...
+
+
+[docs]
+ def __getitem__(
+ self,
+ key: (
+ str
+ | Collection[str]
+ | AgentPolarsMask
+ | tuple[AgentPolarsMask, str]
+ | tuple[
+ AgentPolarsMask,
+ Collection[str],
+ ]
+ ),
+ ) -> pl.Series | pl.DataFrame:
+ attr = super().__getitem__(key)
+ assert isinstance(attr, (pl.Series, pl.DataFrame))
+ return attr
+
+
+
+[docs]
+ def __iter__(self) -> Iterator[dict[str, Any]]:
+ return iter(self._agents.iter_rows(named=True))
+
+
+
+
+
+
+[docs]
+ def __reversed__(self) -> Iterator:
+ return reversed(iter(self._agents.iter_rows(named=True)))
+
+
+ @property
+ def agents(self) -> pl.DataFrame:
+ return self._agents
+
+ @agents.setter
+ def agents(self, agents: pl.DataFrame) -> None:
+ if "unique_id" not in agents.columns:
+ raise KeyError("DataFrame must have a unique_id column.")
+ self._agents = agents
+
+ @property
+ def active_agents(self) -> pl.DataFrame:
+ return self.agents.filter(self._mask)
+
+ @active_agents.setter
+ def active_agents(self, mask: AgentPolarsMask) -> None:
+ self.select(mask=mask, inplace=True)
+
+ @property
+ def inactive_agents(self) -> pl.DataFrame:
+ return self.agents.filter(~self._mask)
+
+ @property
+ def index(self) -> pl.Series:
+ return self._agents["unique_id"]
+
+ @property
+ def pos(self) -> pl.DataFrame:
+ return super().pos
+
+
+"""
+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 mesa_frames.abstract.space import GridDF
+from mesa_frames.concrete.polars.mixin import PolarsMixin
+from mesa_frames.utils import copydoc
+
+
+
+[docs]
+@copydoc(GridDF)
+class GridPolars(GridDF, PolarsMixin):
+ """Polars-based implementation of GridDF."""
+
+ _agents: pl.DataFrame
+ _copy_with_method: dict[str, tuple[str, list[str]]] = {
+ "_agents": ("clone", []),
+ "_cells": ("clone", []),
+ "_cells_capacity": ("copy", []),
+ "_offsets": ("clone", []),
+ }
+ _cells: pl.DataFrame
+ _cells_capacity: np.ndarray
+ _offsets: pl.DataFrame
+
+ def _empty_cell_condition(self, cap: np.ndarray) -> np.ndarray:
+ # Create a boolean mask of the same shape as cap
+ empty_mask = np.ones_like(cap, dtype=bool)
+
+ if not self._agents.is_empty():
+ # Get the coordinates of all agents
+ agent_coords = self._agents[self._pos_col_names].to_numpy()
+
+ # Mark cells containing agents as not empty
+ empty_mask[tuple(agent_coords.T)] = False
+
+ return empty_mask
+
+ def _generate_empty_grid(
+ self, dimensions: Sequence[int], capacity: int
+ ) -> np.ndarray:
+ if not capacity:
+ capacity = np.inf
+ return np.full(dimensions, capacity)
+
+ def _sample_cells(
+ self,
+ n: int | None,
+ with_replacement: bool,
+ condition: Callable[[np.ndarray], np.ndarray],
+ respect_capacity: bool = True,
+ ) -> pl.DataFrame:
+ # Get the coordinates of cells that meet the condition
+ coords = np.array(np.where(condition(self._cells_capacity))).T
+
+ if respect_capacity and condition != self._full_cell_condition:
+ capacities = self._cells_capacity[tuple(coords.T)]
+ else:
+ # If not respecting capacity or for full cells, set capacities to 1
+ capacities = np.ones(len(coords), dtype=int)
+
+ if n is not None:
+ if with_replacement:
+ if respect_capacity and condition != self._full_cell_condition:
+ assert (
+ n <= capacities.sum()
+ ), "Requested sample size exceeds the total available capacity."
+
+ sampled_coords = np.empty((0, coords.shape[1]), dtype=coords.dtype)
+ while len(sampled_coords) < n:
+ remaining_samples = n - len(sampled_coords)
+ sampled_indices = self.random.choice(
+ len(coords),
+ size=remaining_samples,
+ replace=True,
+ )
+ unique_indices, counts = np.unique(
+ sampled_indices, return_counts=True
+ )
+
+ if respect_capacity and condition != self._full_cell_condition:
+ # Calculate valid counts for each unique index
+ valid_counts = np.minimum(counts, capacities[unique_indices])
+ # Update capacities
+ capacities[unique_indices] -= valid_counts
+ else:
+ valid_counts = counts
+
+ # Create array of repeated coordinates
+ new_coords = np.repeat(coords[unique_indices], valid_counts, axis=0)
+ # Extend sampled_coords
+ sampled_coords = np.vstack((sampled_coords, new_coords))
+
+ if respect_capacity and condition != self._full_cell_condition:
+ # Update coords and capacities
+ mask = capacities > 0
+ coords = coords[mask]
+ capacities = capacities[mask]
+
+ sampled_coords = sampled_coords[:n]
+ self.random.shuffle(sampled_coords)
+ else:
+ assert n <= len(
+ coords
+ ), "Requested sample size exceeds the number of available cells."
+ sampled_indices = self.random.choice(len(coords), size=n, replace=False)
+ sampled_coords = coords[sampled_indices]
+ else:
+ sampled_coords = coords
+
+ # Convert the coordinates to a DataFrame
+ sampled_cells = pl.DataFrame(
+ sampled_coords, schema=self._pos_col_names, orient="row"
+ )
+ return sampled_cells
+
+ def _update_capacity_agents(
+ self,
+ agents: pl.DataFrame,
+ operation: Literal["movement", "removal"],
+ ) -> np.ndarray:
+ # Update capacity for agents that were already on the grid
+ masked_df = self._df_get_masked_df(
+ self._agents, index_cols="agent_id", mask=agents
+ )
+
+ if operation == "movement":
+ # Increase capacity at old positions
+ old_positions = tuple(masked_df[self._pos_col_names].to_numpy().T)
+ np.add.at(self._cells_capacity, old_positions, 1)
+
+ # Decrease capacity at new positions
+ new_positions = tuple(agents[self._pos_col_names].to_numpy().T)
+ np.add.at(self._cells_capacity, new_positions, -1)
+ elif operation == "removal":
+ # Increase capacity at the positions of removed agents
+ positions = tuple(masked_df[self._pos_col_names].to_numpy().T)
+ np.add.at(self._cells_capacity, positions, 1)
+ return self._cells_capacity
+
+ def _update_capacity_cells(self, cells: pl.DataFrame) -> np.ndarray:
+ # Get the coordinates of the cells to update
+ coords = cells[self._pos_col_names]
+
+ # Get the current capacity of updatable cells
+ current_capacity = (
+ coords.join(self._cells, on=self._pos_col_names, how="left")
+ .fill_null(self._capacity)["capacity"]
+ .to_numpy()
+ )
+
+ # Calculate the number of agents currently in each cell
+ agents_in_cells = (
+ current_capacity - self._cells_capacity[tuple(zip(*coords.to_numpy()))]
+ )
+
+ # Update the capacity in self._cells_capacity
+ new_capacity = cells["capacity"].to_numpy() - agents_in_cells
+
+ # Assert that no new capacity is negative
+ assert np.all(
+ new_capacity >= 0
+ ), "New capacity of a cell cannot be less than the number of agents in it."
+
+ self._cells_capacity[tuple(zip(*coords.to_numpy()))] = new_capacity
+
+ return self._cells_capacity
+
+ @property
+ def remaining_capacity(self) -> int:
+ if not self._capacity:
+ return np.inf
+ return self._cells_capacity.sum()
+
+Short
+ */ + .o-tooltip--left { + position: relative; + } + + .o-tooltip--left:after { + opacity: 0; + visibility: hidden; + position: absolute; + content: attr(data-tooltip); + padding: .2em; + font-size: .8em; + left: -.2em; + background: grey; + color: white; + white-space: nowrap; + z-index: 2; + border-radius: 2px; + transform: translateX(-102%) translateY(0); + transition: opacity 0.2s cubic-bezier(0.64, 0.09, 0.08, 1), transform 0.2s cubic-bezier(0.64, 0.09, 0.08, 1); +} + +.o-tooltip--left:hover:after { + display: block; + opacity: 1; + visibility: visible; + transform: translateX(-100%) translateY(0); + transition: opacity 0.2s cubic-bezier(0.64, 0.09, 0.08, 1), transform 0.2s cubic-bezier(0.64, 0.09, 0.08, 1); + transition-delay: .5s; +} + +/* By default the copy button shouldn't show up when printing a page */ +@media print { + button.copybtn { + display: none; + } +} diff --git a/api/_static/copybutton.js b/api/_static/copybutton.js new file mode 100644 index 0000000..19f1fbd --- /dev/null +++ b/api/_static/copybutton.js @@ -0,0 +1,248 @@ +// Localization support +const messages = { + 'en': { + 'copy': 'Copy', + 'copy_to_clipboard': 'Copy to clipboard', + 'copy_success': 'Copied!', + 'copy_failure': 'Failed to copy', + }, + 'es' : { + 'copy': 'Copiar', + 'copy_to_clipboard': 'Copiar al portapapeles', + 'copy_success': '¡Copiado!', + 'copy_failure': 'Error al copiar', + }, + 'de' : { + 'copy': 'Kopieren', + 'copy_to_clipboard': 'In die Zwischenablage kopieren', + 'copy_success': 'Kopiert!', + 'copy_failure': 'Fehler beim Kopieren', + }, + 'fr' : { + 'copy': 'Copier', + 'copy_to_clipboard': 'Copier dans le presse-papier', + 'copy_success': 'Copié !', + 'copy_failure': 'Échec de la copie', + }, + 'ru': { + 'copy': 'Скопировать', + 'copy_to_clipboard': 'Скопировать в буфер', + 'copy_success': 'Скопировано!', + 'copy_failure': 'Не удалось скопировать', + }, + 'zh-CN': { + 'copy': '复制', + 'copy_to_clipboard': '复制到剪贴板', + 'copy_success': '复制成功!', + 'copy_failure': '复制失败', + }, + 'it' : { + 'copy': 'Copiare', + 'copy_to_clipboard': 'Copiato negli appunti', + 'copy_success': 'Copiato!', + 'copy_failure': 'Errore durante la copia', + } +} + +let locale = 'en' +if( document.documentElement.lang !== undefined + && messages[document.documentElement.lang] !== undefined ) { + locale = document.documentElement.lang +} + +let doc_url_root = DOCUMENTATION_OPTIONS.URL_ROOT; +if (doc_url_root == '#') { + doc_url_root = ''; +} + +/** + * SVG files for our copy buttons + */ +let iconCheck = `` + +// If the user specified their own SVG use that, otherwise use the default +let iconCopy = ``; +if (!iconCopy) { + iconCopy = `` +} + +/** + * Set up copy/paste for code blocks + */ + +const runWhenDOMLoaded = cb => { + if (document.readyState != 'loading') { + cb() + } else if (document.addEventListener) { + document.addEventListener('DOMContentLoaded', cb) + } else { + document.attachEvent('onreadystatechange', function() { + if (document.readyState == 'complete') cb() + }) + } +} + +const codeCellId = index => `codecell${index}` + +// Clears selected text since ClipboardJS will select the text when copying +const clearSelection = () => { + if (window.getSelection) { + window.getSelection().removeAllRanges() + } else if (document.selection) { + document.selection.empty() + } +} + +// Changes tooltip text for a moment, then changes it back +// We want the timeout of our `success` class to be a bit shorter than the +// tooltip and icon change, so that we can hide the icon before changing back. +var timeoutIcon = 2000; +var timeoutSuccessClass = 1500; + +const temporarilyChangeTooltip = (el, oldText, newText) => { + el.setAttribute('data-tooltip', newText) + el.classList.add('success') + // Remove success a little bit sooner than we change the tooltip + // So that we can use CSS to hide the copybutton first + setTimeout(() => el.classList.remove('success'), timeoutSuccessClass) + setTimeout(() => el.setAttribute('data-tooltip', oldText), timeoutIcon) +} + +// Changes the copy button icon for two seconds, then changes it back +const temporarilyChangeIcon = (el) => { + el.innerHTML = iconCheck; + setTimeout(() => {el.innerHTML = iconCopy}, timeoutIcon) +} + +const addCopyButtonToCodeCells = () => { + // If ClipboardJS hasn't loaded, wait a bit and try again. This + // happens because we load ClipboardJS asynchronously. + if (window.ClipboardJS === undefined) { + setTimeout(addCopyButtonToCodeCells, 250) + return + } + + // Add copybuttons to all of our code cells + const COPYBUTTON_SELECTOR = 'div.highlight pre'; + const codeCells = document.querySelectorAll(COPYBUTTON_SELECTOR) + codeCells.forEach((codeCell, index) => { + const id = codeCellId(index) + codeCell.setAttribute('id', id) + + const clipboardButton = id => + `` + codeCell.insertAdjacentHTML('afterend', clipboardButton(id)) + }) + +function escapeRegExp(string) { + return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string +} + +/** + * Removes excluded text from a Node. + * + * @param {Node} target Node to filter. + * @param {string} exclude CSS selector of nodes to exclude. + * @returns {DOMString} Text from `target` with text removed. + */ +function filterText(target, exclude) { + const clone = target.cloneNode(true); // clone as to not modify the live DOM + if (exclude) { + // remove excluded nodes + clone.querySelectorAll(exclude).forEach(node => node.remove()); + } + return clone.innerText; +} + +// Callback when a copy button is clicked. Will be passed the node that was clicked +// should then grab the text and replace pieces of text that shouldn't be used in output +function formatCopyText(textContent, copybuttonPromptText, isRegexp = false, onlyCopyPromptLines = true, removePrompts = true, copyEmptyLines = true, lineContinuationChar = "", hereDocDelim = "") { + var regexp; + var match; + + // Do we check for line continuation characters and "HERE-documents"? + var useLineCont = !!lineContinuationChar + var useHereDoc = !!hereDocDelim + + // create regexp to capture prompt and remaining line + if (isRegexp) { + regexp = new RegExp('^(' + copybuttonPromptText + ')(.*)') + } else { + regexp = new RegExp('^(' + escapeRegExp(copybuttonPromptText) + ')(.*)') + } + + const outputLines = []; + var promptFound = false; + var gotLineCont = false; + var gotHereDoc = false; + const lineGotPrompt = []; + for (const line of textContent.split('\n')) { + match = line.match(regexp) + if (match || gotLineCont || gotHereDoc) { + promptFound = regexp.test(line) + lineGotPrompt.push(promptFound) + if (removePrompts && promptFound) { + outputLines.push(match[2]) + } else { + outputLines.push(line) + } + gotLineCont = line.endsWith(lineContinuationChar) & useLineCont + if (line.includes(hereDocDelim) & useHereDoc) + gotHereDoc = !gotHereDoc + } else if (!onlyCopyPromptLines) { + outputLines.push(line) + } else if (copyEmptyLines && line.trim() === '') { + outputLines.push(line) + } + } + + // If no lines with the prompt were found then just use original lines + if (lineGotPrompt.some(v => v === true)) { + textContent = outputLines.join('\n'); + } + + // Remove a trailing newline to avoid auto-running when pasting + if (textContent.endsWith("\n")) { + textContent = textContent.slice(0, -1) + } + return textContent +} + + +var copyTargetText = (trigger) => { + var target = document.querySelector(trigger.attributes['data-clipboard-target'].value); + + // get filtered text + let exclude = '.linenos'; + + let text = filterText(target, exclude); + return formatCopyText(text, '>>> |\\.\\.\\. ', true, true, true, true, '', '') +} + + // Initialize with a callback so we can modify the text before copy + const clipboard = new ClipboardJS('.copybtn', {text: copyTargetText}) + + // Update UI with error/success messages + clipboard.on('success', event => { + clearSelection() + temporarilyChangeTooltip(event.trigger, messages[locale]['copy'], messages[locale]['copy_success']) + temporarilyChangeIcon(event.trigger) + }) + + clipboard.on('error', event => { + temporarilyChangeTooltip(event.trigger, messages[locale]['copy'], messages[locale]['copy_failure']) + }) +} + +runWhenDOMLoaded(addCopyButtonToCodeCells) \ No newline at end of file diff --git a/api/_static/copybutton_funcs.js b/api/_static/copybutton_funcs.js new file mode 100644 index 0000000..dbe1aaa --- /dev/null +++ b/api/_static/copybutton_funcs.js @@ -0,0 +1,73 @@ +function escapeRegExp(string) { + return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string +} + +/** + * Removes excluded text from a Node. + * + * @param {Node} target Node to filter. + * @param {string} exclude CSS selector of nodes to exclude. + * @returns {DOMString} Text from `target` with text removed. + */ +export function filterText(target, exclude) { + const clone = target.cloneNode(true); // clone as to not modify the live DOM + if (exclude) { + // remove excluded nodes + clone.querySelectorAll(exclude).forEach(node => node.remove()); + } + return clone.innerText; +} + +// Callback when a copy button is clicked. Will be passed the node that was clicked +// should then grab the text and replace pieces of text that shouldn't be used in output +export function formatCopyText(textContent, copybuttonPromptText, isRegexp = false, onlyCopyPromptLines = true, removePrompts = true, copyEmptyLines = true, lineContinuationChar = "", hereDocDelim = "") { + var regexp; + var match; + + // Do we check for line continuation characters and "HERE-documents"? + var useLineCont = !!lineContinuationChar + var useHereDoc = !!hereDocDelim + + // create regexp to capture prompt and remaining line + if (isRegexp) { + regexp = new RegExp('^(' + copybuttonPromptText + ')(.*)') + } else { + regexp = new RegExp('^(' + escapeRegExp(copybuttonPromptText) + ')(.*)') + } + + const outputLines = []; + var promptFound = false; + var gotLineCont = false; + var gotHereDoc = false; + const lineGotPrompt = []; + for (const line of textContent.split('\n')) { + match = line.match(regexp) + if (match || gotLineCont || gotHereDoc) { + promptFound = regexp.test(line) + lineGotPrompt.push(promptFound) + if (removePrompts && promptFound) { + outputLines.push(match[2]) + } else { + outputLines.push(line) + } + gotLineCont = line.endsWith(lineContinuationChar) & useLineCont + if (line.includes(hereDocDelim) & useHereDoc) + gotHereDoc = !gotHereDoc + } else if (!onlyCopyPromptLines) { + outputLines.push(line) + } else if (copyEmptyLines && line.trim() === '') { + outputLines.push(line) + } + } + + // If no lines with the prompt were found then just use original lines + if (lineGotPrompt.some(v => v === true)) { + textContent = outputLines.join('\n'); + } + + // Remove a trailing newline to avoid auto-running when pasting + if (textContent.endsWith("\n")) { + textContent = textContent.slice(0, -1) + } + return textContent +} diff --git a/api/_static/design-tabs.js b/api/_static/design-tabs.js new file mode 100644 index 0000000..b25bd6a --- /dev/null +++ b/api/_static/design-tabs.js @@ -0,0 +1,101 @@ +// @ts-check + +// Extra JS capability for selected tabs to be synced +// The selection is stored in local storage so that it persists across page loads. + +/** + * @type {Record