implement BackendWrapper registry for unified usage across backends.

Author: r3kste

pyproject.toml

@@ -291,6 +291,7 @@ module = [
     "optimagic.shared.process_user_function",
 
     "optimagic.visualization",
+    "optimagic.visualization.backends",
     "optimagic.visualization.convergence_plot",
     "optimagic.visualization.deviation_plot",
     "optimagic.visualization.history_plots",

src/optimagic/visualization/backends.py

@@ -0,0 +1,108 @@
+import abc
+from dataclasses import dataclass
+from typing import Any
+
+import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+
+
+@dataclass(frozen=True)
+class PlotConfig:
+    template: str
+    plotly_legend: dict[str, Any]
+    matplotlib_legend: dict[str, Any]
+
+
+class BackendWrapper(abc.ABC):
+    def __init__(self, plot_config):
+        self.plot_config = plot_config
+
+    @abc.abstractmethod
+    def create_figure(self):
+        pass
+
+    @abc.abstractmethod
+    def lineplot(self, **kwargs):
+        pass
+
+    @abc.abstractmethod
+    def post_plot(self, **kwargs):
+        pass
+
+    @abc.abstractmethod
+    def return_obj(self):
+        pass
+
+
+class BackendRegistry:
+    _registry: dict[str, BackendWrapper] = {}
+
+    @classmethod
+    def register(cls, backend_name):
+        def decorator(backend_wrapper):
+            cls._registry[backend_name] = backend_wrapper
+            return backend_wrapper
+
+        return decorator
+
+    @classmethod
+    def get_backend_wrapper(cls, backend_name):
+        if backend_name not in cls._registry:
+            raise ValueError(
+                f"Backend '{backend_name}' is not supported. "
+                f"Supported backends are: {', '.join(cls._registry.keys())}."
+            )
+        return cls._registry.get(backend_name)
+
+
+@BackendRegistry.register("plotly")
+class PlotlyBackend(BackendWrapper):
+    def __init__(self, plot_config):
+        super().__init__(plot_config)
+        self.fig = self.create_figure()
+
+    def create_figure(self):
+        fig = go.Figure()
+        return fig
+
+    def lineplot(self, *, x, y, color, name=None, plotly_scatter_kws=None, **kwargs):
+        if plotly_scatter_kws is None:
+            plotly_scatter_kws = {}
+
+        trace = go.Scatter(
+            x=x, y=y, mode="lines", line_color=color, name=name, **plotly_scatter_kws
+        )
+        self.fig.add_trace(trace)
+
+    def post_plot(self, *, xlabel=None, ylabel=None, **kwargs):
+        self.fig.update_layout(
+            template=self.plot_config.template,
+            xaxis_title_text=xlabel,
+            yaxis_title_text=ylabel,
+            legend=self.plot_config.plotly_legend,
+        )
+
+    def return_obj(self):
+        return self.fig
+
+
+@BackendRegistry.register("matplotlib")
+class MatplotlibBackend(BackendWrapper):
+    def __init__(self, plot_config):
+        super().__init__(plot_config)
+        self.fig, self.ax = self.create_figure()
+
+    def create_figure(self):
+        plt.style.use(self.plot_config.template)
+        fig, ax = plt.subplots()
+        return fig, ax
+
+    def lineplot(self, *, x, y, color, name=None, **kwargs):
+        self.ax.plot(x, y, color=color, label=name)
+
+    def post_plot(self, *, xlabel=None, ylabel=None, **kwargs):
+        self.ax.set(xlabel=xlabel, ylabel=ylabel)
+        self.ax.legend(**self.plot_config.matplotlib_legend)
+
+    def return_obj(self):
+        return self.fig

src/optimagic/visualization/history_plots.py

@@ -1,28 +1,27 @@
 import inspect
-import itertools
 from pathlib import Path
 from typing import Any
 
-import matplotlib as mpl
-import matplotlib.pyplot as plt
 import numpy as np
 import plotly.graph_objects as go
 from pybaum import leaf_names, tree_flatten, tree_just_flatten, tree_unflatten
 
-from optimagic.config import PLOT_DEFAULTS, PLOTLY_TEMPLATE
+from optimagic.config import PLOTLY_TEMPLATE
 from optimagic.logging.logger import LogReader, SQLiteLogOptions
 from optimagic.optimization.algorithm import Algorithm
 from optimagic.optimization.history import History
 from optimagic.optimization.optimize_result import OptimizeResult
 from optimagic.parameters.tree_registry import get_registry
 from optimagic.typing import Direction
+from optimagic.visualization.backends import BackendRegistry, PlotConfig
+from optimagic.visualization.plotting_utilities import get_palette, get_template
 
 
 def criterion_plot(
     results,
     names=None,
-    backend="plotly",
     max_evaluations=None,
+    backend="plotly",
     template=None,
     palette=None,
     stack_multistart=False,
@@ -36,8 +35,8 @@ def criterion_plot(
             dict of) optimization results with collected history. If dict, then the
             key is used as the name in a legend.
         names (Union[List[str], str]): Names corresponding to res or entries in res.
-        backend (str): The backend to use for plotting. Default is "plotly".
         max_evaluations (int): Clip the criterion history after that many entries.
+        backend (str): The backend to use for plotting. Default is "plotly".
         template (str): The template for the figure. Default is "plotly_white".
         palette (Union[List[str], str]): The coloring palette for traces. Default is
             "qualitative.Plotly".
@@ -50,7 +49,7 @@ def criterion_plot(
             optimization are visualized. Default is False.
 
     Returns:
-        Figure object returned by the chosen backend.
+        Native figure object returned by the chosen backend.
 
     """
     # ==================================================================================
@@ -59,16 +58,8 @@ def criterion_plot(
 
     results = _harmonize_inputs_to_dict(results, names)
 
-    if template is None:
-        template = PLOT_DEFAULTS[backend]["template"]
-    if palette is None:
-        palette = PLOT_DEFAULTS[backend]["palette"]
-
-    if isinstance(palette, mpl.colors.Colormap):
-        palette = [palette(i) for i in range(palette.N)]
-    if not isinstance(palette, list):
-        palette = [palette]
-    palette = itertools.cycle(palette)
+    template = get_template(backend, template)
+    palette = get_palette(backend, palette)
 
     fun_or_monotone_fun = "monotone_fun" if monotone else "fun"
 
@@ -98,15 +89,22 @@ def criterion_plot(
     # Create figure
     # ==================================================================================
 
-    fig, plot_func, label_func = _get_plot_backend(backend)
-
-    plot_multistart = (
-        len(data) == 1 and data[0]["is_multistart"] and not stack_multistart
+    plot_config = PlotConfig(
+        template=template,
+        plotly_legend={"yanchor": "top", "xanchor": "right", "y": 0.95, "x": 0.95},
+        matplotlib_legend={"loc": "upper right"},
     )
 
+    _backend_wrapper = BackendRegistry.get_backend_wrapper(backend)
+    fig = _backend_wrapper(plot_config)
+
     # ==================================================================================
     # Plot multistart paths
 
+    plot_multistart = (
+        len(data) == 1 and data[0]["is_multistart"] and not stack_multistart
+    )
+
     if plot_multistart:
         scatter_kws = {
             "connectgaps": True,
@@ -119,8 +117,7 @@ def criterion_plot(
             if max_evaluations is not None and len(history) > max_evaluations:
                 history = history[:max_evaluations]
 
-            plot_func(
-                fig,
+            fig.lineplot(
                 x=np.arange(len(history)),
                 y=history,
                 name=None,
@@ -144,30 +141,23 @@ def criterion_plot(
 
         scatter_kws = {
             "connectgaps": True,
-            "showlegend": not plot_multistart,
+            "showlegend": True,
         }
 
-        _color = next(palette)
-
-        plot_func(
-            fig,
+        fig.lineplot(
             x=np.arange(len(history)),
             y=history,
             name="best result" if plot_multistart else _data["name"],
-            color=_color,
+            color=next(palette),
             plotly_scatter_kws=scatter_kws,
         )
 
-    label_func(
-        fig,
-        template=template,
+    fig.post_plot(
         xlabel="No. of criterion evaluations",
         ylabel="Criterion value",
-        plotly_legend={"yanchor": "top", "xanchor": "right", "y": 0.95, "x": 0.95},
-        matplotlib_legend={"loc": "upper right"},
     )
 
-    return fig
+    return fig.return_obj()
 
 
 def _harmonize_inputs_to_dict(results, names):
@@ -461,54 +451,3 @@ def _get_stacked_local_histories(local_histories, direction, history=None):
         task=len(stacked["criterion"]) * [None],
         batches=list(range(len(stacked["criterion"]))),
     )
-
-
-def _get_plot_backend(backend):
-    backends = {
-        "plotly": (
-            go.Figure(),
-            _plot_plotly,
-            _label_plotly,
-        ),
-        "matplotlib": (
-            plt.subplots()[1],
-            _plot_matplotlib,
-            _label_matplotlib,
-        ),
-    }
-
-    if backend not in backends:
-        msg = (
-            f"Backend '{backend}' is not supported. "
-            f"Supported backends are: {', '.join(backends.keys())}."
-        )
-        raise ValueError(msg)
-
-    return backends[backend]
-
-
-def _plot_plotly(fig, *, x, y, name, color, plotly_scatter_kws, **kwargs):
-    trace = go.Scatter(
-        x=x, y=y, mode="lines", name=name, line_color=color, **plotly_scatter_kws
-    )
-    fig.add_trace(trace)
-    return fig
-
-
-def _label_plotly(fig, *, template, xlabel, ylabel, plotly_legend, **kwargs):
-    fig.update_layout(
-        template=template,
-        xaxis_title_text=xlabel,
-        yaxis_title_text=ylabel,
-        legend=plotly_legend,
-    )
-
-
-def _plot_matplotlib(ax, *, x, y, name, color, **kwargs):
-    ax.plot(x, y, label=name, color=color)
-    return ax
-
-
-def _label_matplotlib(ax, *, xlabel, ylabel, matplotlib_legend, **kwargs):
-    ax.set(xlabel=xlabel, ylabel=ylabel)
-    ax.legend(**matplotlib_legend)

src/optimagic/visualization/plotting_utilities.py

@@ -1,11 +1,12 @@
 import itertools
 from copy import deepcopy
 
+import matplotlib as mpl
 import numpy as np
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
-from optimagic.config import PLOTLY_TEMPLATE
+from optimagic.config import PLOT_DEFAULTS, PLOTLY_TEMPLATE
 
 
 def combine_plots(
@@ -328,3 +329,23 @@ def get_layout_kwargs(layout_kwargs, legend_kwargs, title_kwargs, template, show
     if layout_kwargs:
         default_kwargs.update(layout_kwargs)
     return default_kwargs
+
+
+def get_template(backend, template):
+    if template is None:
+        template = PLOT_DEFAULTS[backend]["template"]
+
+    return template
+
+
+def get_palette(backend, palette):
+    if palette is None:
+        palette = PLOT_DEFAULTS[backend]["palette"]
+
+    if isinstance(palette, mpl.colors.Colormap):
+        palette = [palette(i) for i in range(palette.N)]
+    if not isinstance(palette, list):
+        palette = [palette]
+    palette = itertools.cycle(palette)
+
+    return palette