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

mesa/examples/basic/boltzmann_wealth_model/app.py

@@ -1,18 +1,14 @@
 from mesa.examples.basic.boltzmann_wealth_model.model import BoltzmannWealthModel
 from mesa.visualization import (
     SolaraViz,
-    make_plot_measure,
-    make_space_matplotlib,
+    make_plot_component,
+    make_space_component,
 )
 
 
 def agent_portrayal(agent):
-    size = 10
-    color = "tab:red"
-    if agent.wealth > 0:
-        size = 50
-        color = "tab:blue"
-    return {"size": size, "color": color}
+    color = agent.wealth  # we are using a colormap to translate wealth to color
+    return {"color": color}
 
 
 model_params = {
@@ -28,6 +24,11 @@ model_params = {
     "height": 10,
 }
 
+
+def post_process(ax):
+    ax.get_figure().colorbar(ax.collections[0], label="wealth", ax=ax)
+
+
 # Create initial model instance
 model1 = BoltzmannWealthModel(50, 10, 10)
 
@@ -36,8 +37,11 @@ model1 = BoltzmannWealthModel(50, 10, 10)
 # Under the hood these are just classes that receive the model instance.
 # You can also author your own visualization elements, which can also be functions
 # that receive the model instance and return a valid solara component.
-SpaceGraph = make_space_matplotlib(agent_portrayal)
-GiniPlot = make_plot_measure("Gini")
+
+SpaceGraph = make_space_component(
+    agent_portrayal, cmap="viridis", vmin=0, vmax=10, post_process=post_process
+)
+GiniPlot = make_plot_component("Gini")
 
 # Create the SolaraViz page. This will automatically create a server and display the
 # visualization elements in a web browser.

mesa/examples/basic/conways_game_of_life/app.py

@@ -1,12 +1,22 @@
 from mesa.examples.basic.conways_game_of_life.model import ConwaysGameOfLife
 from mesa.visualization import (
     SolaraViz,
-    make_space_matplotlib,
+    make_space_component,
 )
 
 
 def agent_portrayal(agent):
-    return {"color": "white" if agent.state == 0 else "black"}
+    return {
+        "color": "white" if agent.state == 0 else "black",
+        "marker": "s",
+        "size": 25,
+    }
+
+
+def post_process(ax):
+    ax.set_aspect("equal")
+    ax.set_xticks([])
+    ax.set_yticks([])
 
 
 model_params = {
@@ -22,7 +32,9 @@ model1 = ConwaysGameOfLife(50, 50)
 # Under the hood these are just classes that receive the model instance.
 # You can also author your own visualization elements, which can also be functions
 # that receive the model instance and return a valid solara component.
-SpaceGraph = make_space_matplotlib(agent_portrayal)
+SpaceGraph = make_space_component(
+    agent_portrayal, post_process=post_process, draw_grid=False
+)
 
 
 # Create the SolaraViz page. This will automatically create a server and display the

mesa/examples/basic/schelling/app.py

@@ -4,8 +4,8 @@ from mesa.examples.basic.schelling.model import Schelling
 from mesa.visualization import (
     Slider,
     SolaraViz,
-    make_plot_measure,
-    make_space_matplotlib,
+    make_plot_component,
+    make_space_component,
 )
 
 
@@ -28,13 +28,13 @@ model_params = {
 
 model1 = Schelling(20, 20, 0.8, 0.2, 3)
 
-HappyPlot = make_plot_measure("happy")
+HappyPlot = make_plot_component({"happy": "tab:green"})
 
 page = SolaraViz(
     model1,
     components=[
-        make_space_matplotlib(agent_portrayal),
-        make_plot_measure("happy"),
+        make_space_component(agent_portrayal),
+        HappyPlot,
         get_happy_agents,
     ],
     model_params=model_params,

mesa/examples/basic/virus_on_network/app.py

@@ -1,44 +1,27 @@
 import math
 
 import solara
-from matplotlib.figure import Figure
-from matplotlib.ticker import MaxNLocator
 
 from mesa.examples.basic.virus_on_network.model import (
     State,
     VirusOnNetwork,
     number_infected,
 )
-from mesa.visualization import Slider, SolaraViz, make_space_matplotlib
-
+from mesa.visualization import (
+    Slider,
+    SolaraViz,
+    make_plot_component,
+    make_space_component,
+)
 
-def agent_portrayal(graph):
-    def get_agent(node):
-        return graph.nodes[node]["agent"][0]
 
-    edge_width = []
-    edge_color = []
-    for u, v in graph.edges():
-        agent1 = get_agent(u)
-        agent2 = get_agent(v)
-        w = 2
-        ec = "#e8e8e8"
-        if State.RESISTANT in (agent1.state, agent2.state):
-            w = 3
-            ec = "black"
-        edge_width.append(w)
-        edge_color.append(ec)
+def agent_portrayal(agent):
     node_color_dict = {
         State.INFECTED: "tab:red",
         State.SUSCEPTIBLE: "tab:green",
         State.RESISTANT: "tab:gray",
     }
-    node_color = [node_color_dict[get_agent(node).state] for node in graph.nodes()]
-    return {
-        "width": edge_width,
-        "edge_color": edge_color,
-        "node_color": node_color,
-    }
+    return {"color": node_color_dict[agent.state], "size": 10}
 
 
 def get_resistant_susceptible_ratio(model):
@@ -46,25 +29,9 @@ def get_resistant_susceptible_ratio(model):
     ratio_text = r"$\infty$" if ratio is math.inf else f"{ratio:.2f}"
     infected_text = str(number_infected(model))
 
-    return f"Resistant/Susceptible Ratio: {ratio_text}<br>Infected Remaining: {infected_text}"
-
-
-def make_plot(model):
-    # This is for the case when we want to plot multiple measures in 1 figure.
-    fig = Figure()
-    ax = fig.subplots()
-    measures = ["Infected", "Susceptible", "Resistant"]
-    colors = ["tab:red", "tab:green", "tab:gray"]
-    for i, m in enumerate(measures):
-        color = colors[i]
-        df = model.datacollector.get_model_vars_dataframe()
-        ax.plot(df.loc[:, m], label=m, color=color)
-    fig.legend()
-    # Set integer x axis
-    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
-    ax.set_xlabel("Step")
-    ax.set_ylabel("Number of Agents")
-    return solara.FigureMatplotlib(fig)
+    return solara.Markdown(
+        f"Resistant/Susceptible Ratio: {ratio_text}<br>Infected Remaining: {infected_text}"
+    )
 
 
 model_params = {
@@ -119,7 +86,18 @@ model_params = {
     ),
 }
 
-SpacePlot = make_space_matplotlib(agent_portrayal)
+
+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()
 
@@ -127,8 +105,8 @@ page = SolaraViz(
     model1,
     [
         SpacePlot,
-        make_plot,
-        # get_resistant_susceptible_ratio,  # TODO: Fix and uncomment
+        StatePlot,
+        get_resistant_susceptible_ratio,
     ],
     model_params=model_params,
     name="Virus Model",

mesa/space.py

@@ -1287,36 +1287,6 @@ class HexMultiGrid(_HexGrid, MultiGrid):
     """
 
 
-class HexGrid(HexSingleGrid):
-    """Hexagonal Grid: a Grid where neighbors are computed according to a hexagonal tiling of the grid.
-
-    Functions according to odd-q rules.
-    See http://www.redblobgames.com/grids/hexagons/#coordinates for more.
-
-    Properties:
-        width, height: The grid's width and height.
-        torus: Boolean which determines whether to treat the grid as a torus.
-    """
-
-    def __init__(self, width: int, height: int, torus: bool) -> None:
-        """Initializes a HexGrid, deprecated.
-
-        Args:
-            width: the width of the grid
-            height: the height of the grid
-            torus: whether the grid wraps
-        """
-        super().__init__(width, height, torus)
-        warn(
-            (
-                "HexGrid is being deprecated; use instead HexSingleGrid or HexMultiGrid "
-                "depending on your use case."
-            ),
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
-
 class ContinuousSpace:
     """Continuous space where each agent can have an arbitrary position.
 

mesa/visualization/__init__.py

@@ -1,7 +1,17 @@
-"""Solara based visualization for Mesa models."""
+"""Solara based visualization for Mesa models.
 
-from .components.altair import make_space_altair
-from .components.matplotlib import make_plot_measure, make_space_matplotlib
+.. note::
+    SolaraViz is experimental and still in active development for Mesa 3.0. While we attempt to minimize them, there might be API breaking changes between Mesa 3.0 and 3.1.
+
+    There won't be breaking changes between Mesa 3.0.x patch releases.
+"""
+
+from mesa.visualization.mpl_space_drawing import (
+    draw_space,
+)
+
+from .components import make_plot_component, make_space_component
+from .components.altair_components import make_space_altair
 from .solara_viz import JupyterViz, SolaraViz
 from .UserParam import Slider
 
@@ -10,6 +20,7 @@ __all__ = [
     "SolaraViz",
     "Slider",
     "make_space_altair",
-    "make_space_matplotlib",
-    "make_plot_measure",
+    "draw_space",
+    "make_plot_component",
+    "make_space_component",
 ]

mesa/visualization/components/__init__.py

@@ -0,0 +1,83 @@
+"""custom solara components."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from .altair_components import SpaceAltair, make_altair_space
+from .matplotlib_components import (
+    SpaceMatplotlib,
+    make_mpl_plot_component,
+    make_mpl_space_component,
+)
+
+
+def make_space_component(
+    agent_portrayal: Callable | None = None,
+    propertylayer_portrayal: dict | None = None,
+    post_process: Callable | None = None,
+    backend: str = "matplotlib",
+    **space_drawing_kwargs,
+) -> SpaceMatplotlib | SpaceAltair:
+    """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)
+        backend: the backend to use {"matplotlib", "altair"}
+        space_drawing_kwargs : additional keyword arguments to be passed on to the underlying backend specific space drawer function. See
+                               the functions for drawing the various spaces for the appropriate backend further details.
+
+
+    Returns:
+        function: A function that creates a space component
+    """
+    if backend == "matplotlib":
+        return make_mpl_space_component(
+            agent_portrayal,
+            propertylayer_portrayal,
+            post_process,
+            **space_drawing_kwargs,
+        )
+    elif backend == "altair":
+        return make_altair_space(
+            agent_portrayal,
+            propertylayer_portrayal,
+            post_process,
+            **space_drawing_kwargs,
+        )
+    else:
+        raise ValueError(
+            f"unknown backend {backend}, must be one of matplotlib, altair"
+        )
+
+
+def make_plot_component(
+    measure: str | dict[str, str] | list[str] | tuple[str],
+    post_process: Callable | None = None,
+    backend: str = "matplotlib",
+    **plot_drawing_kwargs,
+):
+    """Create a plotting function for a specified measure using the specified backend.
+
+    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.
+        backend: the backend to use {"matplotlib", "altair"}
+        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
+    """
+    if backend == "matplotlib":
+        return make_mpl_plot_component(measure, post_process, **plot_drawing_kwargs)
+    elif backend == "altair":
+        raise NotImplementedError("altair line plots are not yet implemented")
+    else:
+        raise ValueError(
+            f"unknown backend {backend}, must be one of matplotlib, altair"
+        )

mesa/visualization/components/{altair.py → altair_components.py}

@@ -1,6 +1,7 @@
 """Altair based solara components for visualization mesa spaces."""
 
 import contextlib
+import warnings
 
 import solara
 
@@ -12,7 +13,33 @@ from mesa.space import ContinuousSpace, _Grid
 from mesa.visualization.utils import update_counter
 
 
-def make_space_altair(agent_portrayal=None):  # noqa: D103
+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):
@@ -25,7 +52,12 @@ def make_space_altair(agent_portrayal=None):  # noqa: D103
 
 
 @solara.component
-def SpaceAltair(model, agent_portrayal, dependencies: list[any] | None = None):  # noqa: D103
+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:

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
+    )