Mesa 2.3.4__py3-none-any.whl → 2.4.0__py3-none-any.whl

mesa/__init__.py

@@ -26,7 +26,7 @@ __all__ = [
 ]
 
 __title__ = "mesa"
-__version__ = "2.3.4"
+__version__ = "2.4.0"
 __license__ = "Apache 2.0"
 _this_year = datetime.datetime.now(tz=datetime.timezone.utc).date().year
 __copyright__ = f"Copyright {_this_year} Project Mesa Team"

mesa/agent.py

@@ -14,11 +14,11 @@ import operator
 import warnings
 import weakref
 from collections import defaultdict
-from collections.abc import Iterable, Iterator, MutableSet, Sequence
+from collections.abc import Hashable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
 # mypy
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Literal, overload
 
 if TYPE_CHECKING:
     # We ensure that these are not imported during runtime to prevent cyclic
@@ -49,25 +49,12 @@ class Agent:
         self.model = model
         self.pos: Position | None = None
 
-        # register agent
-        try:
-            self.model.agents_[type(self)][self] = None
-        except AttributeError:
-            # model super has not been called
-            self.model.agents_ = defaultdict(dict)
-            self.model.agents_[type(self)][self] = None
-            self.model.agentset_experimental_warning_given = False
-
-            warnings.warn(
-                "The Mesa Model class was not initialized. In the future, you need to explicitly initialize the Model by calling super().__init__() on initialization.",
-                FutureWarning,
-                stacklevel=2,
-            )
+        self.model.register_agent(self)
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
         with contextlib.suppress(KeyError):
-            self.model.agents_[type(self)].pop(self)
+            self.model.deregister_agent(self)
 
     def step(self) -> None:
         """A single step of the agent."""
@@ -129,9 +116,10 @@ class AgentSet(MutableSet, Sequence):
     def select(
         self,
         filter_func: Callable[[Agent], bool] | None = None,
-        n: int = 0,
+        at_most: int | float = float("inf"),
         inplace: bool = False,
         agent_type: type[Agent] | None = None,
+        n: int | None = None,
     ) -> AgentSet:
         """
         Select a subset of agents from the AgentSet based on a filter function and/or quantity limit.
@@ -139,29 +127,47 @@ class AgentSet(MutableSet, Sequence):
         Args:
             filter_func (Callable[[Agent], bool], optional): A function that takes an Agent and returns True if the
                 agent should be included in the result. Defaults to None, meaning no filtering is applied.
-            n (int, optional): The number of agents to select. If 0, all matching agents are selected. Defaults to 0.
+            at_most (int | float, optional): The maximum amount of agents to select. Defaults to infinity.
+              - If an integer, at most the first number of matching agents are selected.
+              - If a float between 0 and 1, at most that fraction of original the agents are selected.
             inplace (bool, optional): If True, modifies the current AgentSet; otherwise, returns a new AgentSet. Defaults to False.
             agent_type (type[Agent], optional): The class type of the agents to select. Defaults to None, meaning no type filtering is applied.
 
         Returns:
             AgentSet: A new AgentSet containing the selected agents, unless inplace is True, in which case the current AgentSet is updated.
+
+        Notes:
+            - at_most just return the first n or fraction of agents. To take a random sample, shuffle() beforehand.
+            - at_most is an upper limit. When specifying other criteria, the number of agents returned can be smaller.
         """
