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

mesa/examples/basic/virus_on_network/app.py

@@ -0,0 +1,114 @@
+import math
+
+import solara
+
+from mesa.examples.basic.virus_on_network.model import (
+    State,
+    VirusOnNetwork,
+    number_infected,
+)
+from mesa.visualization import (
+    Slider,
+    SolaraViz,
+    make_plot_component,
+    make_space_component,
+)
+
+
+def agent_portrayal(agent):
+    node_color_dict = {
+        State.INFECTED: "tab:red",
+        State.SUSCEPTIBLE: "tab:green",
+        State.RESISTANT: "tab:gray",
+    }
+    return {"color": node_color_dict[agent.state], "size": 10}
+
+
+def get_resistant_susceptible_ratio(model):
+    ratio = model.resistant_susceptible_ratio()
+    ratio_text = r"$\infty$" if ratio is math.inf else f"{ratio:.2f}"
+    infected_text = str(number_infected(model))
+
+    return solara.Markdown(
+        f"Resistant/Susceptible Ratio: {ratio_text}<br>Infected Remaining: {infected_text}"
+    )
+
+
+model_params = {
+    "num_nodes": Slider(
+        label="Number of agents",
+        value=10,
+        min=10,
+        max=100,
+        step=1,
+    ),
+    "avg_node_degree": Slider(
+        label="Avg Node Degree",
+        value=3,
+        min=3,
+        max=8,
+        step=1,
+    ),
+    "initial_outbreak_size": Slider(
+        label="Initial Outbreak Size",
+        value=1,
+        min=1,
+        max=10,
+        step=1,
+    ),
+    "virus_spread_chance": Slider(
+        label="Virus Spread Chance",
+        value=0.4,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+    "virus_check_frequency": Slider(
+        label="Virus Check Frequency",
+        value=0.4,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+    "recovery_chance": Slider(
+        label="Recovery Chance",
+        value=0.3,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+    "gain_resistance_chance": Slider(
+        label="Gain Resistance Chance",
+        value=0.5,
+        min=0.0,
+        max=1.0,
+        step=0.1,
+    ),
+}
+
+
+def post_process_lineplot(ax):
+    ax.set_ylim(ymin=0)
+    ax.set_ylabel("# people")
+    ax.legend(bbox_to_anchor=(1.05, 1.0), loc="upper left")
+
+
+SpacePlot = make_space_component(agent_portrayal)
+StatePlot = make_plot_component(
+    {"Infected": "tab:red", "Susceptible": "tab:green", "Resistant": "tab:gray"},
+    post_process=post_process_lineplot,
+)
+
+model1 = VirusOnNetwork()
+
+page = SolaraViz(
+    model1,
+    [
+        SpacePlot,
+        StatePlot,
+        get_resistant_susceptible_ratio,
+    ],
+    model_params=model_params,
+    name="Virus Model",
+)
+page  # noqa

mesa/examples/basic/virus_on_network/model.py

@@ -0,0 +1,96 @@
+import math
+
+import networkx as nx
+
+import mesa
+from mesa import Model
+from mesa.examples.basic.virus_on_network.agents import State, VirusAgent
+
+
+def number_state(model, state):
+    return sum(1 for a in model.grid.get_all_cell_contents() if a.state is state)
+
+
+def number_infected(model):
+    return number_state(model, State.INFECTED)
+
+
+def number_susceptible(model):
+    return number_state(model, State.SUSCEPTIBLE)
+
+
+def number_resistant(model):
+    return number_state(model, State.RESISTANT)
+
+
+class VirusOnNetwork(Model):
+    """A virus model with some number of agents."""
+
+    def __init__(
+        self,
+        num_nodes=10,
+        avg_node_degree=3,
+        initial_outbreak_size=1,
+        virus_spread_chance=0.4,
+        virus_check_frequency=0.4,
+        recovery_chance=0.3,
+        gain_resistance_chance=0.5,
+        seed=None,
+    ):
+        super().__init__(seed=seed)
+        self.num_nodes = num_nodes
+        prob = avg_node_degree / self.num_nodes
+        self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob)
+        self.grid = mesa.space.NetworkGrid(self.G)
+
+        self.initial_outbreak_size = (
+            initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
+        )
+        self.virus_spread_chance = virus_spread_chance
+        self.virus_check_frequency = virus_check_frequency
+        self.recovery_chance = recovery_chance
+        self.gain_resistance_chance = gain_resistance_chance
+
+        self.datacollector = mesa.DataCollector(
+            {
+                "Infected": number_infected,
+                "Susceptible": number_susceptible,
+                "Resistant": number_resistant,
+                "R over S": self.resistant_susceptible_ratio,
+            }
+        )
+
+        # Create agents
+        for node in self.G.nodes():
+            a = VirusAgent(
+                self,
+                State.SUSCEPTIBLE,
+                self.virus_spread_chance,
+                self.virus_check_frequency,
+                self.recovery_chance,
+                self.gain_resistance_chance,
+            )
+
+            # Add the agent to the node
+            self.grid.place_agent(a, node)
+
+        # Infect some nodes
+        infected_nodes = self.random.sample(list(self.G), self.initial_outbreak_size)
+        for a in self.grid.get_cell_list_contents(infected_nodes):
+            a.state = State.INFECTED
+
+        self.running = True
+        self.datacollector.collect(self)
+
+    def resistant_susceptible_ratio(self):
+        try:
+            return number_state(self, State.RESISTANT) / number_state(
+                self, State.SUSCEPTIBLE
+            )
+        except ZeroDivisionError:
+            return math.inf
+
+    def step(self):
+        self.agents.shuffle_do("step")
+        # collect data
+        self.datacollector.collect(self)

mesa/experimental/UserParam.py

@@ -1,7 +1,10 @@
-class UserParam:
+"""helper classes."""
+
+
+class UserParam:  # noqa: D101
     _ERROR_MESSAGE = "Missing or malformed inputs for '{}' Option '{}'"
 
-    def maybe_raise_error(self, param_type, valid):
+    def maybe_raise_error(self, param_type, valid):  # noqa: D102
         if valid:
             return
         msg = self._ERROR_MESSAGE.format(param_type, self.label)
@@ -9,11 +12,9 @@ class UserParam:
 
 
 class Slider(UserParam):
-    """
-    A number-based slider input with settable increment.
+    """A number-based slider input with settable increment.
 
     Example:
-
     slider_option = Slider("My Slider", value=123, min=10, max=200, step=0.1)
 
     Args:
@@ -34,6 +35,16 @@ class Slider(UserParam):
         step=1,
         dtype=None,
     ):
+        """Slider class.
+
+        Args:
+            label: The displayed label in the UI
+            value: The initial value of the slider
+            min: The minimum possible value of the slider
+            max: The maximum possible value of the slider
+            step: The step between min and max for a range of possible values
+            dtype: either int or float
+        """
         self.label = label
         self.value = value
         self.min = min
@@ -47,10 +58,10 @@ class Slider(UserParam):
         if dtype is None:
             self.is_float_slider = self._check_values_are_float(value, min, max, step)
         else:
-            self.is_float_slider = dtype == float
+            self.is_float_slider = dtype is float
 
     def _check_values_are_float(self, value, min, max, step):
         return any(isinstance(n, float) for n in (value, min, max, step))
 
-    def get(self, attr):
+    def get(self, attr):  # noqa: D102
         return getattr(self, attr)

mesa/experimental/__init__.py

@@ -1,5 +1,13 @@
-from .jupyter_viz import JupyterViz, make_text, Slider  # noqa
+"""Experimental init."""
+
 from mesa.experimental import cell_space
 
+try:
+    from .solara_viz import JupyterViz, Slider, SolaraViz, make_text
 
-__all__ = ["JupyterViz", "make_text", "Slider", "cell_space"]
+    __all__ = ["cell_space", "JupyterViz", "Slider", "SolaraViz", "make_text"]
+except ImportError:
+    print(
+        "Could not import SolaraViz. If you need it, install with 'pip install --pre mesa[viz]'"
+    )
+    __all__ = ["cell_space"]

mesa/experimental/cell_space/__init__.py

@@ -1,5 +1,16 @@
+"""Cell spaces.
+
+Cell spaces offer an alternative API for discrete spaces. It is experimental and under development. The API is more
+expressive that the default grids available in `mesa.space`.
+
+"""
+
 from mesa.experimental.cell_space.cell import Cell
-from mesa.experimental.cell_space.cell_agent import CellAgent
+from mesa.experimental.cell_space.cell_agent import (
+    CellAgent,
+    FixedAgent,
+    Grid2DMovingAgent,
+)
 from mesa.experimental.cell_space.cell_collection import CellCollection
 from mesa.experimental.cell_space.discrete_space import DiscreteSpace
 from mesa.experimental.cell_space.grid import (
@@ -9,15 +20,19 @@ from mesa.experimental.cell_space.grid import (
     OrthogonalVonNeumannGrid,
 )
 from mesa.experimental.cell_space.network import Network
+from mesa.experimental.cell_space.voronoi import VoronoiGrid
 
 __all__ = [
     "CellCollection",
     "Cell",
     "CellAgent",
+    "Grid2DMovingAgent",
+    "FixedAgent",
     "DiscreteSpace",
     "Grid",
     "HexGrid",
     "OrthogonalMooreGrid",
     "OrthogonalVonNeumannGrid",
     "Network",
+    "VoronoiGrid",
 ]

mesa/experimental/cell_space/cell.py

@@ -1,13 +1,20 @@
+"""The Cell in a cell space."""
+
 from __future__ import annotations
 
-from functools import cache
+from collections.abc import Callable
+from functools import cache, cached_property
 from random import Random
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from mesa.experimental.cell_space.cell_agent import CellAgent
 from mesa.experimental.cell_space.cell_collection import CellCollection
+from mesa.space import PropertyLayer
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell_agent import CellAgent
+    from mesa.agent import Agent
+
+Coordinate = tuple[int, ...]
 
 
 class Cell:
@@ -24,11 +31,13 @@ class Cell:
 
     __slots__ = [
         "coordinate",
-        "_connections",
+        "connections",
         "agents",
         "capacity",
         "properties",
         "random",
+        "_mesa_property_layers",
+        "__dict__",
     ]
 
     # def __new__(cls,
@@ -42,34 +51,40 @@ class Cell:
 
     def __init__(
         self,
-        coordinate: tuple[int, ...],
-        capacity: float | None = None,
+        coordinate: Coordinate,
+        capacity: int | None = None,
         random: Random | None = None,
     ) -> None:
-        """ "
+        """Initialise the cell.
 
         Args:
-            coordinate:
+            coordinate: coordinates of the cell
             capacity (int) : the capacity of the cell. If None, the capacity is infinite
             random (Random) : the random number generator to use
 
         """
         super().__init__()
         self.coordinate = coordinate
-        self._connections: list[Cell] = []  # TODO: change to CellCollection?
-        self.agents = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
-        self.capacity = capacity
-        self.properties: dict[str, object] = {}
+        self.connections: dict[Coordinate, Cell] = {}
+        self.agents: list[
+            Agent
+        ] = []  # TODO:: change to AgentSet or weakrefs? (neither is very performant, )
+        self.capacity: int | None = capacity
+        self.properties: dict[Coordinate, object] = {}
         self.random = random
+        self._mesa_property_layers: dict[str, PropertyLayer] = {}
 
-    def connect(self, other: Cell) -> None:
+    def connect(self, other: Cell, key: Coordinate | None = None) -> None:
         """Connects this cell to another cell.
 
         Args:
             other (Cell): other cell to connect to
+            key (Tuple[int, ...]): key for the connection. Should resemble a relative coordinate
 
         """
-        self._connections.append(other)
+        if key is None:
+            key = other.coordinate
+        self.connections[key] = other
 
     def disconnect(self, other: Cell) -> None:
         """Disconnects this cell from another cell.
@@ -78,7 +93,9 @@ class Cell:
             other (Cell): other cell to remove from connections
 
         """
-        self._connections.remove(other)
+        keys_to_remove = [k for k, v in self.connections.items() if v == other]
+        for key in keys_to_remove:
+            del self.connections[key]
 
     def add_agent(self, agent: CellAgent) -> None:
         """Adds an agent to the cell.
@@ -104,7 +121,6 @@ class Cell:
 
         """
         self.agents.remove(agent)
-        agent.cell = None
 
     @property
     def is_empty(self) -> bool:
@@ -116,37 +132,91 @@ class Cell:
         """Returns a bool of the contents of a cell."""
         return len(self.agents) == self.capacity
 
-    def __repr__(self):
+    def __repr__(self):  # noqa
         return f"Cell({self.coordinate}, {self.agents})"
 
+    @cached_property
+    def neighborhood(self) -> CellCollection[Cell]:
+        """Returns the direct neighborhood of the cell.
+
+        This is equivalent to cell.get_neighborhood(radius=1)
+
+        """
+        return self.get_neighborhood()
+
     # FIXME: Revisit caching strategy on methods
     @cache  # noqa: B019
-    def neighborhood(self, radius=1, include_center=False):
-        return CellCollection(
+    def get_neighborhood(
+        self, radius: int = 1, include_center: bool = False
+    ) -> CellCollection[Cell]:
+        """Returns a list of all neighboring cells for the given radius.
+
+        For getting the direct neighborhood (i.e., radius=1) you can also use
+        the `neighborhood` property.
+
+        Args:
+            radius (int): the radius of the neighborhood
+            include_center (bool): include the center of the neighborhood
+
+        Returns:
+            a list of all neighboring cells
+
+        """
+        return CellCollection[Cell](
             self._neighborhood(radius=radius, include_center=include_center),
             random=self.random,
         )
 
     # FIXME: Revisit caching strategy on methods
     @cache  # noqa: B019
-    def _neighborhood(self, radius=1, include_center=False):
+    def _neighborhood(
+        self, radius: int = 1, include_center: bool = False
+    ) -> dict[Cell, list[Agent]]:
         # if radius == 0:
         #     return {self: self.agents}
         if radius < 1:
             raise ValueError("radius must be larger than one")
         if radius == 1:
-            neighborhood = {neighbor: neighbor.agents for neighbor in self._connections}
+            neighborhood = {
+                neighbor: neighbor.agents for neighbor in self.connections.values()
+            }
             if not include_center:
                 return neighborhood
             else:
                 neighborhood[self] = self.agents
                 return neighborhood
         else:
-            neighborhood = {}
-            for neighbor in self._connections:
+            neighborhood: dict[Cell, list[Agent]] = {}
+            for neighbor in self.connections.values():
                 neighborhood.update(
                     neighbor._neighborhood(radius - 1, include_center=True)
                 )
             if not include_center:
                 neighborhood.pop(self, None)
             return neighborhood
+
+    # PropertyLayer methods
+    def get_property(self, property_name: str) -> Any:
+        """Get the value of a property."""
+        return self._mesa_property_layers[property_name].data[self.coordinate]
+
+    def set_property(self, property_name: str, value: Any):
+        """Set the value of a property."""
+        self._mesa_property_layers[property_name].set_cell(self.coordinate, value)
+
+    def modify_property(
+        self, property_name: str, operation: Callable, value: Any = None
+    ):
+        """Modify the value of a property."""
+        self._mesa_property_layers[property_name].modify_cell(
+            self.coordinate, operation, value
+        )
+
+    def __getstate__(self):
+        """Return state of the Cell with connections set to empty."""
+        # fixme, once we shift to 3.11, replace this with super. __getstate__
+        state = (self.__dict__, {k: getattr(self, k) for k in self.__slots__})
+        state[1][
+            "connections"
+        ] = {}  # replace this with empty connections to avoid infinite recursion error in pickle/deepcopy
+        return state

mesa/experimental/cell_space/cell_agent.py

@@ -1,37 +1,133 @@
+"""An agent with movement methods for cell spaces."""
+
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
-from mesa import Agent, Model
+from mesa.agent import Agent
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell import Cell
+    from mesa.experimental.cell_space import Cell
 
 
-class CellAgent(Agent):
-    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces
+class HasCellProtocol(Protocol):
+    """Protocol for discrete space cell holders."""
 
+    cell: Cell
 
-    Attributes:
-        unique_id (int): A unique identifier for this agent.
-        model (Model): The model instance to which the agent belongs
-        pos: (Position | None): The position of the agent in the space
-        cell: (Cell | None): the cell which the agent occupies
-    """
 
-    def __init__(self, unique_id: int, model: Model) -> None:
-        """
-        Create a new agent.
+class HasCell:
+    """Descriptor for cell movement behavior."""
 
-        Args:
-            unique_id (int): A unique identifier for this agent.
-            model (Model): The model instance in which the agent exists.
-        """
-        super().__init__(unique_id, model)
-        self.cell: Cell | None = None
+    _mesa_cell: Cell | None = None
 
-    def move_to(self, cell) -> None:
+    @property
+    def cell(self) -> Cell | None:  # noqa: D102
+        return self._mesa_cell
+
+    @cell.setter
+    def cell(self, cell: Cell | None) -> None:
+        # remove from current cell
         if self.cell is not None:
             self.cell.remove_agent(self)
+
+        # update private attribute
+        self._mesa_cell = cell
+
+        # add to new cell
+        if cell is not None:
+            cell.add_agent(self)
+
+
+class BasicMovement:
+    """Mixin for moving agents in discrete space."""
+
+    def move_to(self: HasCellProtocol, cell: Cell) -> None:
+        """Move to a new cell."""
         self.cell = cell
+
+    def move_relative(self: HasCellProtocol, direction: tuple[int, ...]):
+        """Move to a cell relative to the current cell.
+
+        Args:
+            direction: The direction to move in.
+        """
+        new_cell = self.cell.connections.get(direction)
+        if new_cell is not None:
+            self.cell = new_cell
+        else:
+            raise ValueError(f"No cell in direction {direction}")
+
+
+class FixedCell(HasCell):
+    """Mixin for agents that are fixed to a cell."""
+
+    @property
+    def cell(self) -> Cell | None:  # noqa: D102
+        return self._mesa_cell
+
+    @cell.setter
+    def cell(self, cell: Cell) -> None:
+        if self.cell is not None:
+            raise ValueError("Cannot move agent in FixedCell")
+        self._mesa_cell = cell
+
         cell.add_agent(self)
+
+
+class CellAgent(Agent, HasCell, BasicMovement):
+    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces.
+
+    Attributes:
+        cell (Cell): The cell the agent is currently in.
+    """
+
+    def remove(self):
+        """Remove the agent from the model."""
+        super().remove()
+        self.cell = None  # ensures that we are also removed from cell
+
+
+class FixedAgent(Agent, FixedCell):
+    """A patch in a 2D grid."""
+
+    def remove(self):
+        """Remove the agent from the model."""
+        super().remove()
+
+        # fixme we leave self._mesa_cell on the original value
+        #  so you cannot hijack remove() to move patches
+        self.cell.remove_agent(self)
+
+
+class Grid2DMovingAgent(CellAgent):
+    """Mixin for moving agents in 2D grids."""
+
+    # fmt: off
+    DIRECTION_MAP = {
+        "n": (-1, 0), "north": (-1, 0), "up": (-1, 0),
+        "s": (1, 0), "south": (1, 0), "down": (1, 0),
+        "e": (0, 1), "east": (0, 1), "right": (0, 1),
+        "w": (0, -1), "west": (0, -1), "left": (0, -1),
+        "ne": (-1, 1), "northeast": (-1, 1), "upright": (-1, 1),
+        "nw": (-1, -1), "northwest": (-1, -1), "upleft": (-1, -1),
+        "se": (1, 1), "southeast": (1, 1), "downright": (1, 1),
+        "sw": (1, -1), "southwest": (1, -1), "downleft": (1, -1)
+    }
+    # fmt: on
+
+    def move(self, direction: str, distance: int = 1):
+        """Move the agent in a cardinal direction.
+
+        Args:
+            direction: The cardinal direction to move in.
+            distance: The distance to move.
+        """
+        direction = direction.lower()  # Convert direction to lowercase
+
+        if direction not in self.DIRECTION_MAP:
+            raise ValueError(f"Invalid direction: {direction}")
+
+        move_vector = self.DIRECTION_MAP[direction]
+        for _ in range(distance):
+            self.move_relative(move_vector)