move solara component creations into solara_viz with backend option

wang-boyu · wang-boyu · commit 527c0239866d · 2024-11-03T16:20:18.000-05:00
diff --git a/mesa/visualization/__init__.py b/mesa/visualization/__init__.py
@@ -1,8 +1,6 @@
 """Solara based visualization for Mesa models."""
 
-from .components.altair import make_space_altair
-from .components.matplotlib import make_plot_component, make_space_component
-from .solara_viz import JupyterViz, SolaraViz
+from .solara_viz import JupyterViz, SolaraViz, make_plot_component, make_space_component
 from .UserParam import Slider
 
 __all__ = [
diff --git a/mesa/visualization/components/altair.py b/mesa/visualization/components/altair.py
@@ -12,18 +12,6 @@
 from mesa.visualization.utils import update_counter
 
 
-def make_space_altair(agent_portrayal=None):  # noqa: D103
-    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):  # noqa: D103
     update_counter.get()
diff --git a/mesa/visualization/components/matplotlib.py b/mesa/visualization/components/matplotlib.py
@@ -1,6 +1,5 @@
 """Matplotlib based solara components for visualization MESA spaces and plots."""
 
-import warnings
 from collections.abc import Callable
 
 import matplotlib.pyplot as plt
@@ -31,54 +30,6 @@
 Network = NetworkGrid | mesa.experimental.cell_space.Network
 
 
-def make_space_matplotlib(*args, **kwargs):  # noqa: D103
-    warnings.warn(
-        "make_space_matplotlib has been renamed to make_space_component",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-    return make_space_component(*args, **kwargs)
-
-
-def make_space_component(
-    agent_portrayal: Callable | None = None,
-    propertylayer_portrayal: dict | None = None,
-    post_process: Callable | None = None,
-    **space_drawing_kwargs,
-):
-    """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,
@@ -107,39 +58,6 @@ def SpaceMatplotlib(
     )
 
 
-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_plot_component(*args, **kwargs)
-
-
-def make_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,
diff --git a/mesa/visualization/solara_viz.py b/mesa/visualization/solara_viz.py
@@ -25,20 +25,131 @@
 
 import asyncio
 import copy
+import warnings
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Literal
 
 import reacton.core
 import solara
 
-import mesa.visualization.components.altair as components_altair
+from mesa.visualization.components.altair import SpaceAltair
+from mesa.visualization.components.matplotlib import PlotMatplotlib, SpaceMatplotlib
 from mesa.visualization.UserParam import Slider
 from mesa.visualization.utils import force_update, update_counter
 
 if TYPE_CHECKING:
     from mesa.model import Model
 
 
+def make_space_matplotlib(*args, **kwargs):  # noqa: D103
+    warnings.warn(
+        "make_space_matplotlib has been renamed to make_space_component",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return make_space_component(*args, **kwargs)
+
+
+def make_space_altair(agent_portrayal=None):  # noqa: D103
+    warnings.warn(
+        "make_space_altair has been renamed to make_space_component with backend='altair'",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    if agent_portrayal is None:
+
+        def agent_portrayal(a):
+            return {"id": a.unique_id}
+
+    def MakeSpaceAltair(model):
+        return SpaceAltair(model, agent_portrayal)
+
+    return MakeSpaceAltair
+
+
+def make_space_component(
+    agent_portrayal: Callable | None = None,
+    propertylayer_portrayal: dict | None = None,
+    post_process: Callable | None = None,
+    backend: Literal["matplotlib", "altair"] = "matplotlib",
+    **space_drawing_kwargs,
+):
+    """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 for rendering the space. Can be "matplotlib" or "altair".
+        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,
+        )
+
+    def MakeSpaceAltair(model):
+        return SpaceAltair(model, agent_portrayal)
+
+    match backend:
+        case "matplotlib":
+            return MakeSpaceMatplotlib
+        case "altair":
+            return MakeSpaceAltair
+        case _:
+            raise ValueError(f"Invalid backend: {backend}")
+
+
+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_plot_component(*args, **kwargs)
+
+
+def make_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 SolaraViz(
     model: Model | solara.Reactive[Model],
@@ -63,7 +174,7 @@ def SolaraViz(
         components (list[solara.component] | Literal["default"], optional): List of solara
             components or functions that return a solara component.
             These components are used to render different parts of the model visualization.
-            Defaults to "default", which uses the default Altair space visualization.
+            Defaults to "default", which uses the default Matplotlib space visualization.
         play_interval (int, optional): Interval for playing the model steps in milliseconds.
             This controls the speed of the model's automatic stepping. Defaults to 100 ms.
         model_params (dict, optional): Parameters for (re-)instantiating a model.
@@ -87,7 +198,7 @@ def SolaraViz(
           value results in faster stepping, while a higher value results in slower stepping.
     """
     if components == "default":
-        components = [components_altair.make_space_altair()]
+        components = [make_space_component()]
 
     # Convert model to reactive
     if not isinstance(model, solara.Reactive):
diff --git a/tests/test_solara_viz.py b/tests/test_solara_viz.py
@@ -8,8 +8,12 @@
 import mesa
 import mesa.visualization.components.altair
 import mesa.visualization.components.matplotlib
-from mesa.visualization.components.matplotlib import make_space_component
-from mesa.visualization.solara_viz import Slider, SolaraViz, UserInputs
+from mesa.visualization.solara_viz import (
+    Slider,
+    SolaraViz,
+    UserInputs,
+    make_space_component,
+)
 
 
 class TestMakeUserInput(unittest.TestCase):  # noqa: D101
@@ -87,11 +91,9 @@ def Test(user_params):
 
 
 def test_call_space_drawer(mocker):  # noqa: D103
-    mock_space_matplotlib = mocker.spy(
-        mesa.visualization.components.matplotlib, "SpaceMatplotlib"
-    )
+    mock_space_matplotlib = mocker.spy(mesa.visualization.solara_viz, "SpaceMatplotlib")
 
-    mock_space_altair = mocker.spy(mesa.visualization.components.altair, "SpaceAltair")
+    mock_space_altair = mocker.spy(mesa.visualization.solara_viz, "SpaceAltair")
 
     model = mesa.Model()
     mocker.patch.object(mesa.Model, "__init__", return_value=None)
@@ -113,8 +115,8 @@ def test_call_space_drawer(mocker):  # noqa: D103
     mock_space_matplotlib.reset_mock()
     solara.render(SolaraViz(model))
     # should call default method with class instance and agent portrayal
-    assert mock_space_matplotlib.call_count == 0
-    assert mock_space_altair.call_count > 0
+    assert mock_space_matplotlib.call_count > 0
+    assert mock_space_altair.call_count == 0
 
     # specify a custom space method
     class AltSpace:
