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

mesa/__init__.py

@@ -1,14 +1,13 @@
-"""
-Mesa Agent-Based Modeling Framework
+"""Mesa Agent-Based Modeling Framework.
 
 Core Objects: Model, and Agent.
 """
 
 import datetime
 
+import mesa.experimental as experimental
 import mesa.space as space
 import mesa.time as time
-import mesa.visualization as visualization
 from mesa.agent import Agent
 from mesa.batchrunner import batch_run
 from mesa.datacollection import DataCollector
@@ -19,14 +18,13 @@ __all__ = [
     "Agent",
     "time",
     "space",
-    "visualization",
     "DataCollector",
     "batch_run",
     "experimental",
 ]
 
 __title__ = "mesa"
-__version__ = "2.4.0"
+__version__ = "3.0.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

@@ -1,7 +1,6 @@
-"""
-The agent class for Mesa framework.
+"""Agent related classes.
 
-Core Objects: Agent
+Core Objects: Agent and AgentSet.
 """
 
 # Mypy; for the `|` operator purpose
@@ -10,15 +9,19 @@ from __future__ import annotations
 
 import contextlib
 import copy
+import functools
+import itertools
 import operator
 import warnings
 import weakref
 from collections import defaultdict
-from collections.abc import Hashable, Iterable, Iterator, MutableSet, Sequence
+from collections.abc import Callable, Hashable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
 # mypy
-from typing import TYPE_CHECKING, Any, Callable, Literal, overload
+from typing import TYPE_CHECKING, Any, Literal, overload
+
+import numpy as np
 
 if TYPE_CHECKING:
     # We ensure that these are not imported during runtime to prevent cyclic
@@ -28,77 +31,99 @@ if TYPE_CHECKING:
 
 
 class Agent:
-    """
-    Base class for a model agent in Mesa.
+    """Base class for a model agent in Mesa.
 
     Attributes:
-        unique_id (int): A unique identifier for this agent.
         model (Model): A reference to the model instance.
-        self.pos: Position | None = None
+        unique_id (int): A unique identifier for this agent.
+        pos (Position): A reference to the position where this agent is located.
+
+    Notes:
+          unique_id is unique relative to a model instance and starts from 1
+
     """
 
-    def __init__(self, unique_id: int, model: Model) -> None:
-        """
-        Create a new agent.
+    # this is a class level attribute
+    # it is a dictionary, indexed by model instance
+    # so, unique_id is unique relative to a model, and counting starts from 1
+    _ids = defaultdict(functools.partial(itertools.count, 1))
+
+    def __init__(self, model: Model, *args, **kwargs) -> 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.
+            args: passed on to super
+            kwargs: passed on to super
+
+        Notes:
+            to make proper use of python's super, in each class remove the arguments and
+            keyword arguments you need and pass on the rest to super
+
         """
-        self.unique_id = unique_id
-        self.model = model
-        self.pos: Position | None = None
+        super().__init__(*args, **kwargs)
 
+        self.model: Model = model
+        self.unique_id: int = next(self._ids[model])
+        self.pos: Position | None = None
         self.model.register_agent(self)
 
     def remove(self) -> None:
-        """Remove and delete the agent from the model."""
+        """Remove and delete the agent from the model.
+
+        Notes:
+            If you need to do additional cleanup when removing an agent by for example removing
+            it from a space, consider extending this method in your own agent class.
+
+        """
         with contextlib.suppress(KeyError):
             self.model.deregister_agent(self)
 
     def step(self) -> None:
         """A single step of the agent."""
 
-    def advance(self) -> None:
+    def advance(self) -> None:  # noqa: D102
         pass
 
     @property
     def random(self) -> Random:
+        """Return a seeded stdlib rng."""
         return self.model.random
 
+    @property
+    def rng(self) -> np.random.Generator:
+        """Return a seeded np.random rng."""
+        return self.model.rng
+
 
 class AgentSet(MutableSet, Sequence):
-    """
-    A collection class that represents an ordered set of agents within an agent-based model (ABM). This class
-    extends both MutableSet and Sequence, providing set-like functionality with order preservation and
+    """A collection class that represents an ordered set of agents within an agent-based model (ABM).
+
+    This class extends both MutableSet and Sequence, providing set-like functionality with order preservation and
     sequence operations.
 
     Attributes:
         model (Model): The ABM model instance to which this AgentSet belongs.
 
-    Methods:
-        __len__, __iter__, __contains__, select, shuffle, sort, _update, do, get, __getitem__,
-        add, discard, remove, __getstate__, __setstate__, random
-
-    Note:
+    Notes:
         The AgentSet maintains weak references to agents, allowing for efficient management of agent lifecycles
         without preventing garbage collection. It is associated with a specific model instance, enabling
         interactions with the model's environment and other agents.The implementation uses a WeakKeyDictionary to store agents,
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
     """
 
-    agentset_experimental_warning_given = False
-
-    def __init__(self, agents: Iterable[Agent], model: Model):
-        """
-        Initializes the AgentSet with a collection of agents and a reference to the model.
+    def __init__(self, agents: Iterable[Agent], random: Random | None = None):
+        """Initializes the AgentSet with a collection of agents and a reference to the model.
 
         Args:
             agents (Iterable[Agent]): An iterable of Agent objects to be included in the set.
-            model (Model): The ABM model instance to which this AgentSet belongs.
+            random (Random): the random number generator
         """
-
-        self.model = model
+        if random is None:
+            random = (
+                Random()
+            )  # FIXME see issue 1981, how to get the central rng from model
+        self.random = random
         self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
 
     def __len__(self) -> int:
@@ -121,8 +146,7 @@ class AgentSet(MutableSet, Sequence):
         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.
+        """Select a subset of agents from the AgentSet based on a filter function and/or quantity limit.
 
         Args:
             filter_func (Callable[[Agent], bool], optional): A function that takes an Agent and returns True if the
@@ -132,6 +156,7 @@ class AgentSet(MutableSet, Sequence):
               - 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.
+            n (int): deprecated, use at_most instead
 
         Returns:
             AgentSet: A new AgentSet containing the selected agents, unless inplace is True, in which case the current AgentSet is updated.
@@ -169,11 +194,10 @@ class AgentSet(MutableSet, Sequence):
 
         agents = agent_generator(filter_func, agent_type, at_most)
 
-        return AgentSet(agents, self.model) if not inplace else self._update(agents)
+        return AgentSet(agents, self.random) if not inplace else self._update(agents)
 
     def shuffle(self, inplace: bool = False) -> AgentSet:
-        """
-        Randomly shuffle the order of agents in the AgentSet.
+        """Randomly shuffle the order of agents in the AgentSet.
 
         Args:
             inplace (bool, optional): If True, shuffles the agents in the current AgentSet; otherwise, returns a new shuffled AgentSet. Defaults to False.
@@ -193,7 +217,7 @@ class AgentSet(MutableSet, Sequence):
             return self
         else:
             return AgentSet(
-                (agent for ref in weakrefs if (agent := ref()) is not None), self.model
+                (agent for ref in weakrefs if (agent := ref()) is not None), self.random
             )
 
     def sort(
@@ -202,8 +226,7 @@ class AgentSet(MutableSet, Sequence):
         ascending: bool = False,
         inplace: bool = False,
     ) -> AgentSet:
-        """
-        Sort the agents in the AgentSet based on a specified attribute or custom function.
+        """Sort the agents in the AgentSet based on a specified attribute or custom function.
 
         Args:
             key (Callable[[Agent], Any] | str): A function or attribute name based on which the agents are sorted.
@@ -219,22 +242,21 @@ class AgentSet(MutableSet, Sequence):
         sorted_agents = sorted(self._agents.keys(), key=key, reverse=not ascending)
 
         return (
-            AgentSet(sorted_agents, self.model)
+            AgentSet(sorted_agents, self.random)
             if not inplace
             else self._update(sorted_agents)
         )
 
     def _update(self, agents: Iterable[Agent]):
         """Update the AgentSet with a new set of agents.
+
         This is a private method primarily used internally by other methods like select, shuffle, and sort.
         """