+        if n is not None:
+            warnings.warn(
+                "The parameter 'n' is deprecated. Use 'at_most' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            at_most = n
 
-        if filter_func is None and agent_type is None and n == 0:
+        inf = float("inf")
+        if filter_func is None and agent_type is None and at_most == inf:
             return self if inplace else copy.copy(self)
 
-        def agent_generator(filter_func=None, agent_type=None, n=0):
+        # Check if at_most is of type float
+        if at_most <= 1.0 and isinstance(at_most, float):
+            at_most = int(len(self) * at_most)  # Note that it rounds down (floor)
+
+        def agent_generator(filter_func, agent_type, at_most):
             count = 0
             for agent in self:
+                if count >= at_most:
+                    break
                 if (not filter_func or filter_func(agent)) and (
                     not agent_type or isinstance(agent, agent_type)
                 ):
                     yield agent
                     count += 1
-                    if 0 < n <= count:
-                        break
 
-        agents = agent_generator(filter_func, agent_type, n)
+        agents = agent_generator(filter_func, agent_type, at_most)
 
         return AgentSet(agents, self.model) if not inplace else self._update(agents)
 
@@ -226,53 +232,179 @@ class AgentSet(MutableSet, Sequence):
         self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
         return self
 
-    def do(
-        self, method_name: str, *args, return_results: bool = False, **kwargs
-    ) -> AgentSet | list[Any]:
+    def do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
         """
-        Invoke a method on each agent in the AgentSet.
+        Invoke a method or function on each agent in the AgentSet.
 
         Args:
-            method_name (str): The name of the method to call on each agent.
-            return_results (bool, optional): If True, returns the results of the method calls; otherwise, returns the AgentSet itself. Defaults to False, so you can chain method calls.
-            *args: Variable length argument list passed to the method being called.
-            **kwargs: Arbitrary keyword arguments passed to the method being called.
+            method (str, callable): the callable to do on each agent
+
+                                        * in case of str, the name of the method to call on each agent.
+                                        * in case of callable, the function to be called with each agent as first argument
+
+            *args: Variable length argument list passed to the callable being called.
+            **kwargs: Arbitrary keyword arguments passed to the callable being called.
 
         Returns:
-            AgentSet | list[Any]: The results of the method calls if return_results is True, otherwise the AgentSet itself.
+            AgentSet | list[Any]: The results of the callable calls if return_results is True, otherwise the AgentSet itself.
         """
+        try:
+            return_results = kwargs.pop("return_results")
+        except KeyError:
+            return_results = False
+        else:
+            warnings.warn(
+                "Using return_results is deprecated. Use AgenSet.do in case of return_results=False, and "
+                "AgentSet.map in case of return_results=True",
+                stacklevel=2,
+            )
+
+        if return_results:
+            return self.map(method, *args, **kwargs)
+
         # we iterate over the actual weakref keys and check if weakref is alive before calling the method
-        res = [
-            getattr(agent, method_name)(*args, **kwargs)
-            for agentref in self._agents.keyrefs()
-            if (agent := agentref()) is not None
-        ]
+        if isinstance(method, str):
+            for agentref in self._agents.keyrefs():
+                if (agent := agentref()) is not None:
+                    getattr(agent, method)(*args, **kwargs)
+        else:
+            for agentref in self._agents.keyrefs():
+                if (agent := agentref()) is not None:
+                    method(agent, *args, **kwargs)
+
+        return self
 
-        return res if return_results else self
+    def shuffle_do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
+        """Shuffle the agents in the AgentSet and then invoke a method or function on each agent.
 
-    def get(self, attr_names: str | list[str]) -> list[Any]:
+        It's a fast, optimized version of calling shuffle() followed by do().
+        """
+        agents = list(self._agents.keys())
+        self.random.shuffle(agents)
+
+        if isinstance(method, str):
+            for agent in agents:
+                getattr(agent, method)(*args, **kwargs)
+        else:
+            for agent in agents:
+                method(agent, *args, **kwargs)
+
+        return self
+
+    def map(self, method: str | Callable, *args, **kwargs) -> list[Any]:
+        """
+        Invoke a method or function on each agent in the AgentSet and return the results.
+
+        Args:
+            method (str, callable): the callable to apply on each agent
+
+                                        * in case of str, the name of the method to call on each agent.
+                                        * in case of callable, the function to be called with each agent as first argument
+
+            *args: Variable length argument list passed to the callable being called.
+            **kwargs: Arbitrary keyword arguments passed to the callable being called.
+
+        Returns:
+           list[Any]: The results of the callable calls
+        """
+        # we iterate over the actual weakref keys and check if weakref is alive before calling the method
+        if isinstance(method, str):
+            res = [
+                getattr(agent, method)(*args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
+        else:
+            res = [
+                method(agent, *args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
+
+        return res
+
+    def agg(self, attribute: str, func: Callable) -> Any:
+        """
+        Aggregate an attribute of all agents in the AgentSet using a specified function.
+
+        Args:
+            attribute (str): The name of the attribute to aggregate.
+            func (Callable): The function to apply to the attribute values (e.g., min, max, sum, np.mean).
+
+        Returns:
+            Any: The result of applying the function to the attribute values. Often a single value.
+        """
+        values = self.get(attribute)
+        return func(values)
+
+    @overload
+    def get(
+        self,
+        attr_names: str,
+        handle_missing: Literal["error", "default"] = "error",
+        default_value: Any = None,
+    ) -> list[Any]: ...
+
+    @overload
+    def get(
+        self,
+        attr_names: list[str],
+        handle_missing: Literal["error", "default"] = "error",
+        default_value: Any = None,
+    ) -> list[list[Any]]: ...
+
+    def get(
+        self,
+        attr_names,
+        handle_missing="error",
+        default_value=None,
+    ):
         """
         Retrieve the specified attribute(s) from each agent in the AgentSet.
 
         Args:
             attr_names (str | list[str]): The name(s) of the attribute(s) to retrieve from each agent.
+            handle_missing (str, optional): How to handle missing attributes. Can be:
+                                            - 'error' (default): raises an AttributeError if attribute is missing.
+                                            - 'default': returns the specified default_value.
+            default_value (Any, optional): The default value to return if 'handle_missing' is set to 'default'
+                                           and the agent does not have the attribute.
 
         Returns:
-            list[Any]: A list with the attribute value for each agent in the set if attr_names is a str
-            list[list[Any]]: A list with a list of attribute values for each agent in the set if attr_names is a list of str
+            list[Any]: A list with the attribute value for each agent if attr_names is a str.
+            list[list[Any]]: A list with a lists of attribute values for each agent if attr_names is a list of str.
 
         Raises:
-            AttributeError if an agent does not have the specified attribute(s)
-
+            AttributeError: If 'handle_missing' is 'error' and the agent does not have the specified attribute(s).
+            ValueError: If an unknown 'handle_missing' option is provided.
         """
+        is_single_attr = isinstance(attr_names, str)
+
+        if handle_missing == "error":
+            if is_single_attr:
+                return [getattr(agent, attr_names) for agent in self._agents]
+            else:
+                return [
+                    [getattr(agent, attr) for attr in attr_names]
+                    for agent in self._agents
+                ]
+
+        elif handle_missing == "default":
+            if is_single_attr:
+                return [
+                    getattr(agent, attr_names, default_value) for agent in self._agents
+                ]
+            else:
+                return [
+                    [getattr(agent, attr, default_value) for attr in attr_names]
+                    for agent in self._agents
+                ]
 
-        if isinstance(attr_names, str):
-            return [getattr(agent, attr_names) for agent in self._agents]
         else:
-            return [
-                [getattr(agent, attr_name) for attr_name in attr_names]
-                for agent in self._agents
-            ]
+            raise ValueError(
+                f"Unknown handle_missing option: {handle_missing}, "
+                "should be one of 'error' or 'default'"
+            )
 
     def set(self, attr_name: str, value: Any) -> AgentSet:
         """
@@ -371,7 +503,139 @@ class AgentSet(MutableSet, Sequence):
         """
         return self.model.random
 
+    def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:
+        """
+        Group agents by the specified attribute or return from the callable
+
+        Args:
+            by (Callable, str): used to determine what to group agents by
+
+                                * if ``by`` is a callable, it will be called for each agent and the return is used
+                                  for grouping
+                                * if ``by`` is a str, it should refer to an attribute on the agent and the value
+                                  of this attribute will be used for grouping
+            result_type (str, optional): The datatype for the resulting groups {"agentset", "list"}
+        Returns:
+            GroupBy
+
+
+        Notes:
+        There might be performance benefits to using `result_type='list'` if you don't need the advanced functionality
+        of an AgentSet.
+
+        """
+        groups = defaultdict(list)
+
+        if isinstance(by, Callable):
+            for agent in self:
+                groups[by(agent)].append(agent)
+        else:
+            for agent in self:
+                groups[getattr(agent, by)].append(agent)
+
+        if result_type == "agentset":
+            return GroupBy(
+                {k: AgentSet(v, model=self.model) for k, v in groups.items()}
+            )
+        else:
+            return GroupBy(groups)
+
+    # consider adding for performance reasons
+    # for Sequence: __reversed__, index, and count
+    # for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
+
+
+class GroupBy:
+    """Helper class for AgentSet.groupby
+
+
+    Attributes:
+        groups (dict): A dictionary with the group_name as key and group as values
+
+    """
+
+    def __init__(self, groups: dict[Any, list | AgentSet]):
+        self.groups: dict[Any, list | AgentSet] = groups
+
+    def map(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
+        """Apply the specified callable to each group and return the results.
+
+        Args:
+            method (Callable, str): The callable to apply to each group,
+
+                                    * if ``method`` is a callable, it will be called it will be called with the group as first argument
+                                    * if ``method`` is a str, it should refer to a method on the group
+
+                                    Additional arguments and keyword arguments will be passed on to the callable.
+
+        Returns:
+            dict with group_name as key and the return of the method as value
+
+        Notes:
+            this method is useful for methods or functions that do return something. It
+            will break method chaining. For that, use ``do`` instead.
+
+        """
+        if isinstance(method, str):
+            return {
+                k: getattr(v, method)(*args, **kwargs) for k, v in self.groups.items()
+            }
+        else:
+            return {k: method(v, *args, **kwargs) for k, v in self.groups.items()}
+
+    def do(self, method: Callable | str, *args, **kwargs) -> GroupBy:
+        """Apply the specified callable to each group
+
+        Args:
+            method (Callable, str): The callable to apply to each group,
+
+                                    * if ``method`` is a callable, it will be called it will be called with the group as first argument
+                                    * if ``method`` is a str, it should refer to a method on the group
+
+                                    Additional arguments and keyword arguments will be passed on to the callable.
+
+        Returns:
+            the original GroupBy instance
+
+        Notes:
+            this method is useful for methods or functions that don't return anything and/or
+            if you want to chain multiple do calls
+
+        """
+        if isinstance(method, str):
+            for v in self.groups.values():
+                getattr(v, method)(*args, **kwargs)
+        else:
+            for v in self.groups.values():
+                method(v, *args, **kwargs)
+
+        return self
+
+    def count(self) -> dict[Any, int]:
+        """Return the count of agents in each group.
+
+        Returns:
+            dict: A dictionary mapping group names to the number of agents in each group.
+        """
+        return {k: len(v) for k, v in self.groups.items()}
+
+    def agg(self, attr_name: str, func: Callable) -> dict[Hashable, Any]:
+        """Aggregate the values of a specific attribute across each group using the provided function.
+
+        Args:
+            attr_name (str): The name of the attribute to aggregate.
+            func (Callable): The function to apply (e.g., sum, min, max, mean).
+
+        Returns:
+            dict[Hashable, Any]: A dictionary mapping group names to the result of applying the aggregation function.
+        """
+        return {
+            group_name: func([getattr(agent, attr_name) for agent in group])
+            for group_name, group in self.groups.items()
+        }
+
+    def __iter__(self):
+        return iter(self.groups.items())
 
-# consider adding for performance reasons
-# for Sequence: __reversed__, index, and count
-# for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
+    def __len__(self):
+        return len(self.groups)

mesa/batchrunner.py

@@ -1,4 +1,5 @@
 import itertools
+import multiprocessing
 from collections.abc import Iterable, Mapping
 from functools import partial
 from multiprocessing import Pool
@@ -8,6 +9,8 @@ from tqdm.auto import tqdm
 
 from mesa.model import Model
 
+multiprocessing.set_start_method("spawn", force=True)
+
 
 def batch_run(
     model_cls: type[Model],

mesa/datacollection.py

@@ -3,28 +3,30 @@ 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.
 
-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.
 
@@ -36,6 +38,7 @@ The default DataCollector here makes several assumptions:
 import contextlib
 import itertools
 import types
+import warnings
 from copy import deepcopy
 from functools import partial
 
@@ -46,24 +49,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 +91,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 +104,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 +115,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 +130,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 +182,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 +241,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:
@@ -218,6 +295,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 +359,36 @@ 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/devs/examples/epstein_civil_violence.py

@@ -261,7 +261,8 @@ class EpsteinCivilViolence(Model):
         self.active_agents = self.agents
 
     def step(self):
-        self.active_agents.shuffle(inplace=True).do("step")
+        """Run one step of the model."""
+        self.active_agents.shuffle_do("step")
 
 
 if __name__ == "__main__":

mesa/experimental/devs/examples/wolf_sheep.py

@@ -231,8 +231,9 @@ class WolfSheep(mesa.Model):
             self.grid.place_agent(patch, pos)
 
     def step(self):
-        self.get_agents_of_type(Sheep).shuffle(inplace=True).do("step")
-        self.get_agents_of_type(Wolf).shuffle(inplace=True).do("step")
+        """Perform one step of the model."""
+        self.agents_by_type[Sheep].shuffle_do("step")
+        self.agents_by_type[Wolf].shuffle_do("step")
 
 
 if __name__ == "__main__":

mesa/experimental/jupyter_viz.py

@@ -143,14 +143,15 @@ def JupyterViz(
             # otherwise, do nothing (do not draw space)
 
             # 5. Plots
-            for measure in measures:
-                if callable(measure):
-                    # Is a custom object
-                    measure(model)
-                else:
-                    components_matplotlib.PlotMatplotlib(
-                        model, measure, dependencies=dependencies
-                    )
+            if measures:
+                for measure in measures:
+                    if callable(measure):
+                        # Is a custom object
+                        measure(model)
+                    else:
+                        components_matplotlib.PlotMatplotlib(
+                            model, measure, dependencies=dependencies
+                        )
 
     def render_in_browser():
         # if space drawer is disabled, do not include it

mesa/model.py

@@ -8,10 +8,8 @@ Core Objects: Model
 # Remove this __future__ import once the oldest supported Python is 3.10
 from __future__ import annotations
 
-import itertools
 import random
 import warnings
-from collections import defaultdict
 
 # mypy
 from typing import Any, Union
@@ -33,20 +31,28 @@ class Model:
         running: A boolean indicating if the model should continue running.
         schedule: An object to manage the order and execution of agent steps.
         current_id: A counter for assigning unique IDs to agents.
-        agents_: A defaultdict mapping each agent type to a dict of its instances.
-                 This private attribute is used internally to manage agents.
 
     Properties:
-        agents: An AgentSet containing all agents in the model, generated from the _agents attribute.
+        agents: An AgentSet containing all agents in the model
         agent_types: A list of different agent types present in the model.
+        agents_by_type: A dictionary where the keys are agent types and the values are the corresponding AgentSets.
 
     Methods:
         get_agents_of_type: Returns an AgentSet of agents of the specified type.
+            Deprecated: Use agents_by_type[agenttype] instead.
         run_model: Runs the model's simulation until a defined end condition is reached.
         step: Executes a single step of the model's simulation process.
         next_id: Generates and returns the next unique identifier for an agent.
         reset_randomizer: Resets the model's random number generator with a new or existing seed.
         initialize_data_collector: Sets up the data collector for the model, requiring an initialized scheduler and agents.
+        register_agent : register an agent with the model
+        deregister_agent : remove an agent from the model
+
+    Notes:
+        Model.agents returns the AgentSet containing all agents registered with the model. Changing
+        the content of the AgentSet directly can result in strange behavior. If you want change the
+        composition of this AgentSet, ensure you operate on a copy.
+
     """
 
     def __new__(cls, *args: Any, **kwargs: Any) -> Any:
@@ -58,6 +64,7 @@ class Model:
             # advance.
             obj._seed = random.random()
         obj.random = random.Random(obj._seed)
+
         # TODO: Remove these 2 lines just before Mesa 3.0
         obj._steps = 0
         obj._time = 0
@@ -72,7 +79,8 @@ class Model:
         self.running = True
         self.schedule = None
         self.current_id = 0
-        self.agents_: defaultdict[type, dict] = defaultdict(dict)
+
+        self._setup_agent_registration()
 
         self._steps: int = 0
         self._time: TimeT = 0  # the model's clock
@@ -80,33 +88,93 @@ class Model:
     @property
     def agents(self) -> AgentSet:
         """Provides an AgentSet of all agents in the model, combining agents from all types."""
-
-        if hasattr(self, "_agents"):
-            return self._agents
-        else:
-            all_agents = itertools.chain.from_iterable(self.agents_.values())
-            return AgentSet(all_agents, self)
+        return self._all_agents
 
     @agents.setter
     def agents(self, agents: Any) -> None:
         warnings.warn(
-            "You are trying to set model.agents. In a next release, this attribute is used "
-            "by MESA itself so you cannot use it directly anymore."
-            "Please adjust your code to use a different attribute name for custom agent storage",
+            "You are trying to set model.agents. In Mesa 3.0 and higher, this attribute is "
+            "used by Mesa itself, so you cannot use it directly anymore."
+            "Please adjust your code to use a different attribute name for custom agent storage.",
             UserWarning,
             stacklevel=2,
         )
 
-        self._agents = agents
-
     @property
     def agent_types(self) -> list[type]:
-        """Return a list of different agent types."""
-        return list(self.agents_.keys())
+        """Return a list of all unique agent types registered with the model."""
+        return list(self._agents_by_type.keys())
+
+    @property
+    def agents_by_type(self) -> dict[type[Agent], AgentSet]:
+        """A dictionary where the keys are agent types and the values are the corresponding AgentSets."""
+        return self._agents_by_type
 
     def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
-        """Retrieves an AgentSet containing all agents of the specified type."""
-        return AgentSet(self.agents_[agenttype].keys(), self)
+        """Deprecated: Retrieves an AgentSet containing all agents of the specified type."""
+        warnings.warn(
+            f"Model.get_agents_of_type() is deprecated, please replace get_agents_of_type({agenttype})"
+            f"with the property agents_by_type[{agenttype}].",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.agents_by_type[agenttype]
+
+    def _setup_agent_registration(self):
+        """helper method to initialize the agent registration datastructures"""
+        self._agents = {}  # the hard references to all agents in the model
+        self._agents_by_type: dict[
+            type[Agent], AgentSet
+        ] = {}  # a dict with an agentset for each class of agents
+        self._all_agents = AgentSet([], self)  # an agenset with all agents
+
+    def register_agent(self, agent):
+        """Register the agent with the model
+
+        Args:
+            agent: The agent to register.
+
+        Notes:
+            This method is called automatically by ``Agent.__init__``, so there is no need to use this
+            if you are subclassing Agent and calling its super in the ``__init__`` method.
+
+        """
+        if not hasattr(self, "_agents"):
+            self._setup_agent_registration()
+
+            warnings.warn(
+                "The Mesa Model class was not initialized. In the future, you need to explicitly initialize "
+                "the Model by calling super().__init__() on initialization.",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        self._agents[agent] = None
+
+        # because AgentSet requires model, we cannot use defaultdict
+        # tricks with a function won't work because model then cannot be pickled
+        try:
+            self._agents_by_type[type(agent)].add(agent)
+        except KeyError:
+            self._agents_by_type[type(agent)] = AgentSet(
+                [
+                    agent,
+                ],
+                self,
+            )
+
+        self._all_agents.add(agent)
+
+    def deregister_agent(self, agent):
+        """Deregister the agent with the model
+
+        Notes::
+        This method is called automatically by ``Agent.remove``
+
+        """
+        del self._agents[agent]
+        self._agents_by_type[type(agent)].remove(agent)
+        self._all_agents.remove(agent)
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as
@@ -144,6 +212,7 @@ class Model:
         self,
         model_reporters=None,
         agent_reporters=None,
+        agenttype_reporters=None,
         tables=None,
     ) -> None:
         if not hasattr(self, "schedule") or self.schedule is None:
@@ -157,6 +226,7 @@ class Model:
         self.datacollector = DataCollector(
             model_reporters=model_reporters,
             agent_reporters=agent_reporters,
+            agenttype_reporters=agenttype_reporters,
             tables=tables,
         )
         # Collect data for the first time during initialization.

mesa/space.py

@@ -459,15 +459,21 @@ class _Grid:
             elif selection == "closest":
                 current_pos = agent.pos
                 # Find the closest position without sorting all positions
-                closest_pos = None
+                # TODO: See if this method can be optimized further
+                closest_pos = []
                 min_distance = float("inf")
                 agent.random.shuffle(pos)
                 for p in pos:
                     distance = self._distance_squared(p, current_pos)
                     if distance < min_distance:
                         min_distance = distance
-                        closest_pos = p
-                chosen_pos = closest_pos
+                        closest_pos.clear()
+                        closest_pos.append(p)
+                    elif distance == min_distance:
+                        closest_pos.append(p)
+
+                chosen_pos = agent.random.choice(closest_pos)
+
             else:
                 raise ValueError(
                     f"Invalid selection method {selection}. Choose 'random' or 'closest'."

{mesa-2.3.4.dist-info → mesa-2.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: Mesa
-Version: 2.3.4
+Version: 2.4.0
 Summary: Agent-based modeling (ABM) in Python
 Project-URL: homepage, https://github.com/projectmesa/mesa
 Project-URL: repository, https://github.com/projectmesa/mesa
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Artificial Life
@@ -55,7 +56,7 @@ Description-Content-Type: text/markdown
 | Meta | [![linting - Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) [![code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) [![Hatch project](https://img.shields.io/badge/%F0%9F%A5%9A-Hatch-4051b5.svg)](https://github.com/pypa/hatch) |
 | Chat | [![chat](https://img.shields.io/matrix/project-mesa:matrix.org?label=chat&logo=Matrix)](https://matrix.to/#/#project-mesa:matrix.org) |
 
-*This is the `2.3.x-maintenance` branch. Example models for Mesa 2.x can be found [here](https://github.com/projectmesa/mesa-examples/tree/mesa-2.x/examples).*
+*This is the `2.4.x-maintenance` branch. Example models for Mesa 2.x can be found [here](https://github.com/projectmesa/mesa-examples/tree/mesa-2.x/examples).*
 
 Mesa allows users to quickly create agent-based models using built-in
 core components (such as spatial grids and agent schedulers) or
@@ -101,13 +102,13 @@ For resources or help on using Mesa, check out the following:
 
 -   [Intro to Mesa Tutorial](http://mesa.readthedocs.org/en/stable/tutorials/intro_tutorial.html) (An introductory model, the Boltzmann
     Wealth Model, for beginners or those new to Mesa.)
--   [Visualization Tutorial](https://mesa.readthedocs.io/en/stable/tutorials/visualization_tutorial.html) (An introduction into our Solara visualization)
+-   [Visualization Tutorial](https://mesa.readthedocs.io/stable/tutorials/visualization_tutorial.html) (An introduction into our Solara visualization)
 -   [Complexity Explorer Tutorial](https://www.complexityexplorer.org/courses/172-agent-based-models-with-python-an-introduction-to-mesa) (An advanced-beginner model,
     SugarScape with Traders, with instructional videos)
 -   [Mesa Examples](https://github.com/projectmesa/mesa-examples/tree/mesa-2.x/examples) (A repository of seminal ABMs using Mesa and
     examples of employing specific Mesa Features)
 -   [Docs](http://mesa.readthedocs.org/) (Mesa's documentation, API and useful snippets)
-    -   [Development version docs](https://mesa.readthedocs.io/en/latest/) (the latest version docs if you're using a pre-release Mesa version)
+    -   [Development version docs](https://mesa.readthedocs.io/latest/) (the latest version docs if you're using a pre-release Mesa version)
 -   [Discussions](https://github.com/projectmesa/mesa/discussions) (GitHub threaded discussions about Mesa)
 -   [Matrix Chat](https://matrix.to/#/#project-mesa:matrix.org) (Chat Forum via Matrix to talk about Mesa)
 

{mesa-2.3.4.dist-info → mesa-2.4.0.dist-info}/RECORD

@@ -1,10 +1,10 @@
-mesa/__init__.py,sha256=fbfWiSzuVuMgF_znRDKBOgMEBo-DGcrl_NS2ZrNWPwY,680
-mesa/agent.py,sha256=db7xGHY0mWgasp-zqC3NgCz_aQoBMHA8DuErK3o7hfk,13398
-mesa/batchrunner.py,sha256=G-sj2g6N6E0BsBikYq5yKsgTNnKHr4ADBkRWa9PXkl8,6055
-mesa/datacollection.py,sha256=wJwt-bOU8WuXZomUg8JJhtIg4KMhoqXS0Zvyjv9SCAE,11445
+mesa/__init__.py,sha256=s9A2haUb53qfSt3a1x5zpqN3aFxdw-s2SlQ5ctmYA3E,680
+mesa/agent.py,sha256=Vz7RGwIuMSiOgI-CEO3N3cnHJRFp5CIHh5u4t-dWH08,23481
+mesa/batchrunner.py,sha256=kyx0-0LjzoMA5vVMWP6nK9hYWK_M1EvKaAN0t_D9ej4,6133
+mesa/datacollection.py,sha256=9yYiONDdNxltQE_6FUKkRgDkBAHFggJM9F_7sKR6CfU,16344
 mesa/main.py,sha256=7MovfNz88VWNnfXP0kcERB6C3GfkVOh0hb0o32hM9LU,1602
-mesa/model.py,sha256=RxTCJUBfEgRIu3dXiMK9oMlxS3owwNQaQIrVRs6HsZY,5823
-mesa/space.py,sha256=luMN5D6rBluEpkKJGRwmueAgsf4_txzLfnLvMXk3k60,62486
+mesa/model.py,sha256=J_zVRH2caf_LIx8I2fi5whyD6N6fTNfnvXVIMdDKzD0,8454
+mesa/space.py,sha256=PI97fZTHnqNLBzVEIrQ88NMGKkTVWwzukiTewLE3bHU,62722
 mesa/time.py,sha256=G83UKWeMFMnQV9-79Ps2gbD_Qz3hM07IFYLzf5Rvz1w,15243
 mesa/cookiecutter-mesa/cookiecutter.json,sha256=tBSWli39fOWUXGfiDCTKd92M7uKaBIswXbkOdbUufYY,337
 mesa/cookiecutter-mesa/hooks/post_gen_project.py,sha256=8JoXZKIioRYEWJURC0udj8WS3rg0c4So62sOZSGbrMY,294
@@ -16,7 +16,7 @@ mesa/cookiecutter-mesa/{{cookiecutter.snake}}/{{cookiecutter.snake}}/model.pytem
 mesa/cookiecutter-mesa/{{cookiecutter.snake}}/{{cookiecutter.snake}}/server.pytemplate,sha256=nqi6cPjhiyrYw82_Y9hLSrfZtSVCGIhMLOXRB7kKTlQ,823
 mesa/experimental/UserParam.py,sha256=kpFu5yH_RSTI4S2XQxYAGWJgGwQo6j-yCADYJHry4lE,1641
 mesa/experimental/__init__.py,sha256=ncy-EG4IMoSzknYXLMt3Z61nkMwkrH96QbJpGUqOpxQ,168
-mesa/experimental/jupyter_viz.py,sha256=q2Xn-QjH0nqJoxY1RV_SSHC_OgMFZyKWHJjrn5jHCRk,13444
+mesa/experimental/jupyter_viz.py,sha256=UzQjWIR_L_flAInphNzOaBtpl58Xo1JnvnuIE9Fwf78,13501
 mesa/experimental/cell_space/__init__.py,sha256=trFVKf2l5RbkCUyxP09Kox_J3ak2YdM4o3t40Tsjjm4,628
 mesa/experimental/cell_space/cell.py,sha256=AUnvVnXWhdgzr0bLKDRDO9c93v22Zkw6W-tWxhEhGdQ,4578
 mesa/experimental/cell_space/cell_agent.py,sha256=G4u9ht4gW9ns1y2L7pFumF3K4HiP6ROuxwrxHZ-mL1M,1107
@@ -29,8 +29,8 @@ mesa/experimental/components/matplotlib.py,sha256=2toQmjvd1_xj_jxdQPRmpDLeuvWvZP
 mesa/experimental/devs/__init__.py,sha256=CWam15vCj-RD_biMyqv4sJfos1fsL823P7MDEGrbwW8,174
 mesa/experimental/devs/eventlist.py,sha256=H9hufe9VmwvlXQr146wCa7PgbzVvivG4Bk9rlEERZ7A,4880
 mesa/experimental/devs/simulator.py,sha256=NQ3rtBIzykBtMWNslG_Fg04NQn2lYT8cmH-7ndr8v0Y,9530
-mesa/experimental/devs/examples/epstein_civil_violence.py,sha256=KqH9KI-A_BYt7oWi9kaOhTzjrf2pETqzSpAQG8ewud0,9667
-mesa/experimental/devs/examples/wolf_sheep.py,sha256=h5z-eDqMpYeOjrq293N2BcQbs_LDVsgtg9vblXJM7XQ,7697
+mesa/experimental/devs/examples/epstein_civil_violence.py,sha256=vJo5Z2vlDe59V1CazJIgtZU4W22xTRphQE0KoZbOImc,9694
+mesa/experimental/devs/examples/wolf_sheep.py,sha256=TbpLo37uk3JIH4MFQujnSmgqi3KPEVROD_DiSpRlv3I,7706
 mesa/flat/__init__.py,sha256=hSqQDjkfIgDu7B3aYtjPeNEUXdlKPHQuNN8HEV0XR6s,218
 mesa/flat/visualization.py,sha256=5aCm8xDCmZij3hoJZvOVmmpzU9ACXSSSvmQr51buLVg,290
 mesa/visualization/ModularVisualization.py,sha256=gT-FYRnvTFFxtehu-N-8eBLShCF8t_Ut8Au4iVPZIlA,60
@@ -38,8 +38,8 @@ mesa/visualization/TextVisualization.py,sha256=BIP0XcmIdYhz0igqe8yRZXlXeOOqJZeu8
 mesa/visualization/UserParam.py,sha256=D3qxoX-Cpqhyn06IdIO_C5s0u8nlhv3988lVwkBlcGo,49
 mesa/visualization/__init__.py,sha256=5fwVAzgVsmxAzgoLxdC26l2ZE-m2bWj963xPNSDaQEQ,287
 mesa/visualization/modules.py,sha256=pf6K3KECX51VNNqpFCm2EE5KV0A22UYmfXzTVXPnF_o,47
-mesa-2.3.4.dist-info/METADATA,sha256=PTPfEpy2KJuMXXdLUAYn6qEm79MkRFEH9UwDDsIa508,8301
-mesa-2.3.4.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-mesa-2.3.4.dist-info/entry_points.txt,sha256=IOcQtetGF8l4wHpOs_hGb19Rz-FS__BMXOJR10IBPsA,39
-mesa-2.3.4.dist-info/licenses/LICENSE,sha256=OGUgret9fRrm8J3pdsPXETIjf0H8puK_Nmy970ZzT78,572
-mesa-2.3.4.dist-info/RECORD,,
+mesa-2.4.0.dist-info/METADATA,sha256=pLU_lLJDgsowHvvfgh1n-beLMC406r_ArpWKwTyDbm8,8346
+mesa-2.4.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+mesa-2.4.0.dist-info/entry_points.txt,sha256=IOcQtetGF8l4wHpOs_hGb19Rz-FS__BMXOJR10IBPsA,39
+mesa-2.4.0.dist-info/licenses/LICENSE,sha256=OGUgret9fRrm8J3pdsPXETIjf0H8puK_Nmy970ZzT78,572
+mesa-2.4.0.dist-info/RECORD,,