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

mesa/visualization/components/altair_components.py

@@ -0,0 +1,188 @@
+"""Altair based solara components for visualization mesa spaces."""
+
+import contextlib
+import warnings
+
+import solara
+
+with contextlib.suppress(ImportError):
+    import altair as alt
+
+from mesa.experimental.cell_space import DiscreteSpace, Grid
+from mesa.space import ContinuousSpace, _Grid
+from mesa.visualization.utils import update_counter
+
+
+def make_space_altair(*args, **kwargs):  # noqa: D103
+    warnings.warn(
+        "make_space_altair has been renamed to make_altair_space",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return make_altair_space(*args, **kwargs)
+
+
+def make_altair_space(
+    agent_portrayal, propertylayer_portrayal, post_process, **space_drawing_kwargs
+):
+    """Create an Altair-based space visualization component.
+
+    Args:
+        agent_portrayal: Function to portray agents.
+        propertylayer_portrayal: not yet implemented
+        post_process :not yet implemented
+        space_drawing_kwargs : not yet implemented
+
+    ``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.
+
+
+    Returns:
+        function: A function that creates a SpaceMatplotlib component
+    """
+    if agent_portrayal is None:
+
+        def agent_portrayal(a):
+            return {"id": a.unique_id}
+
+    def MakeSpaceAltair(model):
+        return SpaceAltair(model, agent_portrayal)
+
+    return MakeSpaceAltair
+
+
+@solara.component
+def SpaceAltair(model, agent_portrayal, dependencies: list[any] | None = None):
+    """Create an Altair-based space visualization component.
+
+    Returns:
+        a solara FigureAltair instance
+    """
+    update_counter.get()
+    space = getattr(model, "grid", None)
+    if space is None:
+        # Sometimes the space is defined as model.space instead of model.grid
+        space = model.space
+
+    chart = _draw_grid(space, agent_portrayal)
+    solara.FigureAltair(chart)
+
+
+def _get_agent_data_old__discrete_space(space, agent_portrayal):
+    """Format agent portrayal data for old-style discrete spaces.
+
+    Args:
+        space: the mesa.space._Grid instance
+        agent_portrayal: the agent portrayal callable
+
+    Returns:
+        list of dicts
+
+    """
+    all_agent_data = []
+    for content, (x, y) in space.coord_iter():
+        if not content:
+            continue
+        if not hasattr(content, "__iter__"):
+            # Is a single grid
+            content = [content]  # noqa: PLW2901
+        for agent in content:
+            # use all data from agent portrayal, and add x,y coordinates
+            agent_data = agent_portrayal(agent)
+            agent_data["x"] = x
+            agent_data["y"] = y
+            all_agent_data.append(agent_data)
+    return all_agent_data
+
+
+def _get_agent_data_new_discrete_space(space: DiscreteSpace, agent_portrayal):
+    """Format agent portrayal data for new-style discrete spaces.
+
+    Args:
+        space: the mesa.experiment.cell_space.Grid instance
+        agent_portrayal: the agent portrayal callable
+
+    Returns:
+        list of dicts
+
+    """
+    all_agent_data = []
+
+    for cell in space.all_cells:
+        for agent in cell.agents:
+            agent_data = agent_portrayal(agent)
+            agent_data["x"] = cell.coordinate[0]
+            agent_data["y"] = cell.coordinate[1]
+            all_agent_data.append(agent_data)
+    return all_agent_data
+
+
+def _get_agent_data_continuous_space(space: ContinuousSpace, agent_portrayal):
+    """Format agent portrayal data for continuous space.
+
+    Args:
+        space: the ContinuousSpace instance
+        agent_portrayal: the agent portrayal callable
+
+    Returns:
+        list of dicts
+    """
+    all_agent_data = []
+    for agent in space._agent_to_index:
+        agent_data = agent_portrayal(agent)
+        agent_data["x"] = agent.pos[0]
+        agent_data["y"] = agent.pos[1]
+        all_agent_data.append(agent_data)
+    return all_agent_data
+
+
+def _draw_grid(space, agent_portrayal):
+    match space:
+        case Grid():
+            all_agent_data = _get_agent_data_new_discrete_space(space, agent_portrayal)
+        case _Grid():
+            all_agent_data = _get_agent_data_old__discrete_space(space, agent_portrayal)
+        case ContinuousSpace():
+            all_agent_data = _get_agent_data_continuous_space(space, agent_portrayal)
+        case _:
+            raise NotImplementedError(
+                f"visualizing {type(space)} is currently not supported through altair"
+            )
+
+    invalid_tooltips = ["color", "size", "x", "y"]
+
+    x_y_type = "ordinal" if not isinstance(space, ContinuousSpace) else "nominal"
+
+    encoding_dict = {
+        # no x-axis label
+        "x": alt.X("x", axis=None, type=x_y_type),
+        # no y-axis label
+        "y": alt.Y("y", axis=None, type=x_y_type),
+        "tooltip": [
+            alt.Tooltip(key, type=alt.utils.infer_vegalite_type([value]))
+            for key, value in all_agent_data[0].items()
+            if key not in invalid_tooltips
+        ],
+    }
+    has_color = "color" in all_agent_data[0]
+    if has_color:
+        encoding_dict["color"] = alt.Color("color", type="nominal")
+    has_size = "size" in all_agent_data[0]
+    if has_size:
+        encoding_dict["size"] = alt.Size("size", type="quantitative")
+
+    chart = (
+        alt.Chart(
+            alt.Data(values=all_agent_data), encoding=alt.Encoding(**encoding_dict)
+        )
+        .mark_point(filled=True)
+        .properties(width=280, height=280)
+        # .configure_view(strokeOpacity=0)  # hide grid/chart lines
+    )
+    # This is the default value for the marker size, which auto-scales
+    # according to the grid area.
+    if not has_size:
+        length = min(space.width, space.height)
+        chart = chart.mark_point(size=30000 / length**2, filled=True)
+
+    return chart

mesa/visualization/components/matplotlib_components.py

@@ -0,0 +1,176 @@
+"""Matplotlib based solara components for visualization MESA spaces and plots."""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Callable
+
+import matplotlib.pyplot as plt
+import solara
+from matplotlib.figure import Figure
+
+from mesa.visualization.mpl_space_drawing import draw_space
+from mesa.visualization.utils import update_counter
+
+
+def make_space_matplotlib(*args, **kwargs):  # noqa: D103
+    warnings.warn(
+        "make_space_matplotlib has been renamed to make_mpl_space_component",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return make_mpl_space_component(*args, **kwargs)
+
+
+def make_mpl_space_component(
+    agent_portrayal: Callable | None = None,
+    propertylayer_portrayal: dict | None = None,
+    post_process: Callable | None = None,
+    **space_drawing_kwargs,
+) -> SpaceMatplotlib:
+    """Create a Matplotlib-based space visualization component.
+
+    Args:
+        agent_portrayal: Function to portray agents.
+        propertylayer_portrayal: Dictionary of PropertyLayer portrayal specifications
+        post_process : a callable that will be called with the Axes instance. Allows for fine tuning plots (e.g., control ticks)
+        space_drawing_kwargs : additional keyword arguments to be passed on to the underlying space drawer function. See
+                               the functions for drawing the various spaces for further details.
+
+    ``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.
+
+
+    Returns:
+        function: A function that creates a SpaceMatplotlib component
+    """
+    if agent_portrayal is None:
+
+        def agent_portrayal(a):
+            return {}
+
+    def MakeSpaceMatplotlib(model):
+        return SpaceMatplotlib(
+            model,
+            agent_portrayal,
+            propertylayer_portrayal,
+            post_process=post_process,
+            **space_drawing_kwargs,
+        )
+
+    return MakeSpaceMatplotlib
+
+
+@solara.component
+def SpaceMatplotlib(
+    model,
+    agent_portrayal,
+    propertylayer_portrayal,
+    dependencies: list[any] | None = None,
+    post_process: Callable | None = None,
+    **space_drawing_kwargs,
+):
+    """Create a Matplotlib-based space visualization component."""
+    update_counter.get()
+
+    space = getattr(model, "grid", None)
+    if space is None:
+        space = getattr(model, "space", None)
+
+    fig = Figure()
+    ax = fig.add_subplot()
+
+    draw_space(
+        space,
+        agent_portrayal,
+        propertylayer_portrayal=propertylayer_portrayal,
+        ax=ax,
+        **space_drawing_kwargs,
+    )
+
+    if post_process is not None:
+        post_process(ax)
+
+    solara.FigureMatplotlib(
+        fig, format="png", bbox_inches="tight", dependencies=dependencies
+    )
+
+
+def make_plot_measure(*args, **kwargs):  # noqa: D103
+    warnings.warn(
+        "make_plot_measure has been renamed to make_plot_component",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return make_mpl_plot_component(*args, **kwargs)
+
+
+def make_mpl_plot_component(
+    measure: str | dict[str, str] | list[str] | tuple[str],
+    post_process: Callable | None = None,
+    save_format="png",
+):
+    """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.
+        save_format: save format of figure in solara backend
+
+    Returns:
+        function: A function that creates a PlotMatplotlib component.
+    """
+
+    def MakePlotMatplotlib(model):
+        return PlotMatplotlib(
+            model, measure, post_process=post_process, save_format=save_format
+        )
+
+    return MakePlotMatplotlib
+
+
+@solara.component
+def PlotMatplotlib(
+    model,
+    measure,
+    dependencies: list[any] | None = None,
+    post_process: Callable | None = None,
+    save_format="png",
+):
+    """Create a Matplotlib-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.
+        dependencies (list[any] | None): Optional dependencies for the plot.
+        post_process: a user-specified callable to do post-processing called with the Axes instance.
+        save_format: format used for saving the figure.
+
+    Returns:
+        solara.FigureMatplotlib: A component for rendering the plot.
+    """
+    update_counter.get()
+    fig = Figure()
+    ax = fig.subplots()
+    df = model.datacollector.get_model_vars_dataframe()
+    if isinstance(measure, str):
+        ax.plot(df.loc[:, measure])
+        ax.set_ylabel(measure)
+    elif isinstance(measure, dict):
+        for m, color in measure.items():
+            ax.plot(df.loc[:, m], label=m, color=color)
+        ax.legend(loc="best")
+    elif isinstance(measure, list | tuple):
+        for m in measure:
+            ax.plot(df.loc[:, m], label=m)
+        ax.legend(loc="best")
+
+    if post_process is not None:
+        post_process(ax)
+
+    ax.set_xlabel("Step")
+    # Set integer x axis
+    ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True))
+    solara.FigureMatplotlib(
+        fig, format=save_format, bbox_inches="tight", dependencies=dependencies
+    )