Refactor to functional approach for backend plotting. Use hardcoded default palette.

Author: r3kste

src/optimagic/config.py

@@ -11,6 +11,20 @@
 PLOTLY_TEMPLATE = "simple_white"
 PLOTLY_PALETTE = px.colors.qualitative.Set2
 
+# The hex strings are obtained from the Plotly D3 qualitative palette.
+DEFAULT_PALETTE = [
+    "#1F77B4",
+    "#FF7F0E",
+    "#2CA02C",
+    "#D62728",
+    "#9467BD",
+    "#8C564B",
+    "#E377C2",
+    "#7F7F7F",
+    "#BCBD22",
+    "#17BECF",
+]
+
 DEFAULT_N_CORES = 1
 
 CRITERION_PENALTY_SLOPE = 0.1

src/optimagic/visualization/backends.py

@@ -1,7 +1,5 @@
-import abc
-from typing import Any
+from typing import Any, Literal
 
-import plotly.express as px
 import plotly.graph_objects as go
 
 from optimagic.config import IS_MATPLOTLIB_INSTALLED
@@ -20,126 +18,138 @@
         plt.ioff()
 
 
-class PlotBackend(abc.ABC):
-    is_available: bool
-    default_template: str
-
-    @classmethod
-    @abc.abstractmethod
-    def get_default_palette(cls) -> list:
-        pass
-
-    @abc.abstractmethod
-    def __init__(self, template: str | None):
-        if template is None:
-            template = self.default_template
-
-        self.template = template
-        self.figure: Any = None
-
-    @abc.abstractmethod
-    def add_lines(self, lines: list[LineData]) -> None:
-        pass
-
-    @abc.abstractmethod
-    def set_labels(self, xlabel: str | None = None, ylabel: str | None = None) -> None:
-        pass
-
-    @abc.abstractmethod
-    def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
-        pass
-
-
-class PlotlyBackend(PlotBackend):
-    is_available: bool = True
-    default_template: str = "simple_white"
-
-    @classmethod
-    def get_default_palette(cls) -> list:
-        return px.colors.qualitative.Set2
-
-    def __init__(self, template: str | None):
-        super().__init__(template)
-        self._fig = go.Figure()
-        self._fig.update_layout(template=self.template)
-        self.figure = self._fig
-
-    def add_lines(self, lines: list[LineData]) -> None:
-        for line in lines:
-            trace = go.Scatter(
-                x=line.x,
-                y=line.y,
-                name=line.name,
-                mode="lines",
-                line_color=line.color,
-                showlegend=line.show_in_legend,
-                connectgaps=True,
-            )
-            self._fig.add_trace(trace)
-
-    def set_labels(self, xlabel: str | None = None, ylabel: str | None = None) -> None:
-        self._fig.update_layout(xaxis_title_text=xlabel, yaxis_title_text=ylabel)
-
-    def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
-        self._fig.update_layout(legend=legend_properties)
-
-
-class MatplotlibBackend(PlotBackend):
-    is_available: bool = IS_MATPLOTLIB_INSTALLED
-    default_template: str = "default"
-
-    @classmethod
-    def get_default_palette(cls) -> list:
-        return [mpl.colormaps["Set2"](i) for i in range(mpl.colormaps["Set2"].N)]
-
-    def __init__(self, template: str | None):
-        super().__init__(template)
-        plt.style.use(self.template)
-        self._fig, self._ax = plt.subplots()
-        self.figure = self._fig
-
-    def add_lines(self, lines: list[LineData]) -> None:
-        for line in lines:
-            self._ax.plot(
-                line.x,
-                line.y,
-                color=line.color,
-                label=line.name if line.show_in_legend else None,
-            )
+def _line_plot_plotly(
+    lines: list[LineData],
+    *,
+    title: str | None,
+    xlabel: str | None,
+    ylabel: str | None,
+    template: str | None,
+    height: int | None,
+    width: int | None,
+    legend_properties: dict[str, Any] | None,
+) -> go.Figure:
+    fig = go.Figure()
+
+    for line in lines:
+        trace = go.Scatter(
+            x=line.x,
+            y=line.y,
+            name=line.name,
+            line_color=line.color,
+            mode="lines",
+        )
+        fig.add_trace(trace)
+
+    fig.update_layout(
+        title=title,
+        xaxis_title=xlabel,
+        yaxis_title=ylabel,
+        template=template,
+        height=height,
+        width=width,
+    )
+
+    if legend_properties:
+        fig.update_layout(legend=legend_properties)
+
+    return fig
+
+
+def _line_plot_matplotlib(
+    lines: list[LineData],
+    *,
+    title: str | None,
+    xlabel: str | None,
+    ylabel: str | None,
+    template: str | None,
+    height: int | None,
+    width: int | None,
+    legend_properties: dict[str, Any] | None,
+) -> "plt.Figure":
+    if template is not None:
+        plt.style.use(template)
+    fig, ax = plt.subplots(figsize=(width, height) if width and height else None)
+
+    for line in lines:
+        ax.plot(
+            line.x,
+            line.y,
+            label=line.name if line.show_in_legend else None,
+            color=line.color,
+        )
 
