Mesa 3.0.0a4__py3-none-any.whl → 3.0.0b0__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

@@ -1,7 +1,10 @@
-class UserParam:
+"""helper classes."""
+
+
+class UserParam:  # noqa: D101
     _ERROR_MESSAGE = "Missing or malformed inputs for '{}' Option '{}'"
 
-    def maybe_raise_error(self, param_type, valid):
+    def maybe_raise_error(self, param_type, valid):  # noqa: D102
         if valid:
             return
         msg = self._ERROR_MESSAGE.format(param_type, self.label)
@@ -9,11 +12,9 @@ class UserParam:
 
 
 class Slider(UserParam):
-    """
-    A number-based slider input with settable increment.
+    """A number-based slider input with settable increment.
 
     Example:
-
     slider_option = Slider("My Slider", value=123, min=10, max=200, step=0.1)
 
     Args:
@@ -34,6 +35,16 @@ class Slider(UserParam):
         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
@@ -52,5 +63,5 @@ class Slider(UserParam):
     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):
+    def get(self, attr):  # noqa: D102
         return getattr(self, attr)

mesa/experimental/__init__.py

@@ -1,3 +1,5 @@
+"""Experimental init."""
+
 from mesa.experimental import cell_space
 
 from .solara_viz import JupyterViz, Slider, SolaraViz, make_text

mesa/experimental/cell_space/__init__.py

@@ -1,5 +1,16 @@
+"""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_agent import (
+    CellAgent,
+    FixedAgent,
+    Grid2DMovingAgent,
+)
 from mesa.experimental.cell_space.cell_collection import CellCollection
 from mesa.experimental.cell_space.discrete_space import DiscreteSpace
 from mesa.experimental.cell_space.grid import (
@@ -15,6 +26,8 @@ __all__ = [
     "CellCollection",
     "Cell",
     "CellAgent",
+    "Grid2DMovingAgent",
+    "FixedAgent",
     "DiscreteSpace",
     "Grid",
     "HexGrid",

mesa/experimental/cell_space/cell.py

@@ -1,13 +1,20 @@
+"""The Cell in a cell space."""
+
 from __future__ import annotations
 
-from functools import cache
+from collections.abc import Callable
+from functools import cache, cached_property
 from random import Random
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from mesa.experimental.cell_space.cell_agent import CellAgent
 from mesa.experimental.cell_space.cell_collection import CellCollection
+from mesa.space import PropertyLayer
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell_agent import CellAgent
+    from mesa.agent import Agent
+
+Coordinate = tuple[int, ...]
 
 
 class Cell:
@@ -24,11 +31,13 @@ class Cell:
 
     __slots__ = [
         "coordinate",
-        "_connections",
+        "connections",
         "agents",
         "capacity",
         "properties",
         "random",
+        "_mesa_property_layers",
+        "__dict__",
     ]
 
     # def __new__(cls,
@@ -42,34 +51,40 @@ 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 | None = capacity
+        self.properties: dict[Coordinate, object] = {}
         self.random = random
+        self._mesa_property_layers: dict[str, PropertyLayer] = {}
 
-    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 +93,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.
@@ -104,7 +121,6 @@ class Cell:
 
         """
         self.agents.remove(agent)
-        agent.cell = None
 
     @property
     def is_empty(self) -> bool:
@@ -116,37 +132,82 @@ 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[Cell]:
+        """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[Cell]:
+        """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)
                 )
             if not include_center:
                 neighborhood.pop(self, None)
             return neighborhood
+
+    # PropertyLayer methods
+    def get_property(self, property_name: str) -> Any:
+        """Get the value of a property."""
+        return self._mesa_property_layers[property_name].data[self.coordinate]
+
+    def set_property(self, property_name: str, value: Any):
+        """Set the value of a property."""
+        self._mesa_property_layers[property_name].set_cell(self.coordinate, value)
+
+    def modify_property(
+        self, property_name: str, operation: Callable, value: Any = None
+    ):
+        """Modify the value of a property."""
+        self._mesa_property_layers[property_name].modify_cell(
+            self.coordinate, operation, value
+        )