Mesa 3.0.0a3__py3-none-any.whl → 3.0.0a5__py3-none-any.whl

mesa/datacollection.py

@@ -1,20 +1,19 @@
-"""
-Mesa Data Collection Module
-===========================
+"""Mesa Data Collection Module.
 
 DataCollector is meant to provide a simple, standard way to collect data
-generated by a Mesa model. It collects three types of data: model-level data,
-agent-level data, and tables.
+generated by a Mesa model. It collects four types of data: model-level data,
+agent-level data, agent-type-level data, and tables.
 
-A DataCollector is instantiated with two dictionaries of reporter names and
-associated variable names or functions for each, one for model-level data and
-one for agent-level data; a third dictionary provides table names and columns.
-Variable names are converted into functions which retrieve attributes of that
-name.
+A DataCollector is instantiated with three dictionaries of reporter names and
+associated variable names or functions for each, one for model-level data,
+one for agent-level data, and one for agent-type-level data; a fourth dictionary
+provides table names and columns. Variable names are converted into functions
+which retrieve attributes of that name.
 
 When the collect() method is called, each model-level function is called, with
 the model as the argument, and the results associated with the relevant
-variable. Then the agent-level functions are called on each agent.
+variable. Then the agent-level functions are called on each agent, and the
+agent-type-level functions are called on each agent of the specified type.
 
 Additionally, other objects can write directly to tables by passing in an
 appropriate dictionary object for a table row.
@@ -23,19 +22,18 @@ The DataCollector then stores the data it collects in dictionaries:
     * model_vars maps each reporter to a list of its values
     * tables maps each table to a dictionary, with each column as a key with a
       list as its value.
-    * _agent_records maps each model step to a list of each agents id
+    * _agent_records maps each model step to a list of each agent's id
       and its values.
+    * _agenttype_records maps each model step to a dictionary of agent types,
+      each containing a list of each agent's id and its values.
 
 Finally, DataCollector can create a pandas DataFrame from each collection.
-
-The default DataCollector here makes several assumptions:
-    * The model has an agent list called agents
-    * For collecting agent-level variables, agents must have a unique_id
 """
 
 import contextlib
 import itertools
 import types
+import warnings
 from copy import deepcopy
 from functools import partial
 
@@ -46,24 +44,25 @@ with contextlib.suppress(ImportError):
 class DataCollector:
     """Class for collecting data generated by a Mesa model.
 
-    A DataCollector is instantiated with dictionaries of names of model- and
-    agent-level variables to collect, associated with attribute names or
-    functions which actually collect them. When the collect(...) method is
-    called, it collects these attributes and executes these functions one by
-    one and stores the results.
+    A DataCollector is instantiated with dictionaries of names of model-,
+    agent-, and agent-type-level variables to collect, associated with
+    attribute names or functions which actually collect them. When the
+    collect(...) method is called, it collects these attributes and executes
+    these functions one by one and stores the results.
     """
 
     def __init__(
         self,
         model_reporters=None,
         agent_reporters=None,
+        agenttype_reporters=None,
         tables=None,
     ):
-        """
-        Instantiate a DataCollector with lists of model and agent reporters.
-        Both model_reporters and agent_reporters accept a dictionary mapping a
-        variable name to either an attribute name, a function, a method of a class/instance,
-        or a function with parameters placed in a list.
+        """Instantiate a DataCollector with lists of model, agent, and agent-type reporters.
+
+        Both model_reporters, agent_reporters, and agenttype_reporters accept a
+        dictionary mapping a variable name to either an attribute name, a function,
+        a method of a class/instance, or a function with parameters placed in a list.
 
         Model reporters can take four types of arguments:
         1. Lambda function:
@@ -87,6 +86,10 @@ class DataCollector:
         4. Functions with parameters placed in a list:
            {"Agent_Function": [function, [param_1, param_2]]}
 
+        Agenttype reporters take a dictionary mapping agent types to dictionaries
+        of reporter names and attributes/funcs/methods, similar to agent_reporters:
+           {Wolf: {"energy": lambda a: a.energy}}
+
         The tables arg accepts a dictionary mapping names of tables to lists of
         columns. For example, if we want to allow agents to write their age
         when they are destroyed (to keep track of lifespans), it might look
@@ -96,6 +99,8 @@ class DataCollector:
         Args:
             model_reporters: Dictionary of reporter names and attributes/funcs/methods.
             agent_reporters: Dictionary of reporter names and attributes/funcs/methods.
+            agenttype_reporters: Dictionary of agent types to dictionaries of
+                                 reporter names and attributes/funcs/methods.
             tables: Dictionary of table names to lists of column names.
 
         Notes:
@@ -105,9 +110,11 @@ class DataCollector:
         """
         self.model_reporters = {}
         self.agent_reporters = {}
