Enhance availability check for backends. Fix issues with matplotlib in interactive mode.

Author: r3kste

src/optimagic/exceptions.py

@@ -79,6 +79,10 @@ class InvalidAlgoInfoError(OptimagicError):
     """Exception for invalid user provided algorithm information."""
 
 
+class InvalidPlottingBackendError(OptimagicError):
+    """Exception for invalid user provided plotting backend."""
+
+
 class StopOptimizationError(OptimagicError):
     def __init__(self, message, current_status):
         super().__init__(message)

src/optimagic/visualization/backends.py

@@ -5,17 +5,29 @@
 import plotly.graph_objects as go
 
 from optimagic.config import IS_MATPLOTLIB_INSTALLED
-from optimagic.exceptions import NotInstalledError
+from optimagic.exceptions import InvalidPlottingBackendError, NotInstalledError
 from optimagic.visualization.plotting_utilities import LineData
 
 if IS_MATPLOTLIB_INSTALLED:
     import matplotlib as mpl
     import matplotlib.pyplot as plt
 
+    # Handle the case where matplotlib is used in notebooks (inline backend)
+    # to ensure that interactive mode is disabled to avoid double plotting.
+    # (See: https://github.com/matplotlib/matplotlib/issues/26221)
+    if mpl.get_backend() == "module://matplotlib_inline.backend_inline":
+        plt.install_repl_displayhook()
+        plt.ioff()
+
 
 class PlotBackend(abc.ABC):
+    is_available: bool
     default_template: str
-    default_palette: list
+
+    @classmethod
+    @abc.abstractmethod
+    def get_default_palette(cls) -> list:
+        pass
 
     @abc.abstractmethod
     def __init__(self, template: str | None):
@@ -39,8 +51,12 @@ def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
 
 
 class PlotlyBackend(PlotBackend):
+    is_available: bool = True
     default_template: str = "simple_white"
-    default_palette: list = px.colors.qualitative.Set2
+
+    @classmethod
+    def get_default_palette(cls) -> list:
+        return px.colors.qualitative.Set2
 
     def __init__(self, template: str | None):
         super().__init__(template)
@@ -68,41 +84,39 @@ def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
         self._fig.update_layout(legend=legend_properties)
 
 
-if IS_MATPLOTLIB_INSTALLED:
+class MatplotlibBackend(PlotBackend):
+    is_available: bool = IS_MATPLOTLIB_INSTALLED
+    default_template: str = "default"
 
-    class MatplotlibBackend(PlotBackend):
-        default_template: str = "default"
-        default_palette: list = [
-            mpl.colormaps["Set2"](i) for i in range(mpl.colormaps["Set2"].N)
-        ]
+    @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 __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 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 set_labels(
-            self, xlabel: str | None = None, ylabel: str | None = None
-        ) -> None:
-            self._ax.set(xlabel=xlabel, ylabel=ylabel)
+    def set_labels(self, xlabel: str | None = None, ylabel: str | None = None) -> None:
+        self._ax.set(xlabel=xlabel, ylabel=ylabel)
 
-        def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
-            self._ax.legend(**legend_properties)
+    def set_legend_properties(self, legend_properties: dict[str, Any]) -> None:
+        self._ax.legend(**legend_properties)
 
 
 PLOT_BACKEND_CLASSES = {
     "plotly": PlotlyBackend,
-    "matplotlib": MatplotlibBackend if IS_MATPLOTLIB_INSTALLED else None,
+    "matplotlib": MatplotlibBackend,
 }
 
 
@@ -112,15 +126,15 @@ def get_plot_backend_class(backend_name: str) -> type[PlotBackend]:
             f"Invalid backend name '{backend_name}'. "
             f"Supported backends are: {', '.join(PLOT_BACKEND_CLASSES.keys())}."
         )
-        raise ValueError(msg)
+        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]
 
-    if plot_cls is None:
+    if not plot_cls.is_available:
         msg = (
             f"The '{backend_name}' backend is not installed. "
             f"Install the package using either 'pip install {backend_name}' or "

src/optimagic/visualization/history_plots.py

@@ -2,7 +2,7 @@
 import itertools
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import plotly.graph_objects as go
@@ -38,7 +38,7 @@ def criterion_plot(
     results: ResultOrPath | list[ResultOrPath] | dict[str, ResultOrPath],
     names: list[str] | str | None = None,
     max_evaluations: int | None = None,
-    backend: str = "plotly",
+    backend: Literal["plotly", "matplotlib"] = "plotly",
     template: str | None = None,
     palette: list[str] | str | None = None,
     stack_multistart: bool = False,
@@ -75,7 +75,7 @@ def criterion_plot(
     # Process inputs
 
     if palette is None:
-        palette = plot_cls.default_palette
+        palette = plot_cls.get_default_palette()
     palette_cycle = get_palette_cycle(palette)
 
     dict_of_optimize_results_or_paths = _harmonize_inputs_to_dict(results, names)

tests/optimagic/visualization/test_history_plots.py

@@ -6,6 +6,7 @@
 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
@@ -144,6 +145,15 @@ 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):
+    res = minimize_result[False][0]
+    criterion_plot(res, backend=backend)
+
 
 def test_harmonize_inputs_to_dict_single_result():
     res = minimize(fun=lambda x: x @ x, params=np.arange(5), algorithm="scipy_lbfgsb")