-
         self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
         return self
 
     def do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
-        """
-        Invoke a method or function on each agent in the AgentSet.
+        """Invoke a method or function on each agent in the AgentSet.
 
         Args:
             method (str, callable): the callable to do on each agent
@@ -279,21 +301,22 @@ class AgentSet(MutableSet, Sequence):
 
         It's a fast, optimized version of calling shuffle() followed by do().
         """
-        agents = list(self._agents.keys())
-        self.random.shuffle(agents)
+        weakrefs = list(self._agents.keyrefs())
+        self.random.shuffle(weakrefs)
 
         if isinstance(method, str):
-            for agent in agents:
-                getattr(agent, method)(*args, **kwargs)
+            for ref in weakrefs:
+                if (agent := ref()) is not None:
+                    getattr(agent, method)(*args, **kwargs)
         else:
-            for agent in agents:
-                method(agent, *args, **kwargs)
+            for ref in weakrefs:
+                if (agent := ref()) is not None:
+                    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.
+        """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
@@ -324,8 +347,7 @@ class AgentSet(MutableSet, Sequence):
         return res
 
     def agg(self, attribute: str, func: Callable) -> Any:
-        """
-        Aggregate an attribute of all agents in the AgentSet using a specified function.
+        """Aggregate an attribute of all agents in the AgentSet using a specified function.
 
         Args:
             attribute (str): The name of the attribute to aggregate.
@@ -359,8 +381,7 @@ class AgentSet(MutableSet, Sequence):
         handle_missing="error",
         default_value=None,
     ):
-        """
-        Retrieve the specified attribute(s) from each agent in the AgentSet.
+        """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.
@@ -407,8 +428,7 @@ class AgentSet(MutableSet, Sequence):
             )
 
     def set(self, attr_name: str, value: Any) -> AgentSet:
-        """
-        Set a specified attribute to a given value for all agents in the AgentSet.
+        """Set a specified attribute to a given value for all agents in the AgentSet.
 
         Args:
             attr_name (str): The name of the attribute to set.
@@ -422,8 +442,7 @@ class AgentSet(MutableSet, Sequence):
         return self
 
     def __getitem__(self, item: int | slice) -> Agent:
-        """
-        Retrieve an agent or a slice of agents from the AgentSet.
+        """Retrieve an agent or a slice of agents from the AgentSet.
 
         Args:
             item (int | slice): The index or slice for selecting agents.
@@ -434,8 +453,7 @@ class AgentSet(MutableSet, Sequence):
         return list(self._agents.keys())[item]
 
     def add(self, agent: Agent):
-        """
-        Add an agent to the AgentSet.
+        """Add an agent to the AgentSet.
 
         Args:
             agent (Agent): The agent to add to the set.
@@ -446,8 +464,7 @@ class AgentSet(MutableSet, Sequence):
         self._agents[agent] = None
 
     def discard(self, agent: Agent):
-        """
-        Remove an agent from the AgentSet if it exists.
+        """Remove an agent from the AgentSet if it exists.
 
         This method does not raise an error if the agent is not present.
 
@@ -461,8 +478,7 @@ class AgentSet(MutableSet, Sequence):
             del self._agents[agent]
 
     def remove(self, agent: Agent):
-        """
-        Remove an agent from the AgentSet.
+        """Remove an agent from the AgentSet.
 
         This method raises an error if the agent is not present.
 
@@ -475,37 +491,24 @@ class AgentSet(MutableSet, Sequence):
         del self._agents[agent]
 
     def __getstate__(self):
-        """
-        Retrieve the state of the AgentSet for serialization.
+        """Retrieve the state of the AgentSet for serialization.
 
         Returns:
             dict: A dictionary representing the state of the AgentSet.
         """
-        return {"agents": list(self._agents.keys()), "model": self.model}
+        return {"agents": list(self._agents.keys()), "random": self.random}
 
     def __setstate__(self, state):
-        """
-        Set the state of the AgentSet during deserialization.
+        """Set the state of the AgentSet during deserialization.
 
         Args:
             state (dict): A dictionary representing the state to restore.
         """