+        self.agenttype_reporters = {}
 
         self.model_vars = {}
         self._agent_records = {}
+        self._agenttype_records = {}
         self.tables = {}
 
         if model_reporters is not None:
@@ -118,6 +125,11 @@ class DataCollector:
             for name, reporter in agent_reporters.items():
                 self._new_agent_reporter(name, reporter)
 
+        if agenttype_reporters is not None:
+            for agent_type, reporters in agenttype_reporters.items():
+                for name, reporter in reporters.items():
+                    self._new_agenttype_reporter(agent_type, name, reporter)
+
         if tables is not None:
             for name, columns in tables.items():
                 self._new_table(name, columns)
@@ -165,6 +177,38 @@ class DataCollector:
 
         self.agent_reporters[name] = reporter
 
+    def _new_agenttype_reporter(self, agent_type, name, reporter):
+        """Add a new agent-type-level reporter to collect.
+
+        Args:
+            agent_type: The type of agent to collect data for.
+            name: Name of the agent-type-level variable to collect.
+            reporter: Attribute string, function object, method of a class/instance, or
+                      function with parameters placed in a list that returns the
+                      variable when given an agent instance.
+        """
+        if agent_type not in self.agenttype_reporters:
+            self.agenttype_reporters[agent_type] = {}
+
+        # Use the same logic as _new_agent_reporter
+        if isinstance(reporter, str):
+            attribute_name = reporter
+
+            def attr_reporter(agent):
+                return getattr(agent, attribute_name, None)
+
+            reporter = attr_reporter
+
+        elif isinstance(reporter, list):
+            func, params = reporter[0], reporter[1]
+
+            def func_with_params(agent):
+                return func(agent, *params)
+
+            reporter = func_with_params
+
+        self.agenttype_reporters[agent_type][name] = reporter
+
     def _new_table(self, table_name, table_columns):
         """Add a new table that objects can write to.
 
@@ -192,6 +236,34 @@ class DataCollector:
         )
         return agent_records
 
+    def _record_agenttype(self, model, agent_type):
+        """Record agent-type data in a mapping of functions and agents."""
+        rep_funcs = self.agenttype_reporters[agent_type].values()
+
+        def get_reports(agent):
+            _prefix = (agent.model.steps, agent.unique_id)
+            reports = tuple(rep(agent) for rep in rep_funcs)
+            return _prefix + reports
+
+        agent_types = model.agent_types
+        if agent_type in agent_types:
+            agents = model.agents_by_type[agent_type]
+        else:
+            from mesa import Agent
+
+            if issubclass(agent_type, Agent):
+                agents = [
+                    agent for agent in model.agents if isinstance(agent, agent_type)
+                ]
+            else:
+                # Raise error if agent_type is not in model.agent_types
+                raise ValueError(
+                    f"Agent type {agent_type} is not recognized as an Agent type in the model or Agent subclass. Use an Agent (sub)class, like {agent_types}."
+                )
+
+        agenttype_records = map(get_reports, agents)
+        return agenttype_records
+
     def collect(self, model):
         """Collect all the data for the given model object."""
         if self.model_reporters:
@@ -210,7 +282,6 @@ class DataCollector:
                 elif isinstance(reporter, list):
                     self.model_vars[var].append(deepcopy(reporter[0](*reporter[1])))
                 # Assume it's a callable otherwise (e.g., method)
-                # TODO: Check if method of a class explicitly
                 else:
                     self.model_vars[var].append(deepcopy(reporter()))
 
@@ -218,6 +289,14 @@ class DataCollector:
             agent_records = self._record_agents(model)
             self._agent_records[model.steps] = list(agent_records)
 
+        if self.agenttype_reporters:
+            self._agenttype_records[model.steps] = {}
+            for agent_type in self.agenttype_reporters:
+                agenttype_records = self._record_agenttype(model, agent_type)
+                self._agenttype_records[model.steps][agent_type] = list(
+                    agenttype_records
+                )
+
     def add_table_row(self, table_name, row, ignore_missing=False):
         """Add a row dictionary to a specific table.
 
@@ -274,6 +353,38 @@ class DataCollector:
         )
         return df
 
