Mesa 3.0.0b2__py3-none-any.whl → 3.0.0rc0__py3-none-any.whl

mesa/visualization/mpl_space_drawing.py

@@ -0,0 +1,558 @@
+"""Helper functions for drawing mesa spaces with matplotlib.
+
+These functions are used by the provided matplotlib components, but can also be used to quickly visualize
+a space with matplotlib for example when creating a mp4 of a movie run or when needing a figure
+for a paper.
+
+"""
+
+import itertools
+import math
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+import networkx as nx
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.axes import Axes
+from matplotlib.cm import ScalarMappable
+from matplotlib.collections import PatchCollection
+from matplotlib.colors import LinearSegmentedColormap, Normalize, to_rgba
+from matplotlib.patches import RegularPolygon
+
+import mesa
+from mesa.experimental.cell_space import (
+    OrthogonalMooreGrid,
+    OrthogonalVonNeumannGrid,
+    VoronoiGrid,
+)
+from mesa.space import (
+    ContinuousSpace,
+    HexMultiGrid,
+    HexSingleGrid,
+    MultiGrid,
+    NetworkGrid,
+    PropertyLayer,
+    SingleGrid,
+)
+
+OrthogonalGrid = SingleGrid | MultiGrid | OrthogonalMooreGrid | OrthogonalVonNeumannGrid
+HexGrid = HexSingleGrid | HexMultiGrid | mesa.experimental.cell_space.HexGrid
+Network = NetworkGrid | mesa.experimental.cell_space.Network
+
+
+def collect_agent_data(
+    space: OrthogonalGrid | HexGrid | Network | ContinuousSpace | VoronoiGrid,
+    agent_portrayal: Callable,
+    color="tab:blue",
+    size=25,
+    marker="o",
+    zorder: int = 1,
+):
+    """Collect the plotting data for all agents in the space.
+
+    Args:
+        space: The space containing the Agents.
+        agent_portrayal: A callable that is called with the agent and returns a dict
+        color: default color
+        size: default size
+        marker: default marker
+        zorder: default zorder
+
+    agent_portrayal should return a dict, limited to size (size of marker), color (color of marker), zorder (z-order),
+    and marker (marker style)
+
+    """
+    arguments = {"s": [], "c": [], "marker": [], "zorder": [], "loc": []}
+
+    for agent in space.agents:
+        portray = agent_portrayal(agent)
+        loc = agent.pos
+        if loc is None:
+            loc = agent.cell.coordinate
+
+        arguments["loc"].append(loc)
+        arguments["s"].append(portray.pop("size", size))
+        arguments["c"].append(portray.pop("color", color))
+        arguments["marker"].append(portray.pop("marker", marker))
+        arguments["zorder"].append(portray.pop("zorder", zorder))
+
+        if len(portray) > 0:
+            ignored_fields = list(portray.keys())
+            msg = ", ".join(ignored_fields)
+            warnings.warn(
+                f"the following fields are not used in agent portrayal and thus ignored: {msg}.",
+                stacklevel=2,
+            )
+
+    return {k: np.asarray(v) for k, v in arguments.items()}
+
+
+def draw_space(
+    space,
+    agent_portrayal: Callable,
+    propertylayer_portrayal: dict | None = None,
+    ax: Axes | None = None,
+    **space_drawing_kwargs,
+):
+    """Draw a Matplotlib-based visualization of the space.
+
+    Args:
+        space: the space of the mesa model
+        agent_portrayal: A callable that returns a dict specifying how to show the agent
+        propertylayer_portrayal: a dict specifying how to show propertylayer(s)
+        ax: the axes upon which to draw the plot
+        post_process: a callable called with the Axes instance
+        space_drawing_kwargs: any additional keyword arguments to be passed on to the underlying function for drawing the space.
+
+    Returns:
+        Returns the Axes object with the plot drawn onto it.
+
+    ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color",
+    "size", "marker", and "zorder". Other field are ignored and will result in a user warning.
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    # https://stackoverflow.com/questions/67524641/convert-multiple-isinstance-checks-to-structural-pattern-matching
+    match space:
+        case mesa.space._Grid() | OrthogonalMooreGrid() | OrthogonalVonNeumannGrid():
+            draw_orthogonal_grid(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
+        case HexSingleGrid() | HexMultiGrid() | mesa.experimental.cell_space.HexGrid():
+            draw_hex_grid(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
+        case mesa.space.NetworkGrid() | mesa.experimental.cell_space.Network():
+            draw_network(space, agent_portrayal, ax=ax, **space_drawing_kwargs)
+        case mesa.space.ContinuousSpace():
+            draw_continuous_space(space, agent_portrayal, ax=ax)
+        case VoronoiGrid():
+            draw_voroinoi_grid(space, agent_portrayal, ax=ax)
+
+    if propertylayer_portrayal:
+        draw_property_layers(space, propertylayer_portrayal, ax=ax)
+
+    return ax
+
+
+def draw_property_layers(
+    space, propertylayer_portrayal: dict[str, dict[str, Any]], ax: Axes
+):
+    """Draw PropertyLayers on the given axes.
+
+    Args:
+        space (mesa.space._Grid): The space containing the PropertyLayers.
+        propertylayer_portrayal (dict): the key is the name of the layer, the value is a dict with
+                                        fields specifying how the layer is to be portrayed
+        ax (matplotlib.axes.Axes): The axes to draw on.
+
+    Notes:
+        valid fields in in the inner dict of propertylayer_portrayal are "alpha", "vmin", "vmax", "color" or "colormap", and "colorbar"
+        so you can do `{"some_layer":{"colormap":'viridis', 'alpha':.25, "colorbar":False}}`
+
+    """
+    try:
+        # old style spaces
+        property_layers = space.properties
+    except AttributeError:
+        # new style spaces
+        property_layers = space.property_layers
+
+    for layer_name, portrayal in propertylayer_portrayal.items():
+        layer = property_layers.get(layer_name, None)
+        if not isinstance(layer, PropertyLayer):
+            continue
+
+        data = layer.data.astype(float) if layer.data.dtype == bool else layer.data
+        width, height = data.shape if space is None else (space.width, space.height)
+
+        if space and data.shape != (width, height):
+            warnings.warn(
+                f"Layer {layer_name} dimensions ({data.shape}) do not match space dimensions ({width}, {height}).",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        # Get portrayal properties, or use defaults
+        alpha = portrayal.get("alpha", 1)
+        vmin = portrayal.get("vmin", np.min(data))
+        vmax = portrayal.get("vmax", np.max(data))
+        colorbar = portrayal.get("colorbar", True)
+
+        # Draw the layer
+        if "color" in portrayal:
+            rgba_color = to_rgba(portrayal["color"])
+            normalized_data = (data - vmin) / (vmax - vmin)
+            rgba_data = np.full((*data.shape, 4), rgba_color)
+            rgba_data[..., 3] *= normalized_data * alpha
+            rgba_data = np.clip(rgba_data, 0, 1)
+            cmap = LinearSegmentedColormap.from_list(
+                layer_name, [(0, 0, 0, 0), (*rgba_color[:3], alpha)]
+            )
+            im = ax.imshow(
+                rgba_data.transpose(1, 0, 2),
+                origin="lower",
+            )
+            if colorbar:
+                norm = Normalize(vmin=vmin, vmax=vmax)
+                sm = ScalarMappable(norm=norm, cmap=cmap)
+                sm.set_array([])
+                ax.figure.colorbar(sm, ax=ax, orientation="vertical")
+
+        elif "colormap" in portrayal:
+            cmap = portrayal.get("colormap", "viridis")
+            if isinstance(cmap, list):
+                cmap = LinearSegmentedColormap.from_list(layer_name, cmap)
+            im = ax.imshow(
+                data.T,
+                cmap=cmap,
+                alpha=alpha,
+                vmin=vmin,
+                vmax=vmax,
+                origin="lower",
+            )
+            if colorbar:
+                plt.colorbar(im, ax=ax, label=layer_name)
+        else:
+            raise ValueError(
+                f"PropertyLayer {layer_name} portrayal must include 'color' or 'colormap'."
+            )
+
+
+def draw_orthogonal_grid(
+    space: OrthogonalGrid,
+    agent_portrayal: Callable,
+    ax: Axes | None = None,
+    draw_grid: bool = True,
+    **kwargs,
+):
+    """Visualize a orthogonal grid.
+
+    Args:
+        space: the space to visualize
+        agent_portrayal: a callable that is called with the agent and returns a dict
+        ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        draw_grid: whether to draw the grid
+        kwargs: additional keyword arguments passed to ax.scatter
+
+    Returns:
+        Returns the Axes object with the plot drawn onto it.
+
+    ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color",
+    "size", "marker", and "zorder". Other field are ignored and will result in a user warning.
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    # gather agent data
+    s_default = (180 / max(space.width, space.height)) ** 2
+    arguments = collect_agent_data(space, agent_portrayal, size=s_default)
+
+    # plot the agents
+    _scatter(ax, arguments, **kwargs)
+
+    # further styling
+    ax.set_xlim(-0.5, space.width - 0.5)
+    ax.set_ylim(-0.5, space.height - 0.5)
+
+    if draw_grid:
+        # Draw grid lines
+        for x in np.arange(-0.5, space.width - 0.5, 1):
+            ax.axvline(x, color="gray", linestyle=":")
+        for y in np.arange(-0.5, space.height - 0.5, 1):
+            ax.axhline(y, color="gray", linestyle=":")
+
+    return ax
+
+
+def draw_hex_grid(
+    space: HexGrid,
+    agent_portrayal: Callable,
+    ax: Axes | None = None,
+    draw_grid: bool = True,
+    **kwargs,
+):
+    """Visualize a hex grid.
+
+    Args:
+        space: the space to visualize
+        agent_portrayal: a callable that is called with the agent and returns a dict
+        ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        draw_grid: whether to draw the grid
+        kwargs: additional keyword arguments passed to ax.scatter
+
+    Returns:
+        Returns the Axes object with the plot drawn onto it.
+
+    ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color",
+    "size", "marker", and "zorder". Other field are ignored and will result in a user warning.
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    # gather data
+    s_default = (180 / max(space.width, space.height)) ** 2
+    arguments = collect_agent_data(space, agent_portrayal, size=s_default)
+
+    # for hexgrids we have to go from logical coordinates to visual coordinates
+    # this is a bit messy.
+
+    # give all even rows an offset in the x direction
+    # give all rows an offset in the y direction
+
+    # numbers here are based on a distance of 1 between centers of hexes
+    offset = math.sqrt(0.75)
+
+    loc = arguments["loc"].astype(float)
+
+    logical = np.mod(loc[:, 1], 2) == 0
+    loc[:, 0][logical] += 0.5
+    loc[:, 1] *= offset
+    arguments["loc"] = loc
+
+    # plot the agents
+    _scatter(ax, arguments, **kwargs)
+
+    # further styling and adding of grid
+    ax.set_xlim(-1, space.width + 0.5)
+    ax.set_ylim(-offset, space.height * offset)
+
+    def setup_hexmesh(
+        width,
+        height,
+    ):
+        """Helper function for creating the hexmaesh."""
+        # fixme: this should be done once, rather than in each update
+        # fixme check coordinate system in hexgrid (see https://www.redblobgames.com/grids/hexagons/#coordinates-offset)
+
+        patches = []
+        for x, y in itertools.product(range(width), range(height)):
+            if y % 2 == 0:
+                x += 0.5  # noqa: PLW2901
+            y *= offset  # noqa: PLW2901
+            hex = RegularPolygon(
+                (x, y),
+                numVertices=6,
+                radius=math.sqrt(1 / 3),
+                orientation=np.radians(120),
+            )
+            patches.append(hex)
+        mesh = PatchCollection(
+            patches, edgecolor="k", facecolor=(1, 1, 1, 0), linestyle="dotted", lw=1
+        )
+        return mesh
+
+    if draw_grid:
+        # add grid
+        ax.add_collection(
+            setup_hexmesh(
+                space.width,
+                space.height,
+            )
+        )
+    return ax
+
+
+def draw_network(
+    space: Network,
+    agent_portrayal: Callable,
+    ax: Axes | None = None,
+    draw_grid: bool = True,
+    layout_alg=nx.spring_layout,
+    layout_kwargs=None,
+    **kwargs,
+):
+    """Visualize a network space.
+
+    Args:
+        space: the space to visualize
+        agent_portrayal: a callable that is called with the agent and returns a dict
+        ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        draw_grid: whether to draw the grid
+        layout_alg: a networkx layout algorithm or other callable with the same behavior
+        layout_kwargs: a dictionary of keyword arguments for the layout algorithm
+        kwargs: additional keyword arguments passed to ax.scatter
+
+    Returns:
+        Returns the Axes object with the plot drawn onto it.
+
+    ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color",
+    "size", "marker", and "zorder". Other field are ignored and will result in a user warning.
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+    if layout_kwargs is None:
+        layout_kwargs = {"seed": 0}
+
+    # gather locations for nodes in network
+    graph = space.G
+    pos = layout_alg(graph, **layout_kwargs)
+    x, y = list(zip(*pos.values()))
+    xmin, xmax = min(x), max(x)
+    ymin, ymax = min(y), max(y)
+
+    width = xmax - xmin
+    height = ymax - ymin
+    x_padding = width / 20
+    y_padding = height / 20
+
+    # gather agent data
+    s_default = (180 / max(width, height)) ** 2
+    arguments = collect_agent_data(space, agent_portrayal, size=s_default)
+
+    # this assumes that nodes are identified by an integer
+    # which is true for default nx graphs but might user changeable
+    pos = np.asarray(list(pos.values()))
+    arguments["loc"] = pos[arguments["loc"]]
+
+    # plot the agents
+    _scatter(ax, arguments, **kwargs)
+
+    # further styling
+    ax.set_axis_off()
+    ax.set_xlim(xmin=xmin - x_padding, xmax=xmax + x_padding)
+    ax.set_ylim(ymin=ymin - y_padding, ymax=ymax + y_padding)
+
+    if draw_grid:
+        # fixme we need to draw the empty nodes as well
+        edge_collection = nx.draw_networkx_edges(
+            graph, pos, ax=ax, alpha=0.5, style="--"
+        )
+        edge_collection.set_zorder(0)
+
+    return ax
+
+
+def draw_continuous_space(
+    space: ContinuousSpace, agent_portrayal: Callable, ax: Axes | None = None, **kwargs
+):
+    """Visualize a continuous space.
+
+    Args:
+        space: the space to visualize
+        agent_portrayal: a callable that is called with the agent and returns a dict
+        ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        kwargs: additional keyword arguments passed to ax.scatter
+
+    Returns:
+        Returns the Axes object with the plot drawn onto it.
+
+    ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color",
+    "size", "marker", and "zorder". Other field are ignored and will result in a user warning.
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    # space related setup
+    width = space.x_max - space.x_min
+    x_padding = width / 20
+    height = space.y_max - space.y_min
+    y_padding = height / 20
+
+    # gather agent data
+    s_default = (180 / max(width, height)) ** 2
+    arguments = collect_agent_data(space, agent_portrayal, size=s_default)
+
+    # plot the agents
+    _scatter(ax, arguments, **kwargs)
+
+    # further visual styling
+    border_style = "solid" if not space.torus else (0, (5, 10))
+    for spine in ax.spines.values():
+        spine.set_linewidth(1.5)
+        spine.set_color("black")
+        spine.set_linestyle(border_style)
+
+    ax.set_xlim(space.x_min - x_padding, space.x_max + x_padding)
+    ax.set_ylim(space.y_min - y_padding, space.y_max + y_padding)
+
+    return ax
+
+
+def draw_voroinoi_grid(
+    space: VoronoiGrid, agent_portrayal: Callable, ax: Axes | None = None, **kwargs
+):
+    """Visualize a voronoi grid.
+
+    Args:
+        space: the space to visualize
+        agent_portrayal: a callable that is called with the agent and returns a dict
+        ax: a Matplotlib Axes instance. If none is provided a new figure and ax will be created using plt.subplots
+        kwargs: additional keyword arguments passed to ax.scatter
+
+    Returns:
+        Returns the Axes object with the plot drawn onto it.
+
+    ``agent_portrayal`` is called with an agent and should return a dict. Valid fields in this dict are "color",
+    "size", "marker", and "zorder". Other field are ignored and will result in a user warning.
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    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
+
+    s_default = (180 / max(width, height)) ** 2
+    arguments = collect_agent_data(space, agent_portrayal, size=s_default)
+
+    ax.set_xlim(x_min - x_padding, x_max + x_padding)
+    ax.set_ylim(y_min - y_padding, y_max + y_padding)
+
+    _scatter(ax, arguments, **kwargs)
+
+    for cell in space.all_cells:
+        polygon = cell.properties["polygon"]
+        ax.fill(
+            *zip(*polygon),
+            alpha=min(1, cell.properties[space.cell_coloring_property]),
+            c="red",
+            zorder=0,
+        )  # Plot filled polygon
+        ax.plot(*zip(*polygon), color="black")  # Plot polygon edges in black
+
+    return ax
+
+
+def _scatter(ax: Axes, arguments, **kwargs):
+    """Helper function for plotting the agents.
+
+    Args:
+        ax: a Matplotlib Axes instance
+        arguments: the agents specific arguments for platting
+        kwargs: additional keyword arguments for ax.scatter
+
+    """
+    loc = arguments.pop("loc")
+
+    x = loc[:, 0]
+    y = loc[:, 1]
+    marker = arguments.pop("marker")
+    zorder = arguments.pop("zorder")
+
+    for mark in np.unique(marker):
+        mark_mask = marker == mark
+        for z_order in np.unique(zorder):
+            zorder_mask = z_order == zorder
+            logical = mark_mask & zorder_mask
+            ax.scatter(
+                x[logical],
+                y[logical],
+                marker=mark,
+                zorder=z_order,
+                **{k: v[logical] for k, v in arguments.items()},
+                **kwargs,
+            )