-        self.model = state["model"]
+        self.random = state["random"]
         self._update(state["agents"])
 
-    @property
-    def random(self) -> Random:
-        """
-        Provide access to the model's random number generator.
-
-        Returns:
-            Random: The random number generator associated with the model.
-        """
-        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
+        """Group agents by the specified attribute or return from the callable.
 
         Args:
             by (Callable, str): used to determine what to group agents by
@@ -515,6 +518,7 @@ class AgentSet(MutableSet, Sequence):
                                 * 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
 
@@ -535,7 +539,7 @@ class AgentSet(MutableSet, Sequence):
 
         if result_type == "agentset":
             return GroupBy(
-                {k: AgentSet(v, model=self.model) for k, v in groups.items()}
+                {k: AgentSet(v, random=self.random) for k, v in groups.items()}
             )
         else:
             return GroupBy(groups)
@@ -546,8 +550,7 @@ class AgentSet(MutableSet, Sequence):
 
 
 class GroupBy:
-    """Helper class for AgentSet.groupby
-
+    """Helper class for AgentSet.groupby.
 
     Attributes:
         groups (dict): A dictionary with the group_name as key and group as values
@@ -555,6 +558,12 @@ class GroupBy:
     """
 
     def __init__(self, groups: dict[Any, list | AgentSet]):
