Implementation of matplotlib backend for criterion_plot() (#599)

r3kste · web-flow · commit b0f7a8e50d9a · 2025-09-05T21:04:38.000+02:00
diff --git a/.tools/envs/testenv-linux.yml b/.tools/envs/testenv-linux.yml
@@ -19,6 +19,7 @@ dependencies:
   - numpy >= 2  # run, tests
   - pandas  # run, tests
   - plotly>=6.2  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/.tools/envs/testenv-nevergrad.yml b/.tools/envs/testenv-nevergrad.yml
@@ -17,6 +17,7 @@ dependencies:
   - numpy >= 2  # run, tests
   - pandas  # run, tests
   - plotly>=6.2  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/.tools/envs/testenv-numpy.yml b/.tools/envs/testenv-numpy.yml
@@ -17,6 +17,7 @@ dependencies:
   - cloudpickle  # run, tests
   - joblib  # run, tests
   - plotly>=6.2  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/.tools/envs/testenv-others.yml b/.tools/envs/testenv-others.yml
@@ -17,6 +17,7 @@ dependencies:
   - numpy >= 2  # run, tests
   - pandas  # run, tests
   - plotly>=6.2  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/.tools/envs/testenv-pandas.yml b/.tools/envs/testenv-pandas.yml
@@ -17,6 +17,7 @@ dependencies:
   - cloudpickle  # run, tests
   - joblib  # run, tests
   - plotly>=6.2  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/.tools/envs/testenv-plotly.yml b/.tools/envs/testenv-plotly.yml
@@ -17,6 +17,7 @@ dependencies:
   - joblib  # run, tests
   - numpy >= 2  # run, tests
   - pandas  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/environment.yml b/environment.yml
@@ -21,6 +21,7 @@ dependencies:
   - numpy >= 2  # run, tests
   - pandas  # run, tests
   - plotly>=6.2  # run, tests
+  - matplotlib  # tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
diff --git a/pyproject.toml b/pyproject.toml
@@ -290,6 +290,7 @@ module = [
 
     "optimagic.visualization",
     "optimagic.visualization.convergence_plot",
+    "optimagic.visualization.backends",
     "optimagic.visualization.deviation_plot",
     "optimagic.visualization.history_plots",
     "optimagic.visualization.plotting_utilities",
@@ -346,6 +347,8 @@ module = [
     "plotly.graph_objects",
     "plotly.express",
     "plotly.subplots",
+    "matplotlib",
+    "matplotlib.pyplot",
     "cyipopt",
     "nlopt",
     "bokeh",
diff --git a/src/optimagic/config.py b/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
@@ -23,7 +37,7 @@ def _is_installed(module_name: str) -> bool:
 
 
 # ======================================================================================
-# Check Available Packages
+# Check Available Optimization Packages
 # ======================================================================================
 
 IS_PETSC4PY_INSTALLED = _is_installed("petsc4py")
@@ -42,6 +56,12 @@ def _is_installed(module_name: str) -> bool:
 IS_GRADIENT_FREE_OPTIMIZERS_INSTALLED = _is_installed("gradient_free_optimizers")
 IS_PYGAD_INSTALLED = _is_installed("pygad")
 
+# ======================================================================================
+# Check Available Visualization Packages
+# ======================================================================================
+
+IS_MATPLOTLIB_INSTALLED = _is_installed("matplotlib")
+
 
 # ======================================================================================
 # Check if pandas version is newer or equal to version 2.1.0
diff --git a/src/optimagic/exceptions.py b/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)
diff --git a/src/optimagic/visualization/backends.py b/src/optimagic/visualization/backends.py
@@ -0,0 +1,170 @@
+from typing import Any, Literal, Protocol, runtime_checkable
+
+import plotly.graph_objects as go
+
+from optimagic.config import IS_MATPLOTLIB_INSTALLED
+from optimagic.exceptions import InvalidPlottingBackendError, NotInstalledError
+from optimagic.visualization.plotting_utilities import LineData
+
+if IS_MATPLOTLIB_INSTALLED:
+    import matplotlib.pyplot as plt
+
+
+@runtime_checkable
+class LinePlotFunction(Protocol):
+    def __call__(
+        self,
+        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,
+    ) -> Any: ...
+
+
+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:
+    if template is None:
+        template = "plotly"
+
+    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.Axes":
+    if template is None:
+        template = "default"
+
+    with plt.style.context(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,
+            )
+
+        ax.set(title=title, xlabel=xlabel, ylabel=ylabel)
+        if legend_properties:
+            ax.legend(**legend_properties)
+
+    return ax
+
+
+BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION: dict[
+    str, tuple[bool, LinePlotFunction]
+] = {
+    "plotly": (True, _line_plot_plotly),
+    "matplotlib": (IS_MATPLOTLIB_INSTALLED, _line_plot_matplotlib),
+}
+
+
+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 plotting backend '{backend}'. "
+            f"Available backends: "
+            f"{', '.join(BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION.keys())}"
+        )
+        raise InvalidPlottingBackendError(msg)
+
+    _is_backend_available, _line_plot_backend_function = (
+        BACKEND_AVAILABILITY_AND_LINE_PLOT_FUNCTION[backend]
+    )
+
+    if not _is_backend_available:
+        msg = (
+            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)
+
+    fig = _line_plot_backend_function(
+        lines,
+        title=title,
+        xlabel=xlabel,
+        ylabel=ylabel,
+        template=template,
+        height=height,
+        width=width,
+        legend_properties=legend_properties,
+    )
+
+    return fig
diff --git a/src/optimagic/visualization/history_plots.py b/src/optimagic/visualization/history_plots.py
diff --git a/src/optimagic/visualization/plotting_utilities.py b/src/optimagic/visualization/plotting_utilities.py
diff --git a/tests/optimagic/visualization/test_backends.py b/tests/optimagic/visualization/test_backends.py
diff --git a/tests/optimagic/visualization/test_history_plots.py b/tests/optimagic/visualization/test_history_plots.py
