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

mesa/experimental/components/altair.py

@@ -1,5 +1,6 @@
+"""Altair components."""
+
 import contextlib
-from typing import Optional
 
 import solara
 
@@ -8,7 +9,15 @@ with contextlib.suppress(ImportError):
 
 
 @solara.component
-def SpaceAltair(model, agent_portrayal, dependencies: Optional[list[any]] = None):
+def SpaceAltair(model, agent_portrayal, dependencies: list[any] | None = None):
+    """A component that renders a Space using Altair.
+
+    Args:
+        model: a model instance
+        agent_portrayal: agent portray specification
+        dependencies: optional list of dependencies (currently not used)
+
+    """
     space = getattr(model, "grid", None)
     if space is None:
         # Sometimes the space is defined as model.space instead of model.grid

mesa/experimental/components/matplotlib.py

@@ -1,4 +1,6 @@
-from typing import Optional
+"""Support for using matplotlib to draw spaces."""
+
+from collections import defaultdict
 
 import networkx as nx
 import solara
@@ -6,10 +8,19 @@ from matplotlib.figure import Figure
 from matplotlib.ticker import MaxNLocator
 
 import mesa
+from mesa.experimental.cell_space import VoronoiGrid
 
 
 @solara.component
-def SpaceMatplotlib(model, agent_portrayal, dependencies: Optional[list[any]] = None):
+def SpaceMatplotlib(model, agent_portrayal, dependencies: list[any] | None = None):
+    """A component for rendering a space using Matplotlib.
+
+    Args:
+        model: a model instance
+        agent_portrayal: a specification of how to portray an agent.
+        dependencies: list of dependencies.
+
+    """
     space_fig = Figure()
     space_ax = space_fig.subplots()
     space = getattr(model, "grid", None)