+    def get_agenttype_vars_dataframe(self, agent_type):
+        """Create a pandas DataFrame from the agent-type variables for a specific agent type.
+
+        The DataFrame has one column for each variable, with two additional
+        columns for tick and agent_id.
+
+        Args:
+            agent_type: The type of agent to get the data for.
+        """
+        # Check if self.agenttype_reporters dictionary is empty for this agent type, if so return empty DataFrame
+        if agent_type not in self.agenttype_reporters:
+            warnings.warn(
+                f"No agent-type reporters have been defined for {agent_type} in the DataCollector, returning empty DataFrame.",
+                UserWarning,
+                stacklevel=2,
+            )
+            return pd.DataFrame()
+
+        all_records = itertools.chain.from_iterable(
+            records[agent_type]
+            for records in self._agenttype_records.values()
+            if agent_type in records
+        )
+        rep_names = list(self.agenttype_reporters[agent_type])
+
+        df = pd.DataFrame.from_records(
+            data=all_records,
+            columns=["Step", "AgentID", *rep_names],
+            index=["Step", "AgentID"],
+        )
+        return df
+
     def get_table_dataframe(self, table_name):
         """Create a pandas DataFrame from a particular table.
 

mesa/experimental/UserParam.py

@@ -0,0 +1,67 @@
+"""helper classes."""
+
+
+class UserParam:  # noqa: D101
+    _ERROR_MESSAGE = "Missing or malformed inputs for '{}' Option '{}'"
+
+    def maybe_raise_error(self, param_type, valid):  # noqa: D102
+        if valid:
+            return
+        msg = self._ERROR_MESSAGE.format(param_type, self.label)
+        raise ValueError(msg)
+
+
+class Slider(UserParam):
+    """A number-based slider input with settable increment.
+
+    Example:
+    slider_option = Slider("My Slider", value=123, min=10, max=200, step=0.1)
+
+    Args:
+        label: The displayed label in the UI
+        value: The initial value of the slider
+        min: The minimum possible value of the slider
+        max: The maximum possible value of the slider
+        step: The step between min and max for a range of possible values
+        dtype: either int or float
+    """
+
+    def __init__(
+        self,
+        label="",
+        value=None,
+        min=None,
+        max=None,
+        step=1,
+        dtype=None,
+    ):
+        """Slider class.
+
+        Args:
+            label: The displayed label in the UI
+            value: The initial value of the slider
+            min: The minimum possible value of the slider
+            max: The maximum possible value of the slider
+            step: The step between min and max for a range of possible values
+            dtype: either int or float
+        """
+        self.label = label
+        self.value = value
+        self.min = min
+        self.max = max
+        self.step = step
+
+        # Validate option type to make sure values are supplied properly
+        valid = not (self.value is None or self.min is None or self.max is None)
+        self.maybe_raise_error("slider", valid)
+
+        if dtype is None:
+            self.is_float_slider = self._check_values_are_float(value, min, max, step)
+        else:
+            self.is_float_slider = dtype is float
+
+    def _check_values_are_float(self, value, min, max, step):
+        return any(isinstance(n, float) for n in (value, min, max, step))
+
+    def get(self, attr):  # noqa: D102
+        return getattr(self, attr)

mesa/experimental/__init__.py

@@ -1,3 +1,7 @@
+"""Experimental init."""
+
 from mesa.experimental import cell_space
 
-__all__ = ["cell_space"]
+from .solara_viz import JupyterViz, Slider, SolaraViz, make_text
+
+__all__ = ["cell_space", "JupyterViz", "SolaraViz", "make_text", "Slider"]

mesa/experimental/cell_space/__init__.py

@@ -1,3 +1,10 @@
+"""Cell spaces.
+
+Cell spaces offer an alternative API for discrete spaces. It is experimental and under development. The API is more
+expressive that the default grids available in `mesa.space`.
+
+"""
+
 from mesa.experimental.cell_space.cell import Cell
 from mesa.experimental.cell_space.cell_agent import CellAgent
 from mesa.experimental.cell_space.cell_collection import CellCollection

mesa/experimental/cell_space/cell.py

@@ -1,14 +1,19 @@
+"""The Cell in a cell space."""
+
 from __future__ import annotations
 
-from functools import cache
+from functools import cache, cached_property
 from random import Random
 from typing import TYPE_CHECKING
 
 from mesa.experimental.cell_space.cell_collection import CellCollection
 
 if TYPE_CHECKING:
+    from mesa.agent import Agent
     from mesa.experimental.cell_space.cell_agent import CellAgent
 
+Coordinate = tuple[int, ...]
+
 
 class Cell:
     """The cell represents a position in a discrete space.