mesa/visualization/solara_viz.py

@@ -25,13 +25,14 @@ from __future__ import annotations
 
 import asyncio
 import copy
+import inspect
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Literal
 
 import reacton.core
 import solara
 
-import mesa.visualization.components.altair as components_altair
+import mesa.visualization.components.altair_components as components_altair
 from mesa.visualization.UserParam import Slider
 from mesa.visualization.utils import force_update, update_counter
 
@@ -299,9 +300,12 @@ def ModelCreator(model, model_params, seed=1):
         - The component provides an interface for adjusting user-defined parameters and reseeding the model.
 
     """
-    user_params, fixed_params = split_model_params(model_params)
+    solara.use_effect(
+        lambda: _check_model_params(model.value.__class__.__init__, fixed_params),
+        [model.value],
+    )
 
-    reactive_seed = solara.use_reactive(seed)
+    user_params, fixed_params = split_model_params(model_params)
 
     model_parameters, set_model_parameters = solara.use_state(
         {
@@ -310,29 +314,35 @@ def ModelCreator(model, model_params, seed=1):
         }
     )
 
-    def do_reseed():
-        """Update the random seed for the model."""
-        reactive_seed.value = model.value.random.random()
-
     def on_change(name, value):
-        set_model_parameters({**model_parameters, name: value})
+        new_model_parameters = {**model_parameters, name: value}
+        model.value = model.value.__class__(**new_model_parameters)
+        set_model_parameters(new_model_parameters)
 
-    def create_model():
-        model.value = model.value.__class__(**model_parameters)
-        model.value._seed = reactive_seed.value
+    UserInputs(user_params, on_change=on_change)
 
-    solara.use_effect(create_model, [model_parameters, reactive_seed.value])
 
-    with solara.Row(justify="space-between"):
-        solara.InputText(
-            label="Seed",
-            value=reactive_seed,
-            continuous_update=True,
-        )
+def _check_model_params(init_func, model_params):
+    """Check if model parameters are valid for the model's initialization function.
 
-        solara.Button(label="Reseed", color="primary", on_click=do_reseed)
+    Args:
+        init_func: Model initialization function
+        model_params: Dictionary of model parameters
 
-    UserInputs(user_params, on_change=on_change)
+    Raises:
+        ValueError: If a parameter is not valid for the model's initialization function
+    """
+    model_parameters = inspect.signature(init_func).parameters
+    for name in model_parameters:
+        if (
+            model_parameters[name].default == inspect.Parameter.empty
+            and name not in model_params
+            and name != "self"
+        ):
+            raise ValueError(f"Missing required model parameter: {name}")
+    for name in model_params:
+        if name not in model_parameters:
+            raise ValueError(f"Invalid model parameter: {name}")
 
 
 @solara.component

{mesa-3.0.0b2.dist-info → mesa-3.0.0rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: Mesa
-Version: 3.0.0b2
+Version: 3.0.0rc0
 Summary: Agent-based modeling (ABM) in Python
 Project-URL: homepage, https://github.com/projectmesa/mesa
 Project-URL: repository, https://github.com/projectmesa/mesa