Mesa 3.2.0.dev0__py3-none-any.whl → 3.3.0__py3-none-any.whl

mesa/visualization/backends/matplotlib_backend.py

@@ -0,0 +1,419 @@
+# noqa: D100
+import os
+import warnings
+from dataclasses import fields
+
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.cm import ScalarMappable
+from matplotlib.collections import PolyCollection
+from matplotlib.colors import LinearSegmentedColormap, Normalize, to_rgba
+from matplotlib.offsetbox import AnnotationBbox, OffsetImage
+from PIL import Image
+
+import mesa
+from mesa.discrete_space import (
+    OrthogonalMooreGrid,
+    OrthogonalVonNeumannGrid,
+)
+from mesa.space import (
+    HexMultiGrid,
+    HexSingleGrid,
+    MultiGrid,
+    SingleGrid,
+)
+from mesa.visualization.backends.abstract_renderer import AbstractRenderer
+
+OrthogonalGrid = SingleGrid | MultiGrid | OrthogonalMooreGrid | OrthogonalVonNeumannGrid
+HexGrid = HexSingleGrid | HexMultiGrid | mesa.discrete_space.HexGrid
+
+
+CORRECTION_FACTOR_MARKER_ZOOM = 0.01
+
+
+class MatplotlibBackend(AbstractRenderer):
+    """Matplotlib-based renderer for Mesa spaces.
+
+    Provides visualization capabilities using Matplotlib for rendering
+    space structures, agents, and property layers.
+    """
+
+    def __init__(self, space_drawer):
+        """Initialize the Matplotlib backend.
+
+        Args:
+            space_drawer: An instance of a SpaceDrawer class that handles
+                the drawing of the space structure.
+        """
+        super().__init__(space_drawer)
+
+        self._active_colorbars = []
+
+    def initialize_canvas(self, ax=None):
+        """Initialize the matplotlib canvas.
+
+        Args:
+            ax (matplotlib.axes.Axes, optional): Existing axes to use.
+                If None, creates new figure and axes.
+        """
+        if ax is None:
+            fig, ax = plt.subplots(constrained_layout=True)
+            self.fig = fig
+        self.ax = ax
+
+    def draw_structure(self, **kwargs):
+        """Draw the space structure using matplotlib.
+
+        Args:
+            **kwargs: Additional arguments passed to the space drawer.
+            Checkout respective `SpaceDrawer` class on details how to pass **kwargs.
+
+        Returns:
+            The matplotlib axes with the drawn structure.
+        """
+        return self.space_drawer.draw_matplotlib(self.ax, **kwargs)
+
+    def collect_agent_data(self, space, agent_portrayal, default_size=None):
+        """Collect plotting data for all agents in the space.
+
+        Args:
+            space: The Mesa space containing agents.
+            agent_portrayal (Callable): Function that returns AgentPortrayalStyle for each agent.
+            default_size (float, optional): Default marker size if not specified in portrayal.
+
+        Returns:
+            dict: Dictionary containing agent plotting data arrays.
+        """
+        # Initialize data collection arrays
+        arguments = {
+            "loc": [],
+            "s": [],
+            "c": [],
+            "marker": [],
+            "zorder": [],
+            "alpha": [],
+            "edgecolors": [],
+            "linewidths": [],
+        }
+        # Import here to prevent circular imports
+        from mesa.visualization.components import AgentPortrayalStyle  # noqa: PLC0415
+
+        # Get default values from AgentPortrayalStyle
+        style_fields = {f.name: f.default for f in fields(AgentPortrayalStyle)}
+        class_default_size = style_fields.get("size")
+
+        for agent in space.agents:
+            portray_input = agent_portrayal(agent)
+
+            if isinstance(portray_input, dict):
+                warnings.warn(
+                    "Returning a dict from agent_portrayal is deprecated. "
+                    "Please return an AgentPortrayalStyle instance instead.",
+                    PendingDeprecationWarning,
+                    stacklevel=2,
+                )
+                # Handle legacy dict input
+                dict_data = portray_input.copy()
+                agent_x, agent_y = self._get_agent_pos(agent, space)
+
+                # Extract values with defaults
+                aps = AgentPortrayalStyle(
+                    x=agent_x,
+                    y=agent_y,
+                    size=dict_data.pop("size", style_fields.get("size")),
+                    color=dict_data.pop("color", style_fields.get("color")),
+                    marker=dict_data.pop("marker", style_fields.get("marker")),
+                    zorder=dict_data.pop("zorder", style_fields.get("zorder")),
+                    alpha=dict_data.pop("alpha", style_fields.get("alpha")),
+                    edgecolors=dict_data.pop("edgecolors", None),
+                    linewidths=dict_data.pop(
+                        "linewidths", style_fields.get("linewidths")
+                    ),
+                )
+
+                # Warn about unused keys
+                if dict_data:
+                    ignored_keys = list(dict_data.keys())
+                    warnings.warn(
+                        f"The following keys were ignored: {', '.join(ignored_keys)}",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+            else:
+                aps = portray_input
+                # Set defaults if not provided
+                if aps.x is None and aps.y is None:
+                    aps.x, aps.y = self._get_agent_pos(agent, space)
+
+            # Collect agent data
+            arguments["loc"].append((aps.x, aps.y))
+
+            # Determine final size
+            size_to_collect = aps.size or default_size or class_default_size
+            arguments["s"].append(size_to_collect)
+            arguments["c"].append(aps.color)
+            arguments["marker"].append(aps.marker)
+            arguments["zorder"].append(aps.zorder)
+            arguments["alpha"].append(aps.alpha)
+            if aps.edgecolors is not None:
+                arguments["edgecolors"].append(aps.edgecolors)
+            arguments["linewidths"].append(aps.linewidths)
+
+        # Convert to numpy arrays
+        data = {
+            k: (np.asarray(v, dtype=object) if k == "marker" else np.asarray(v))
+            for k, v in arguments.items()
+        }
+
+        # Handle marker array specially to preserve tuples
+        arr = np.empty(len(arguments["marker"]), dtype=object)
+        arr[:] = arguments["marker"]
+        data["marker"] = arr
+
+        return data
+
+    def _get_zoom_factor(self, ax, img):
+        """Calculate zoom factor only once and cache the result."""
+        ax.get_figure().canvas.draw()
+
+        bbox = ax.get_window_extent().transformed(
+            ax.get_figure().dpi_scale_trans.inverted()
+        )
+        width, height = (
+            bbox.width * ax.get_figure().dpi,
+            bbox.height * ax.get_figure().dpi,
+        )
+
+        xr = ax.get_xlim()
+        yr = ax.get_ylim()
+
+        x_pixel_per_data = width / (xr[1] - xr[0])
+        y_pixel_per_data = height / (yr[1] - yr[0])
+
+        zoom_x = (x_pixel_per_data / img.width) * CORRECTION_FACTOR_MARKER_ZOOM
+        zoom_y = (y_pixel_per_data / img.height) * CORRECTION_FACTOR_MARKER_ZOOM
+
+        return min(zoom_x, zoom_y)
+
+    def draw_agents(self, arguments, **kwargs):
+        """Draw agents on the backend's axes - optimized version.
+
+        Args:
+            arguments: Dictionary containing agent data arrays.
+            **kwargs: Additional keyword arguments for customization.
+            Checkout respective `SpaceDrawer` class on details how to pass **kwargs.
+
+        Returns:
+            matplotlib.axes.Axes: The Matplotlib Axes with the agents drawn upon it.
+        """
+        if arguments["loc"].size == 0:
+            return None
+
+        loc = arguments.pop("loc")
+        loc_x, loc_y = loc[:, 0], loc[:, 1]
+        marker = arguments.pop("marker")
+        zorder = arguments.pop("zorder")
+        malpha = arguments["alpha"]
+        msize = arguments["s"]
+
+        # Validate edge arguments
+        for entry in ["edgecolors", "linewidths"]:
+            if len(arguments[entry]) == 0:
+                arguments.pop(entry)
+            elif entry in kwargs:
+                raise ValueError(
+                    f"{entry} is specified in agent portrayal and via plotting kwargs, "
+                    "you can only use one or the other"
+                )
+
+        # Cache for loaded images and their zoom factors
+        image_cache = {}
+
+        # Separate image and non-image markers
+        unique_markers = set(marker)
+        image_markers = set()
+        regular_markers = set()
+
+        for mark in unique_markers:
+            if isinstance(mark, str | os.PathLike) and os.path.isfile(mark):
+                image_markers.add(mark)
+            else:
+                regular_markers.add(mark)
+
+        self.ax.get_figure().canvas.draw()
+
+        for mark in image_markers:
+            if mark not in image_cache:
+                image = Image.open(mark)
+                base_zoom = self._get_zoom_factor(self.ax, image)
+                image_cache[mark] = {
+                    "image": image,
+                    "base_zoom": base_zoom,
+                    "offset_images": {},
+                }
+
+            cache_entry = image_cache[mark]
+            mask_marker = marker == mark
+
+            unique_sizes = np.unique(msize[mask_marker])
+
+            for m_size in unique_sizes:
+                if m_size not in cache_entry["offset_images"]:
+                    im = OffsetImage(
+                        cache_entry["image"],
+                        zoom=cache_entry["base_zoom"] * m_size,
+                    )
+                    im.image.axes = self.ax
+                    cache_entry["offset_images"][m_size] = im
+
+                im = cache_entry["offset_images"][m_size]
+
+                size_marker_mask = mask_marker & (msize == m_size)
+
+                unique_zorders = np.unique(zorder[size_marker_mask])
+                unique_alphas = np.unique(malpha[size_marker_mask])
+
+                for z_order in unique_zorders:
+                    for m_alpha in unique_alphas:
+                        final_mask = (
+                            (zorder == z_order) & (malpha == m_alpha) & size_marker_mask
+                        )
+                        positions = loc[final_mask]
+
+                        for x, y in positions:
+                            ab = AnnotationBbox(
+                                im,
+                                (x, y),
+                                frameon=False,
+                                pad=0.0,
+                                zorder=z_order,
+                                **kwargs,
+                            )
+                            self.ax.add_artist(ab)
+
+        for mark in regular_markers:
+            mask_marker = marker == mark
+
+            unique_zorders = np.unique(zorder[mask_marker])
+
+            for z_order in unique_zorders:
+                zorder_mask = (zorder == z_order) & mask_marker
+
+                scatter_args = {k: v[zorder_mask] for k, v in arguments.items()}
+
+                self.ax.scatter(
+                    loc_x[zorder_mask],
+                    loc_y[zorder_mask],
+                    marker=mark,
+                    zorder=z_order,
+                    **scatter_args,
+                    **kwargs,
+                )
+
+        return self.ax
+
+    def draw_propertylayer(self, space, property_layers, propertylayer_portrayal):
+        """Draw property layers using matplotlib backend.
+
+        Args:
+            space: The Mesa space object.
+            property_layers (dict): Dictionary of property layers to visualize.
+            propertylayer_portrayal (Callable): Function that returns PropertyLayerStyle.
+
+        Returns:
+            tuple: (matplotlib.axes.Axes, colorbar) - The matplotlib axes and colorbar objects.
+        """
+        # Draw each layer
+        for layer_name in property_layers:
+            if layer_name == "empty":
+                continue
+
+            layer = property_layers.get(layer_name)
+            portrayal = propertylayer_portrayal(layer)
+
+            if portrayal is None:
+                continue
+
+            data = layer.data.astype(float) if layer.data.dtype == bool else layer.data
+
+            # Check dimensions
+            if (space.width, space.height) != data.shape:
+                warnings.warn(
+                    f"Layer {layer_name} dimensions ({data.shape}) "
+                    f"don't match space dimensions ({space.width}, {space.height})",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                continue
+
+            # Get portrayal parameters
+            color = portrayal.color
+            colormap = portrayal.colormap
+            alpha = portrayal.alpha
+            vmin = portrayal.vmin if portrayal.vmin is not None else np.min(data)
+            vmax = portrayal.vmax if portrayal.vmax is not None else np.max(data)
+
+            # Set up colormap
+            if color:
+                rgba_color = to_rgba(color)
+                cmap = LinearSegmentedColormap.from_list(
+                    layer_name, [(0, 0, 0, 0), (*rgba_color[:3], alpha)]
+                )
+            elif colormap:
+                cmap = colormap
+                if isinstance(cmap, list):
+                    cmap = LinearSegmentedColormap.from_list(layer_name, cmap)
+                elif isinstance(cmap, str):
+                    cmap = plt.get_cmap(cmap)
+            else:
+                raise ValueError(
+                    f"PropertyLayer {layer_name} must include 'color' or 'colormap'"
+                )
+
+            # Draw based on space type
+            if isinstance(space, OrthogonalGrid):
+                if color:
+                    data = data.T
+                    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)
+                    self.ax.imshow(rgba_data, origin="lower")
+                else:
+                    self.ax.imshow(
+                        data.T,
+                        cmap=cmap,
+                        alpha=alpha,
+                        vmin=vmin,
+                        vmax=vmax,
+                        origin="lower",
+                    )
+            elif isinstance(space, HexGrid):
+                hexagons = self.space_drawer.hexagons
+                norm = Normalize(vmin=vmin, vmax=vmax)
+                colors = data.ravel()
+
+                if color:
+                    normalized_colors = np.clip(norm(colors), 0, 1)
+                    rgba_colors = np.full((len(colors), 4), rgba_color)
+                    rgba_colors[:, 3] = normalized_colors * alpha
+                else:
+                    rgba_colors = cmap(norm(colors))
+                    rgba_colors[..., 3] *= alpha
+
+                collection = PolyCollection(hexagons, facecolors=rgba_colors, zorder=-1)
+                self.ax.add_collection(collection)
+            else:
+                raise NotImplementedError(
+                    f"PropertyLayer visualization not implemented for {type(space)}"
+                )
+
+            # Add colorbar if requested
+            cbar = None
+            if portrayal.colorbar:
+                norm = Normalize(vmin=vmin, vmax=vmax)
+                sm = ScalarMappable(norm=norm, cmap=cmap)
+                sm.set_array([])
+                cbar = plt.colorbar(sm, ax=self.ax, label=layer_name)
+                self._active_colorbars.append(cbar)
+        return self.ax, cbar

mesa/visualization/components/__init__.py

@@ -1,15 +1,32 @@
-"""custom solara components."""
+"""Custom visualization components."""
 
 from __future__ import annotations
 
 from collections.abc import Callable
 
-from .altair_components import SpaceAltair, make_altair_space
+from .altair_components import (
+    SpaceAltair,
+    make_altair_plot_component,
+    make_altair_space,
+)
 from .matplotlib_components import (
     SpaceMatplotlib,
     make_mpl_plot_component,
     make_mpl_space_component,
 )
+from .portrayal_components import AgentPortrayalStyle, PropertyLayerStyle
+
+__all__ = [
+    "AgentPortrayalStyle",
+    "PropertyLayerStyle",
+    "SpaceAltair",
+    "SpaceMatplotlib",
+    "make_altair_space",
+    "make_mpl_plot_component",
+    "make_mpl_space_component",
+    "make_plot_component",
+    "make_space_component",
+]
 
 
 def make_space_component(
@@ -57,6 +74,7 @@ def make_plot_component(
     measure: str | dict[str, str] | list[str] | tuple[str],
     post_process: Callable | None = None,
     backend: str = "matplotlib",
+    page: int = 0,
     **plot_drawing_kwargs,
 ):
     """Create a plotting function for a specified measure using the specified backend.
@@ -65,18 +83,20 @@ def make_plot_component(
         measure (str | dict[str, str] | list[str] | tuple[str]): Measure(s) to plot.
         post_process: a user-specified callable to do post-processing called with the Axes instance.
         backend: the backend to use {"matplotlib", "altair"}
+        page: Page number where the plot should be displayed (default 0).
         plot_drawing_kwargs: additional keyword arguments to pass onto the backend specific function for making a plotting component
 
-    Notes:
-        altair plotting backend is not yet implemented and planned for mesa 3.1.
-
     Returns:
-        function: A function that creates a plot component
+        (function, page): A tuple of a function and page number that creates a plot component on that specific page.
     """
     if backend == "matplotlib":
-        return make_mpl_plot_component(measure, post_process, **plot_drawing_kwargs)
+        return make_mpl_plot_component(
+            measure, post_process, page, **plot_drawing_kwargs
+        )
     elif backend == "altair":
-        raise NotImplementedError("altair line plots are not yet implemented")
+        return make_altair_plot_component(
+            measure, post_process, page, **plot_drawing_kwargs
+        )
     else:
         raise ValueError(
             f"unknown backend {backend}, must be one of matplotlib, altair"

mesa/visualization/components/altair_components.py

@@ -1,6 +1,7 @@
 """Altair based solara components for visualization mesa spaces."""
 
 import warnings
+from collections.abc import Callable
 
 import altair as alt
 import numpy as np
@@ -448,3 +449,88 @@ def chart_property_layers(space, propertylayer_portrayal, chart_width, chart_hei
                 f"PropertyLayer {layer_name} portrayal must include 'color' or 'colormap'."
             )
     return base, bar_chart
+
+
+def make_altair_plot_component(
+    measure: str | dict[str, str] | list[str] | tuple[str],
+    post_process: Callable | None = None,
+    page: int = 0,
+    grid=False,
+):
+    """Create a plotting function for a specified measure.
+
+    Args:
+        measure (str | dict[str, str] | list[str] | tuple[str]): Measure(s) to plot.
+        post_process: a user-specified callable to do post-processing called with the Axes instance.
+        page: Page number where the plot should be displayed.
+        grid: Bool to draw grid or not.
+
+    Returns:
+        (function, page): A tuple of a function that creates a PlotAltair component and a page number.
+    """
+
+    def MakePlotAltair(model):
+        return PlotAltair(model, measure, post_process=post_process, grid=grid)
+
+    return (MakePlotAltair, page)
+
+
+@solara.component
+def PlotAltair(model, measure, post_process: Callable | None = None, grid=False):
+    """Create an Altair-based plot for a measure or measures.
+
+    Args:
+        model (mesa.Model): The model instance.
+        measure (str | dict[str, str] | list[str] | tuple[str]): Measure(s) to plot.
+            If a dict is given, keys are measure names and values are colors.
+        post_process: A user-specified callable for post-processing, called
+            with the Altair Chart instance.
+        grid: Bool to draw grid or not.
+
+    Returns:
+        solara.FigureAltair: A component for rendering the plot.
+    """
+    update_counter.get()
+    df = model.datacollector.get_model_vars_dataframe().reset_index()
+    df = df.rename(columns={"index": "Step"})
+
+    y_title = "Value"
+    if isinstance(measure, str):
+        measures_to_plot = [measure]
+        y_title = measure
+    elif isinstance(measure, list | tuple):
+        measures_to_plot = list(measure)
+    elif isinstance(measure, dict):
+        measures_to_plot = list(measure.keys())
+
+    df_long = df.melt(
+        id_vars=["Step"],
+        value_vars=measures_to_plot,
+        var_name="Measure",
+        value_name="Value",
+    )
+
+    chart = (
+        alt.Chart(df_long)
+        .mark_line()
+        .encode(
+            x=alt.X("Step:Q", axis=alt.Axis(tickMinStep=1, title="Step", grid=grid)),
+            y=alt.Y("Value:Q", axis=alt.Axis(title=y_title, grid=grid)),
+            tooltip=["Step", "Measure", "Value"],
+        )
+        .properties(width=450, height=350)
+        .interactive()
+    )
+
+    if len(measures_to_plot) > 0:
+        color_args = {}
+        if isinstance(measure, dict):
+            color_args["scale"] = alt.Scale(
+                domain=list(measure.keys()), range=list(measure.values())
+            )
+        chart = chart.encode(color=alt.Color("Measure:N", **color_args))
+
+    if post_process is not None:
+        chart = post_process(chart)
+
+    return solara.FigureAltair(chart)

mesa/visualization/components/matplotlib_components.py

@@ -107,6 +107,7 @@ def make_plot_measure(*args, **kwargs):  # noqa: D103
 def make_mpl_plot_component(
     measure: str | dict[str, str] | list[str] | tuple[str],
     post_process: Callable | None = None,
+    page: int = 0,
     save_format="png",
 ):
     """Create a plotting function for a specified measure.
@@ -114,10 +115,11 @@ def make_mpl_plot_component(
     Args:
         measure (str | dict[str, str] | list[str] | tuple[str]): Measure(s) to plot.
         post_process: a user-specified callable to do post-processing called with the Axes instance.
+        page: Page number where the plot should be displayed.
         save_format: save format of figure in solara backend
 
     Returns:
-        function: A function that creates a PlotMatplotlib component.
+        (function, page): A tuple of a function that creates a PlotMatplotlib component and a page number.
     """
 
     def MakePlotMatplotlib(model):
@@ -125,7 +127,7 @@ def make_mpl_plot_component(
             model, measure, post_process=post_process, save_format=save_format
         )
 
-    return MakePlotMatplotlib
+    return (MakePlotMatplotlib, page)
 
 
 @solara.component