-    def set_labels(self, xlabel: str | None = None, ylabel: str | None = None) -> None:
-        self._ax.set(xlabel=xlabel, ylabel=ylabel)
+    ax.set(title=title, xlabel=xlabel, ylabel=ylabel)
+    if legend_properties:
+        ax.legend(**legend_properties)
 
-    def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
-        self._ax.legend(**legend_properties)
+    return fig
 
 
-PLOT_BACKEND_CLASSES = {
-    "plotly": PlotlyBackend,
-    "matplotlib": MatplotlibBackend,
+BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION = {
+    "plotly": (True, _line_plot_plotly),
+    "matplotlib": (IS_MATPLOTLIB_INSTALLED, _line_plot_matplotlib),
 }
 
 
-def get_plot_backend_class(backend_name: str) -> type[PlotBackend]:
-    if backend_name not in PLOT_BACKEND_CLASSES:
+def line_plot(
+    lines: list[LineData],
+    backend: Literal["plotly", "matplotlib"] = "plotly",
+    *,
+    title: str | None = None,
+    xlabel: str | None = None,
+    ylabel: str | None = None,
+    template: str | None = None,
+    height: int | None = None,
+    width: int | None = None,
+    legend_properties: dict[str, Any] | None = None,
+) -> Any:
+    """Create a line plot corresponding to the specified backend.
+
+    Args:
+        lines: List of objects each containing data for a line in the plot.
+        backend: The backend to use for plotting.
+        title: Title of the plot.
+        xlabel: Label for the x-axis.
+        ylabel: Label for the y-axis.
+        template: Backend-specific template for styling the plot.
+        height: Height of the plot (in pixels).
+        width: Width of the plot (in pixels).
+        legend_properties: Backend-specific properties for the legend.
+
+    Returns:
+        A figure object corresponding to the specified backend.
+
+    """
+    if backend not in BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION:
         msg = (
-            f"Invalid backend name '{backend_name}'. "
-            f"Supported backends are: {', '.join(PLOT_BACKEND_CLASSES.keys())}."
+            f"Invalid plotting backend '{backend}'. "
+            f"Available backends: "
+            f"{', '.join(BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION.keys())}"
         )
         raise InvalidPlottingBackendError(msg)
 
-    return _get_backend_if_installed(backend_name)
-
-
-def _get_backend_if_installed(backend_name: str) -> type[PlotBackend]:
-    plot_cls = PLOT_BACKEND_CLASSES[backend_name]
+    _is_backend_available, _line_plot_backend_function = (
+        BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION[backend]
+    )
 
-    if not plot_cls.is_available:
+    if not _is_backend_available:
         msg = (
-            f"The '{backend_name}' backend is not installed. "
-            f"Install the package using either 'pip install {backend_name}' or "
-            f"'conda install -c conda-forge {backend_name}'"
+            f"The {backend} backend is not installed. "
+            f"Install the package using either 'pip install {backend}' or "
+            f"'conda install -c conda-forge {backend}'"
         )
         raise NotInstalledError(msg)
 
-    return plot_cls
+    fig = _line_plot_backend_function(
+        lines,
+        title=title,
+        xlabel=xlabel,
+        ylabel=ylabel,
+        template=template,
+        height=height,
+        width=width,
+        legend_properties=legend_properties,
+    )
+
+    return fig

src/optimagic/visualization/history_plots.py

@@ -8,14 +8,14 @@
 import plotly.graph_objects as go
 from pybaum import leaf_names, tree_flatten, tree_just_flatten, tree_unflatten
 
-from optimagic.config import PLOTLY_TEMPLATE
+from optimagic.config import DEFAULT_PALETTE, 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 IterationHistory, PyTree
-from optimagic.visualization.backends import get_plot_backend_class
+from optimagic.visualization.backends import line_plot
 from optimagic.visualization.plotting_utilities import LineData, get_palette_cycle
 
 BACKEND_TO_CRITERION_PLOT_LEGEND_PROPERTIES: dict[str, dict[str, Any]] = {
@@ -40,7 +40,7 @@ def criterion_plot(
     max_evaluations: int | None = None,
     backend: Literal["plotly", "matplotlib"] = "plotly",
     template: str | None = None,
-    palette: list[str] | str | None = None,
+    palette: list[str] | str = DEFAULT_PALETTE,
     stack_multistart: bool = False,
     monotone: bool = False,
     show_exploration: bool = False,
@@ -66,16 +66,9 @@ def criterion_plot(
         The figure object containing the criterion plot.
 
     """
-    # ==================================================================================
-    # Get Plot Backend class
-
-    plot_cls = get_plot_backend_class(backend)
-
     # ==================================================================================
     # Process inputs
 
-    if palette is None:
-        palette = plot_cls.get_default_palette()
     palette_cycle = get_palette_cycle(palette)
 
     dict_of_optimize_results_or_paths = _harmonize_inputs_to_dict(results, names)
@@ -100,13 +93,16 @@ def criterion_plot(
     # ==================================================================================
     # Generate the figure
 
-    plot = plot_cls(template)
-
-    plot.add_lines(lines + multistart_lines)
-    plot.set_labels(xlabel="No. of criterion evaluations", ylabel="Criterion value")
-    plot.set_legend_properties(BACKEND_TO_CRITERION_PLOT_LEGEND_PROPERTIES[backend])
+    fig = line_plot(
+        lines=lines + multistart_lines,
+        backend=backend,
+        xlabel="No. of criterion evaluations",
+        ylabel="Criterion value",
+        template=template,
+        legend_properties=BACKEND_TO_CRITERION_PLOT_LEGEND_PROPERTIES[backend],
+    )
 
-    return plot.figure
+    return fig
 
 
 def _harmonize_inputs_to_dict(

tests/optimagic/visualization/test_history_plots.py

@@ -6,7 +6,6 @@
 from numpy.testing import assert_array_equal
 
 import optimagic as om
-from optimagic.exceptions import InvalidPlottingBackendError
 from optimagic.logging import SQLiteLogOptions
 from optimagic.optimization.optimize import minimize
 from optimagic.parameters.bounds import Bounds
@@ -145,9 +144,6 @@ def test_criterion_plot_wrong_inputs():
     with pytest.raises(ValueError):
         criterion_plot(["bla", "bla"], names="blub")
 
-    with pytest.raises(InvalidPlottingBackendError):
-        criterion_plot("bla", backend="blub")
-
 
 @pytest.mark.parametrize("backend", ["plotly", "matplotlib"])
 def test_criterion_plot_different_backends(minimize_result, backend):