@@ -24,11 +29,12 @@ class Cell:
 
     __slots__ = [
         "coordinate",
-        "_connections",
+        "connections",
         "agents",
         "capacity",
         "properties",
         "random",
+        "__dict__",
     ]
 
     # def __new__(cls,
@@ -42,34 +48,39 @@ class Cell:
 
     def __init__(
         self,
-        coordinate: tuple[int, ...],
-        capacity: float | None = None,
+        coordinate: Coordinate,
+        capacity: int | None = None,
         random: Random | None = None,
     ) -> None:
-        """ "
+        """Initialise the cell.
 
         Args:
-            coordinate:
+            coordinate: coordinates of the cell
             capacity (int) : the capacity of the cell. If None, the capacity is infinite
             random (Random) : the random number generator to use
 
         """
         super().__init__()
         self.coordinate = coordinate
-        self._connections: list[Cell] = []  # TODO: change to CellCollection?
-        self.agents = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
-        self.capacity = capacity
-        self.properties: dict[str, object] = {}
+        self.connections: dict[Coordinate, Cell] = {}
+        self.agents: list[
+            Agent
+        ] = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
+        self.capacity: int = capacity
+        self.properties: dict[Coordinate, object] = {}
         self.random = random
 
-    def connect(self, other: Cell) -> None:
+    def connect(self, other: Cell, key: Coordinate | None = None) -> None:
         """Connects this cell to another cell.
 
         Args:
             other (Cell): other cell to connect to
+            key (Tuple[int, ...]): key for the connection. Should resemble a relative coordinate
 
         """
-        self._connections.append(other)
+        if key is None:
+            key = other.coordinate
+        self.connections[key] = other
 
     def disconnect(self, other: Cell) -> None:
         """Disconnects this cell from another cell.
@@ -78,7 +89,9 @@ class Cell:
             other (Cell): other cell to remove from connections
 
         """
-        self._connections.remove(other)
+        keys_to_remove = [k for k, v in self.connections.items() if v == other]
+        for key in keys_to_remove:
+            del self.connections[key]
 
     def add_agent(self, agent: CellAgent) -> None:
         """Adds an agent to the cell.
@@ -116,34 +129,62 @@ class Cell:
         """Returns a bool of the contents of a cell."""
         return len(self.agents) == self.capacity
 
-    def __repr__(self):
+    def __repr__(self):  # noqa
         return f"Cell({self.coordinate}, {self.agents})"
 
+    @cached_property
+    def neighborhood(self) -> CellCollection:
+        """Returns the direct neighborhood of the cell.
+
+        This is equivalent to cell.get_neighborhood(radius=1)
+
+        """
+        return self.get_neighborhood()
+
     # FIXME: Revisit caching strategy on methods
     @cache  # noqa: B019