@@ -20,17 +31,51 @@ def SpaceMatplotlib(model, agent_portrayal, dependencies: Optional[list[any]] =
         _draw_network_grid(space, space_ax, agent_portrayal)
     elif isinstance(space, mesa.space.ContinuousSpace):
         _draw_continuous_space(space, space_ax, agent_portrayal)
+    elif isinstance(space, VoronoiGrid):
+        _draw_voronoi(space, space_ax, agent_portrayal)
     else:
         _draw_grid(space, space_ax, agent_portrayal)
     solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
 
 
+# matplotlib scatter does not allow for multiple shapes in one call
+def _split_and_scatter(portray_data, space_ax):
+    grouped_data = defaultdict(lambda: {"x": [], "y": [], "s": [], "c": []})
+
+    # Extract data from the dictionary
+    x = portray_data["x"]
+    y = portray_data["y"]
+    s = portray_data["s"]
+    c = portray_data["c"]
+    m = portray_data["m"]
+
+    if not (len(x) == len(y) == len(s) == len(c) == len(m)):
+        raise ValueError(
+            "Length mismatch in portrayal data lists: "
+            f"x: {len(x)}, y: {len(y)}, size: {len(s)}, "
+            f"color: {len(c)}, marker: {len(m)}"
+        )
+
+    # Group the data by marker
+    for i in range(len(x)):
+        marker = m[i]
+        grouped_data[marker]["x"].append(x[i])
+        grouped_data[marker]["y"].append(y[i])
+        grouped_data[marker]["s"].append(s[i])
+        grouped_data[marker]["c"].append(c[i])
+
+    # Plot each group with the same marker
+    for marker, data in grouped_data.items():
+        space_ax.scatter(data["x"], data["y"], s=data["s"], c=data["c"], marker=marker)
+
+
 def _draw_grid(space, space_ax, agent_portrayal):
     def portray(g):
         x = []
         y = []
         s = []  # size
         c = []  # color
+        m = []  # shape
         for i in range(g.width):
             for j in range(g.height):
                 content = g._grid[i][j]
@@ -43,23 +88,23 @@ def _draw_grid(space, space_ax, agent_portrayal):
                     data = agent_portrayal(agent)
                     x.append(i)
                     y.append(j)
-                    if "size" in data:
-                        s.append(data["size"])
-                    if "color" in data:
-                        c.append(data["color"])
-        out = {"x": x, "y": y}
-        # This is the default value for the marker size, which auto-scales
-        # according to the grid area.
-        out["s"] = (180 / max(g.width, g.height)) ** 2
-        if len(s) > 0:
-            out["s"] = s
-        if len(c) > 0:
-            out["c"] = c
+
+                    # This is the default value for the marker size, which auto-scales
+                    # according to the grid area.
+                    default_size = (180 / max(g.width, g.height)) ** 2
+                    # establishing a default prevents misalignment if some agents are not given size, color, etc.
+                    size = data.get("size", default_size)
+                    s.append(size)
+                    color = data.get("color", "tab:blue")
+                    c.append(color)
+                    mark = data.get("shape", "o")
+                    m.append(mark)
+        out = {"x": x, "y": y, "s": s, "c": c, "m": m}
         return out
 
     space_ax.set_xlim(-1, space.width)
     space_ax.set_ylim(-1, space.height)
-    space_ax.scatter(**portray(space))
+    _split_and_scatter(portray(space), space_ax)
 
 
 def _draw_network_grid(space, space_ax, agent_portrayal):
@@ -79,20 +124,23 @@ def _draw_continuous_space(space, space_ax, agent_portrayal):
         y = []
         s = []  # size
         c = []  # color
+        m = []  # shape
         for agent in space._agent_to_index:
             data = agent_portrayal(agent)
             _x, _y = agent.pos
             x.append(_x)
             y.append(_y)
-            if "size" in data:
-                s.append(data["size"])
-            if "color" in data:
-                c.append(data["color"])
-        out = {"x": x, "y": y}
-        if len(s) > 0:
-            out["s"] = s
-        if len(c) > 0:
-            out["c"] = c
+
+            # This is matplotlib's default marker size
+            default_size = 20
+            # establishing a default prevents misalignment if some agents are not given size, color, etc.
+            size = data.get("size", default_size)
+            s.append(size)
+            color = data.get("color", "tab:blue")
+            c.append(color)
+            mark = data.get("shape", "o")
+            m.append(mark)
+        out = {"x": x, "y": y, "s": s, "c": c, "m": m}
         return out
 
     # Determine border style based on space.torus
@@ -112,11 +160,69 @@ def _draw_continuous_space(space, space_ax, agent_portrayal):
     space_ax.set_ylim(space.y_min - y_padding, space.y_max + y_padding)
 
     # Portray and scatter the agents in the space
+    _split_and_scatter(portray(space), space_ax)
+
+
+def _draw_voronoi(space, space_ax, agent_portrayal):
+    def portray(g):
+        x = []
+        y = []
+        s = []  # size
+        c = []  # color
+
+        for cell in g.all_cells:
+            for agent in cell.agents:
+                data = agent_portrayal(agent)
+                x.append(cell.coordinate[0])
+                y.append(cell.coordinate[1])
+                if "size" in data:
+                    s.append(data["size"])
+                if "color" in data:
+                    c.append(data["color"])
+        out = {"x": x, "y": y}
+        # This is the default value for the marker size, which auto-scales
+        # according to the grid area.
+        out["s"] = s
+        if len(c) > 0:
+            out["c"] = c
+
+        return out
+
+    x_list = [i[0] for i in space.centroids_coordinates]
+    y_list = [i[1] for i in space.centroids_coordinates]
+    x_max = max(x_list)
+    x_min = min(x_list)
+    y_max = max(y_list)
+    y_min = min(y_list)
+
+    width = x_max - x_min
+    x_padding = width / 20
+    height = y_max - y_min
+    y_padding = height / 20
+    space_ax.set_xlim(x_min - x_padding, x_max + x_padding)
+    space_ax.set_ylim(y_min - y_padding, y_max + y_padding)
     space_ax.scatter(**portray(space))
 
+    for cell in space.all_cells:
+        polygon = cell.properties["polygon"]
+        space_ax.fill(
+            *zip(*polygon),
+            alpha=min(1, cell.properties[space.cell_coloring_property]),
+            c="red",
+        )  # Plot filled polygon
+        space_ax.plot(*zip(*polygon), color="black")  # Plot polygon edges in red
+
 
 @solara.component
-def PlotMatplotlib(model, measure, dependencies: Optional[list[any]] = None):
+def PlotMatplotlib(model, measure, dependencies: list[any] | None = None):
+    """A solara component for creating a matplotlib figure.
+
+    Args:
+        model: Model instance
+        measure: measure to plot
+        dependencies: list of additional dependencies
+
+    """
     fig = Figure()
     ax = fig.subplots()
     df = model.datacollector.get_model_vars_dataframe()
@@ -127,7 +233,7 @@ def PlotMatplotlib(model, measure, dependencies: Optional[list[any]] = None):
         for m, color in measure.items():
             ax.plot(df.loc[:, m], label=m, color=color)
         fig.legend()
-    elif isinstance(measure, (list, tuple)):
+    elif isinstance(measure, list | tuple):
         for m in measure:
             ax.plot(df.loc[:, m], label=m)
         fig.legend()

mesa/experimental/devs/__init__.py

@@ -1,3 +1,5 @@
+"""Support for event scheduling."""
+
 from .eventlist import Priority, SimulationEvent
 from .simulator import ABMSimulator, DEVSimulator
 

mesa/experimental/devs/eventlist.py

@@ -1,23 +1,28 @@
+"""Eventlist which is at the core of event scheduling."""
+
 from __future__ import annotations
 
 import itertools
+from collections.abc import Callable
 from enum import IntEnum
 from heapq import heapify, heappop, heappush
 from types import MethodType
-from typing import Any, Callable
+from typing import Any
 from weakref import WeakMethod, ref
 
 
 class Priority(IntEnum):
+    """Enumeration of priority levels."""
+
     LOW = 10
     DEFAULT = 5
     HIGH = 1
 
 
 class SimulationEvent:
-    """A simulation event
+    """A simulation event.
 
-    the callable is wrapped using weakref, so there is no need to explicitly cancel event if e.g., an agent
+    The callable is wrapped using weakref, so there is no need to explicitly cancel event if e.g., an agent
     is removed from the simulation.
 
     Attributes:
@@ -28,12 +33,18 @@ class SimulationEvent:
         function_args (list[Any]): Argument for the function
         function_kwargs (Dict[str, Any]): Keyword arguments for the function
 
+
+    Notes:
+        simulation events use a weak reference to the callable. Therefore, you cannot pass a lambda function in fn.
+        A simulation event where the callable no longer exists (e.g., because the agent has been removed from the model)
+        will fail silently.
+
     """
 
     _ids = itertools.count()
 
     @property
-    def CANCELED(self) -> bool:
+    def CANCELED(self) -> bool:  # noqa: D102
         return self._canceled
 
     def __init__(
@@ -44,6 +55,15 @@ class SimulationEvent:
         function_args: list[Any] | None = None,
         function_kwargs: dict[str, Any] | None = None,
     ) -> None:
+        """Initialize a simulation event.
+
+        Args:
+            time: the instant of time of the simulation event
+            function: the callable to invoke
+            priority: the priority of the event
+            function_args: arguments for callable
+            function_kwargs: keyword arguments for the callable
+        """
         super().__init__()
         if not callable(function):
             raise Exception()
@@ -63,20 +83,20 @@ class SimulationEvent:
         self.function_kwargs = function_kwargs if function_kwargs else {}
 
     def execute(self):
-        """execute this event"""
+        """Execute this event."""
         if not self._canceled:
             fn = self.fn()
             if fn is not None:
                 fn(*self.function_args, **self.function_kwargs)
 
     def cancel(self) -> None:
-        """cancel this event"""
+        """Cancel this event."""
         self._canceled = True
         self.fn = None
         self.function_args = []
         self.function_kwargs = {}
 
-    def __lt__(self, other):
+    def __lt__(self, other):  # noqa
         # Define a total ordering for events to be used by the heapq
         return (self.time, self.priority, self.unique_id) < (
             other.time,
@@ -86,30 +106,31 @@ class SimulationEvent:
 
 
 class EventList:
-    """An event list
+    """An event list.
 
     This is a heap queue sorted list of events. Events are always removed from the left, so heapq is a performant and
     appropriate data structure. Events are sorted based on their time stamp, their priority, and their unique_id
     as a tie-breaker, guaranteeing a complete ordering.
 
+
     """
 
     def __init__(self):
+        """Initialize an event list."""
         self._events: list[SimulationEvent] = []
         heapify(self._events)
 
     def add_event(self, event: SimulationEvent):
-        """Add the event to the event list
+        """Add the event to the event list.
 
         Args:
             event (SimulationEvent): The event to be added
 
         """
-
         heappush(self._events, event)
 
     def peak_ahead(self, n: int = 1) -> list[SimulationEvent]:
-        """Look at the first n non-canceled event in the event list
+        """Look at the first n non-canceled event in the event list.
 
         Args:
             n (int): The number of events to look ahead
@@ -138,7 +159,7 @@ class EventList:
         return peek
 
     def pop_event(self) -> SimulationEvent:
-        """pop the first element from the event list"""
+        """Pop the first element from the event list."""
         while self._events:
             event = heappop(self._events)
             if not event.CANCELED:
@@ -146,16 +167,33 @@ class EventList:
         raise IndexError("Event list is empty")
 
     def is_empty(self) -> bool:
+        """Return whether the event list is empty."""
         return len(self) == 0
 
-    def __contains__(self, event: SimulationEvent) -> bool:
+    def __contains__(self, event: SimulationEvent) -> bool:  # noqa
         return event in self._events
 
-    def __len__(self) -> int:
+    def __len__(self) -> int:  # noqa
         return len(self._events)
 
+    def __repr__(self) -> str:
+        """Return a string representation of the event list."""
+        events_str = ", ".join(
+            [
+                f"Event(time={e.time}, priority={e.priority}, id={e.unique_id})"
+                for e in self._events
+                if not e.CANCELED
+            ]
+        )
+        return f"EventList([{events_str}])"
+
     def remove(self, event: SimulationEvent) -> None:
-        """remove an event from the event list"""
+        """Remove an event from the event list.
+
+        Args:
+            event (SimulationEvent): The event to be removed
+
+        """
         # we cannot simply remove items from _eventlist because this breaks
         # heap structure invariant. So, we use a form of lazy deletion.
         # SimEvents have a CANCELED flag that we set to True, while popping and peak_ahead
@@ -163,4 +201,5 @@ class EventList:
         event.cancel()
 
     def clear(self):
+        """Clear the event list."""
         self._events.clear()

mesa/experimental/devs/examples/epstein_civil_violence.py

@@ -1,3 +1,5 @@
+"""Epstein civil violence example using ABMSimulator."""
+
 import enum
 import math
 
@@ -7,21 +9,32 @@ from mesa.space import SingleGrid
 
 
 class EpsteinAgent(Agent):
-    def __init__(self, unique_id, model, vision, movement):
-        super().__init__(unique_id, model)
+    """Epstein Agent."""
+
+    def __init__(self, model, vision, movement):
+        """Initialize the agent.
+
+        Args:
+            model: a model instance
+            vision: size of neighborhood
+            movement: boolean whether agent can move or not
+        """
+        super().__init__(model)
         self.vision = vision
         self.movement = movement
 
 
 class AgentState(enum.IntEnum):
+    """Agent states."""
+
     QUIESCENT = enum.auto()
     ARRESTED = enum.auto()
     ACTIVE = enum.auto()
 
 
 class Citizen(EpsteinAgent):
-    """
-    A member of the general population, may or may not be in active rebellion.
+    """A member of the general population, may or may not be in active rebellion.
+
     Summary of rule: If grievance - risk > threshold, rebel.
 
     Attributes:
@@ -46,7 +59,6 @@ class Citizen(EpsteinAgent):
 
     def __init__(
         self,
-        unique_id,
         model,
         vision,
         movement,
@@ -56,11 +68,13 @@ class Citizen(EpsteinAgent):
         threshold,
         arrest_prob_constant,
     ):
-        """
-        Create a new Citizen.
+        """Create a new Citizen.
+
         Args:
-            unique_id: unique int
             model : model instance
+            vision: number of cells in each direction (N, S, E and W) that
+                agent can inspect. Exogenous.
+            movement: whether agent can move or not
             hardship: Agent's 'perceived hardship (i.e., physical or economic
                 privation).' Exogenous, drawn from U(0,1).
             regime_legitimacy: Agent's perception of regime legitimacy, equal
@@ -68,10 +82,10 @@ class Citizen(EpsteinAgent):
             risk_aversion: Exogenous, drawn from U(0,1).
             threshold: if (grievance - (risk_aversion * arrest_probability)) >
                 threshold, go/remain Active
-            vision: number of cells in each direction (N, S, E and W) that
-                agent can inspect. Exogenous.
+            arrest_prob_constant : agent's assessment of arrest probability
+
         """
-        super().__init__(unique_id, model, vision, movement)
+        super().__init__(model, vision, movement)
         self.hardship = hardship
         self.regime_legitimacy = regime_legitimacy
         self.risk_aversion = risk_aversion
@@ -82,9 +96,7 @@ class Citizen(EpsteinAgent):
         self.arrest_prob_constant = arrest_prob_constant
 
     def step(self):
-        """
-        Decide whether to activate, then move if applicable.
-        """
+        """Decide whether to activate, then move if applicable."""
         self.update_neighbors()
         self.update_estimated_arrest_probability()
         net_risk = self.risk_aversion * self.arrest_probability
@@ -97,9 +109,7 @@ class Citizen(EpsteinAgent):
             self.model.grid.move_agent(self, new_pos)
 
     def update_neighbors(self):
-        """
-        Look around and see who my neighbors are
-        """
+        """Look around and see who my neighbors are."""
         self.neighborhood = self.model.grid.get_neighborhood(
             self.pos, moore=True, radius=self.vision
         )
@@ -109,10 +119,7 @@ class Citizen(EpsteinAgent):
         ]
 
     def update_estimated_arrest_probability(self):