+        """Initialize a GroupBy instance.
+
+        Args:
+            groups (dict): A dictionary with the group_name as key and group as values
+
+        """
         self.groups: dict[Any, list | AgentSet] = groups
 
     def map(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
@@ -567,6 +576,8 @@ class GroupBy:
                                     * 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.
+            args: arguments to pass to the callable
+            kwargs: keyword arguments to pass to the callable
 
         Returns:
             dict with group_name as key and the return of the method as value
@@ -584,7 +595,7 @@ class GroupBy:
             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
+        """Apply the specified callable to each group.
 
         Args:
             method (Callable, str): The callable to apply to each group,
@@ -593,6 +604,8 @@ class GroupBy:
                                     * 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.
+            args: arguments to pass to the callable
+            kwargs: keyword arguments to pass to the callable
 
         Returns:
             the original GroupBy instance
@@ -634,8 +647,8 @@ class GroupBy:
             for group_name, group in self.groups.items()
         }
 
-    def __iter__(self):
+    def __iter__(self):  # noqa: D105
         return iter(self.groups.items())
 
-    def __len__(self):
+    def __len__(self):  # noqa: D105
         return len(self.groups)

mesa/batchrunner.py

@@ -1,9 +1,36 @@
+"""batchrunner for running a factorial experiment design over a model.
+
+To take advantage of parallel execution of experiments, `batch_run` uses
+multiprocessing if ``number_processes`` is larger than 1. It is strongly advised
+to only run in parallel using a normal python file (so don't try to do it in a
+jupyter notebook). Moreover, best practice when using multiprocessing is to
+put the code inside an ``if __name__ == '__main__':`` code black as shown below::
+
+    from mesa.batchrunner import batch_run
+
+    params = {"width": 10, "height": 10, "N": range(10, 500, 10)}
+
+    if __name__ == '__main__':
+        results = batch_run(
+            MoneyModel,
+            parameters=params,
+            iterations=5,
+            max_steps=100,
+            number_processes=None,
+            data_collection_period=1,
+            display_progress=True,
+        )
+
+
+
+"""
+
 import itertools
 import multiprocessing
 from collections.abc import Iterable, Mapping
 from functools import partial
 from multiprocessing import Pool
-from typing import Any, Optional, Union
+from typing import Any
 
 from tqdm.auto import tqdm
 
@@ -14,9 +41,9 @@ multiprocessing.set_start_method("spawn", force=True)
 
 def batch_run(
     model_cls: type[Model],
-    parameters: Mapping[str, Union[Any, Iterable[Any]]],
+    parameters: Mapping[str, Any | Iterable[Any]],
     # We still retain the Optional[int] because users may set it to None (i.e. use all CPUs)
-    number_processes: Optional[int] = 1,
+    number_processes: int | None = 1,
     iterations: int = 1,
     data_collection_period: int = -1,
     max_steps: int = 1000,
@@ -24,29 +51,22 @@ def batch_run(
 ) -> list[dict[str, Any]]:
     """Batch run a mesa model with a set of parameter values.
 
-    Parameters
-    ----------
-    model_cls : Type[Model]
-        The model class to batch-run
-    parameters : Mapping[str, Union[Any, Iterable[Any]]],
-        Dictionary with model parameters over which to run the model. You can either pass single values or iterables.
-    number_processes : int, optional
-        Number of processes used, by default 1. Set this to None if you want to use all CPUs.
-    iterations : int, optional
-        Number of iterations for each parameter combination, by default 1
-    data_collection_period : int, optional
-        Number of steps after which data gets collected, by default -1 (end of episode)
-    max_steps : int, optional
-        Maximum number of model steps after which the model halts, by default 1000
-    display_progress : bool, optional
-        Display batch run process, by default True
+    Args:
+        model_cls (Type[Model]): The model class to batch-run
+        parameters (Mapping[str, Union[Any, Iterable[Any]]]): Dictionary with model parameters over which to run the model. You can either pass single values or iterables.
+        number_processes (int, optional): Number of processes used, by default 1. Set this to None if you want to use all CPUs.
+        iterations (int, optional): Number of iterations for each parameter combination, by default 1
+        data_collection_period (int, optional): Number of steps after which data gets collected, by default -1 (end of episode)
+        max_steps (int, optional): Maximum number of model steps after which the model halts, by default 1000
+        display_progress (bool, optional): Display batch run process, by default True
 
-    Returns
-    -------
-    List[Dict[str, Any]]
-        [description]
-    """
+    Returns:
+        List[Dict[str, Any]]
 
+    Notes:
+        batch_run assumes the model has a `datacollector` attribute that has a DataCollector object initialized.
+
+    """
     runs_list = []
     run_id = 0
     for iteration in range(iterations):
@@ -79,7 +99,7 @@ def batch_run(
 
 
 def _make_model_kwargs(
-    parameters: Mapping[str, Union[Any, Iterable[Any]]],
+    parameters: Mapping[str, Any | Iterable[Any]],
 ) -> list[dict[str, Any]]:
     """Create model kwargs from parameters dictionary.
 
@@ -88,7 +108,7 @@ def _make_model_kwargs(
     parameters : Mapping[str, Union[Any, Iterable[Any]]]
         Single or multiple values for each model parameter name
 
-    Returns
+    Returns:
     -------
     List[Dict[str, Any]]
         A list of all kwargs combinations.
@@ -128,21 +148,21 @@ def _model_run_func(
     data_collection_period : int
         Number of steps after which data gets collected
 
-    Returns
+    Returns:
     -------
     List[Dict[str, Any]]
         Return model_data, agent_data from the reporters
     """
     run_id, iteration, kwargs = run
     model = model_cls(**kwargs)
-    while model.running and model._steps <= max_steps:
+    while model.running and model.steps <= max_steps:
         model.step()
 
     data = []
 
-    steps = list(range(0, model._steps, data_collection_period))
-    if not steps or steps[-1] != model._steps - 1:
-        steps.append(model._steps - 1)
+    steps = list(range(0, model.steps, data_collection_period))
+    if not steps or steps[-1] != model.steps - 1:
+        steps.append(model.steps - 1)
 
     for step in steps:
         model_data, all_agents_data = _collect_data(model, step)
@@ -181,6 +201,10 @@ def _collect_data(
     step: int,
 ) -> tuple[dict[str, Any], list[dict[str, Any]]]:
     """Collect model and agent data from a model using mesas datacollector."""
+    if not hasattr(model, "datacollector"):
+        raise AttributeError(
+            "The model does not have a datacollector attribute. Please add a DataCollector to your model."
+        )
     dc = model.datacollector
 
     model_data = {param: values[step] for param, values in dc.model_vars.items()}