-    def neighborhood(self, radius=1, include_center=False):
-        return CellCollection(
+    def get_neighborhood(
+        self, radius: int = 1, include_center: bool = False
+    ) -> CellCollection:
+        """Returns a list of all neighboring cells for the given radius.
+
+        For getting the direct neighborhood (i.e., radius=1) you can also use
+        the `neighborhood` property.
+
+        Args:
+            radius (int): the radius of the neighborhood
+            include_center (bool): include the center of the neighborhood
+
+        Returns:
+            a list of all neighboring cells
+
+        """
+        return CellCollection[Cell](
             self._neighborhood(radius=radius, include_center=include_center),
             random=self.random,
         )
 
     # FIXME: Revisit caching strategy on methods
     @cache  # noqa: B019
-    def _neighborhood(self, radius=1, include_center=False):
+    def _neighborhood(
+        self, radius: int = 1, include_center: bool = False
+    ) -> dict[Cell, list[Agent]]:
         # if radius == 0:
         #     return {self: self.agents}
         if radius < 1:
             raise ValueError("radius must be larger than one")
         if radius == 1:
-            neighborhood = {neighbor: neighbor.agents for neighbor in self._connections}
+            neighborhood = {
+                neighbor: neighbor.agents for neighbor in self.connections.values()
+            }
             if not include_center:
                 return neighborhood
             else:
                 neighborhood[self] = self.agents
                 return neighborhood
         else:
-            neighborhood = {}
-            for neighbor in self._connections:
+            neighborhood: dict[Cell, list[Agent]] = {}
+            for neighbor in self.connections.values():
                 neighborhood.update(
                     neighbor._neighborhood(radius - 1, include_center=True)
                 )

mesa/experimental/cell_space/cell_agent.py

@@ -1,3 +1,5 @@
+"""An agent with movement methods for cell spaces."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
@@ -9,8 +11,7 @@ if TYPE_CHECKING:
 
 
 class CellAgent(Agent):
-    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces
-
+    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces.
 
     Attributes:
         unique_id (int): A unique identifier for this agent.
@@ -19,18 +20,22 @@ class CellAgent(Agent):
         cell: (Cell | None): the cell which the agent occupies
     """
 
-    def __init__(self, unique_id: int, model: Model) -> None:
-        """
-        Create a new agent.
+    def __init__(self, model: Model) -> None:
+        """Create a new agent.
 
         Args:
-            unique_id (int): A unique identifier for this agent.
             model (Model): The model instance in which the agent exists.
         """
-        super().__init__(unique_id, model)
+        super().__init__(model)
         self.cell: Cell | None = None
 
     def move_to(self, cell) -> None:
+        """Move agent to cell.
+
+        Args:
+            cell: cell to which agent is to move
+
+        """
         if self.cell is not None:
             self.cell.remove_agent(self)
         self.cell = cell

mesa/experimental/cell_space/cell_collection.py

@@ -1,3 +1,5 @@
+"""CellCollection class."""
+
 from __future__ import annotations
 
 import itertools
@@ -14,7 +16,7 @@ T = TypeVar("T", bound="Cell")
 
 
 class CellCollection(Generic[T]):
-    """An immutable collection of cells
+    """An immutable collection of cells.
 
     Attributes:
         cells (List[Cell]): The list of cells this collection represents
@@ -28,6 +30,12 @@ class CellCollection(Generic[T]):
         cells: Mapping[T, list[CellAgent]] | Iterable[T],
         random: Random | None = None,
     ) -> None:
+        """Initialize a CellCollection.
+
+        Args:
+            cells: cells to add to the collection
+            random: a seeded random number generator.
+        """
         if isinstance(cells, dict):
             self._cells = cells
         else:
@@ -40,42 +48,71 @@ class CellCollection(Generic[T]):
             random = Random()  # FIXME
         self.random = random
 
-    def __iter__(self):
+    def __iter__(self):  # noqa
         return iter(self._cells)
 
-    def __getitem__(self, key: T) -> Iterable[CellAgent]:
+    def __getitem__(self, key: T) -> Iterable[CellAgent]:  # noqa
         return self._cells[key]
 
     # @cached_property
-    def __len__(self) -> int:
+    def __len__(self) -> int:  # noqa
         return len(self._cells)
 
-    def __repr__(self):
+    def __repr__(self):  # noqa
         return f"CellCollection({self._cells})"
 
     @cached_property
-    def cells(self) -> list[T]:
+    def cells(self) -> list[T]:  # noqa
         return list(self._cells.keys())
 
     @property
-    def agents(self) -> Iterable[CellAgent]:
+    def agents(self) -> Iterable[CellAgent]:  # noqa
         return itertools.chain.from_iterable(self._cells.values())
 
     def select_random_cell(self) -> T:
+        """Select a random cell."""
         return self.random.choice(self.cells)
 
     def select_random_agent(self) -> CellAgent:
+        """Select a random agent.
+
+        Returns:
+            CellAgent instance
+
+
+        """
         return self.random.choice(list(self.agents))
 
-    def select(self, filter_func: Callable[[T], bool] | None = None, n=0):
-        # FIXME: n is not considered
-        if filter_func is None and n == 0:
+    def select(
+        self,
+        filter_func: Callable[[T], bool] | None = None,
+        at_most: int | float = float("inf"),
+    ):
+        """Select cells based on filter function.
+
+        Args:
+            filter_func: filter function
+            at_most: The maximum amount of cells to select. Defaults to infinity.
+              - If an integer, at most the first number of matching cells is selected.
+              - If a float between 0 and 1, at most that fraction of original number of cells
+
+        Returns:
+            CellCollection
+
+        """
+        if filter_func is None and at_most == float("inf"):
             return self
 
-        return CellCollection(
-            {
-                cell: agents
-                for cell, agents in self._cells.items()
-                if filter_func is None or filter_func(cell)
-            }
-        )
+        if at_most <= 1.0 and isinstance(at_most, float):
+            at_most = int(len(self) * at_most)  # Note that it rounds down (floor)
+
+        def cell_generator(filter_func, at_most):
+            count = 0
+            for cell in self:
+                if count >= at_most:
+                    break
+                if not filter_func or filter_func(cell):
+                    yield cell
+                    count += 1
+
+        return CellCollection(cell_generator(filter_func, at_most))