-        """
-        Based on the ratio of cops to actives in my neighborhood, estimate the
-        p(Arrest | I go active).
-        """
+        """Based on the ratio of cops to actives in my neighborhood, estimate the p(Arrest | I go active)."""
         cops_in_vision = len([c for c in self.neighbors if isinstance(c, Cop)])
         actives_in_vision = 1.0  # citizen counts herself
         for c in self.neighbors:
@@ -123,18 +130,25 @@ class Citizen(EpsteinAgent):
         )
 
     def sent_to_jail(self, value):
+        """Sent agent to jail.
+
+        Args:
+            value: duration of jail sentence
+
+        """
         self.model.active_agents.remove(self)
         self.condition = AgentState.ARRESTED
         self.model.simulator.schedule_event_relative(self.release_from_jail, value)
 
     def release_from_jail(self):
+        """Release agent from jail."""
         self.model.active_agents.add(self)
         self.condition = AgentState.QUIESCENT
 
 
 class Cop(EpsteinAgent):
-    """
-    A cop for life.  No defection.
+    """A cop for life.  No defection.
+
     Summary of rule: Inspect local vision and arrest a random active agent.
 
     Attributes:
@@ -144,15 +158,20 @@ class Cop(EpsteinAgent):
             able to inspect
     """
 
-    def __init__(self, unique_id, model, vision, movement, max_jail_term):
-        super().__init__(unique_id, model, vision, movement)
+    def __init__(self, model, vision, movement, max_jail_term):
+        """Initialize a Cop agent.
+
+        Args:
+            model: a model instance
+            vision: size of neighborhood
+            movement: whether agent can move or not
+            max_jail_term: maximum jail sentence
+        """
+        super().__init__(model, vision, movement)
         self.max_jail_term = max_jail_term
 
     def step(self):
-        """
-        Inspect local vision and arrest a random active agent. Move if
-        applicable.
-        """
+        """Inspect local vision and arrest a random active agent. Move if applicable."""
         self.update_neighbors()
         active_neighbors = []
         for agent in self.neighbors:
@@ -166,9 +185,7 @@ class Cop(EpsteinAgent):
             self.model.grid.move_agent(self, new_pos)
 
     def update_neighbors(self):
-        """
-        Look around and see who my neighbors are.
-        """
+        """Look around and see who my neighbors are."""
         self.neighborhood = self.model.grid.get_neighborhood(
             self.pos, moore=True, radius=self.vision
         )
@@ -179,9 +196,8 @@ class Cop(EpsteinAgent):
 
 
 class EpsteinCivilViolence(Model):
-    """
-    Model 1 from "Modeling civil violence: An agent-based computational
-    approach," by Joshua Epstein.
+    """Model 1 from "Modeling civil violence: An agent-based computational approach," by Joshua Epstein.
+
     http://www.pnas.org/content/99/suppl_3/7243.full
     Attributes:
         height: grid height
@@ -220,6 +236,23 @@ class EpsteinCivilViolence(Model):
         max_iters=1000,
         seed=None,
     ):
+        """Initialize the Eppstein civil violence model.
+
+        Args:
+            width: the width of the grid
+            height: the height of the grid
+            citizen_density: density of citizens
+            cop_density: density of cops
+            citizen_vision: size of citizen vision
+            cop_vision: size of cop vision
+            legitimacy: perceived legitimacy
+            max_jail_term: maximum jail term
+            active_threshold: threshold for citizen to become active
+            arrest_prob_constant: arrest probability
+            movement: allow agent movement or not
+            max_iters: number of iterations
+            seed: seed for random number generator
+        """
         super().__init__(seed)
         if cop_density + citizen_density > 1:
             raise ValueError("Cop density + citizen density must be less than 1")
@@ -236,7 +269,6 @@ class EpsteinCivilViolence(Model):
         for _, pos in self.grid.coord_iter():
             if self.random.random() < self.cop_density:
                 agent = Cop(
-                    self.next_id(),
                     self,
                     cop_vision,
                     movement,
@@ -244,7 +276,6 @@ class EpsteinCivilViolence(Model):
                 )
             elif self.random.random() < (self.cop_density + self.citizen_density):
                 agent = Citizen(
-                    self.next_id(),
                     self,
                     citizen_vision,
                     movement,
@@ -271,4 +302,4 @@ if __name__ == "__main__":
 
     simulator.setup(model)
 
-    simulator.run(time_delta=100)
+    simulator.run_for(time_delta=100)