fix(plots): export per-point colors and the plot title to matplotlib

bec_widgets/widgets/plots/toolbar_components/plot_export.py

@@ -2,6 +2,11 @@ from __future__ import annotations
 
 import traceback
 
+import numpy as np
+import pyqtgraph as pg
+from pyqtgraph import functions as fn
+from qtpy.QtCore import Qt
+
 from bec_widgets.utils.error_popups import SafeSlot, WarningPopupUtility
 from bec_widgets.utils.toolbars.actions import MaterialIconAction, SwitchableToolBarAction
 from bec_widgets.utils.toolbars.bundles import ToolbarBundle, ToolbarComponents
@@ -48,6 +53,184 @@ def plot_export_bundle(components: ToolbarComponents) -> ToolbarBundle:
     return bundle
 
 
+# pyqtgraph symbol names mapped to matplotlib markers, as in pyqtgraph's own MatplotlibExporter
+_SYMBOL_PG_TO_MPL = {
+    "o": "o",
+    "s": "s",
+    "t": "v",
+    "t1": "^",
+    "t2": ">",
+    "t3": "<",
+    "d": "d",
+    "+": "P",
+    "x": "X",
+    "p": "p",
+    "h": "h",
+    "star": "*",
+    "arrow_up": 6,
+    "arrow_right": 5,
+    "arrow_down": 7,
+    "arrow_left": 4,
+    "crosshair": "o",
+    "|": "|",
+    "_": "_",
+}
+
+
+def _is_per_point(spec) -> bool:
+    """Return True if a pen/brush option holds one entry per data point."""
+    return (
+        isinstance(spec, (list, np.ndarray))
+        and len(spec) > 0
+        and not isinstance(spec[0], (int, float, np.number))
+    )
+
+
+def _brush_colors(spec, default=None) -> list[tuple] | tuple | None:
+    """Convert a pyqtgraph brush option (single or per point) to matplotlib RGBA colors."""
+    if spec is None:
+        return default
+    if _is_per_point(spec):
+        return [_brush_colors(entry, default) for entry in spec]
+    brush = fn.mkBrush(spec)
+    if brush.style() == Qt.BrushStyle.NoBrush:
+        return "none"
+    return brush.color().getRgbF()
+
+
+def _pen_colors(spec, default=None) -> list[tuple] | tuple | None:
+    """Convert a pyqtgraph pen option (single or per point) to matplotlib RGBA colors."""
+    if spec is None:
+        return default
+    if _is_per_point(spec):
+        return [_pen_colors(entry, default) for entry in spec]
+    pen = fn.mkPen(spec)
+    if pen.style() == Qt.PenStyle.NoPen:
+        return "none"
+    return pen.color().getRgbF()
+
+
+def export_to_matplotlib(plot_item: pg.PlotItem):
+    """
+    Open the curves and scatter points of a plot item in a Matplotlib window.
+
+    Unlike pyqtgraph's MatplotlibExporter, this also handles points that each have their own
+    color (ScatterWaveform, MotorMap). Requires matplotlib to be installed.
+
+    Args:
+        plot_item (pg.PlotItem): The plot item to export.
+
+    Returns:
+        MatplotlibWindow: The opened Matplotlib window.
+    """
+    from pyqtgraph.exporters.Matplotlib import MatplotlibExporter, MatplotlibWindow
+
+    window = MatplotlibWindow()
+    MatplotlibExporter.windows.append(window)
+    try:
+        _draw_plot_item(window.getFigure(), plot_item)
+    except Exception:
+        window.close()
+        raise
+    window.draw()
+    return window
+
+
+def _mpl_marker(symbol) -> str | int:
+    """Map a pyqtgraph symbol to a matplotlib marker; other symbols become circles."""
+    return _SYMBOL_PG_TO_MPL.get(symbol, "o") if isinstance(symbol, str) else "o"
+
+
+def _draw_plot_item(fig, plot_item: pg.PlotItem):
+    """Draw the visible data items of a plot item into a new axes of a matplotlib figure."""
+    x_axis = plot_item.getAxis("bottom")
+    y_axis = plot_item.getAxis("left")
+    # Scale the data like pyqtgraph does so matplotlib doesn't add its own scale factor label
+    x_scale = x_axis.autoSIPrefixScale if x_axis.autoSIPrefix else 1.0
+    y_scale = y_axis.autoSIPrefixScale if y_axis.autoSIPrefix else 1.0
+
+    ax = fig.add_subplot(111)
+    ax.set_title(plot_item.titleLabel.text)
+    ax.spines["right"].set_color("none")
+    ax.spines["top"].set_color("none")
+    ax.xaxis.set_ticks_position("bottom")
+
+    for item in plot_item.curves:
+        if not item.isVisible():
+            continue
+        x, y = item.getData()
+        if x is None or y is None or len(x) == 0:
+            continue
+        x = np.asarray(x) * x_scale
+        y = np.asarray(y) * y_scale
+        if isinstance(item, pg.ScatterPlotItem):
+            _draw_scatter_item(ax, item, x, y)
+        else:
+            _draw_data_item(ax, item, x, y)
+
+    (x_min, x_max), (y_min, y_max) = plot_item.viewRange()
+    ax.set_xbound(x_min * x_scale, x_max * x_scale)
+    ax.set_ybound(y_min * y_scale, y_max * y_scale)
+    ax.set_xlabel(x_axis.label.toPlainText())
+    ax.set_ylabel(y_axis.label.toPlainText())
+
+
+def _draw_data_item(ax, item: pg.PlotDataItem, x: np.ndarray, y: np.ndarray):
+    """Draw a curve (line, fill and symbols) of a PlotDataItem or PlotCurveItem."""
+    opts = item.opts
+    pen = fn.mkPen(opts.get("pen"))
+    if pen.style() != Qt.PenStyle.NoPen:
+        if opts.get("fillLevel") is not None and opts.get("fillBrush") is not None:
+            ax.fill_between(
+                x=x, y1=y, y2=opts["fillLevel"], facecolor=_brush_colors(opts["fillBrush"])
+            )
+        ax.plot(x, y, color=pen.color().getRgbF(), linewidth=pen.widthF(), linestyle="-")
+
+    symbol = opts.get("symbol")
+    if symbol is None:
+        return
+    face = _brush_colors(opts.get("symbolBrush"))
+    edge = _pen_colors(opts.get("symbolPen"))
+    size = opts.get("symbolSize", 10)
+    if _is_per_point(face) or _is_per_point(edge) or np.ndim(size) > 0:
+        ax.scatter(
+            x,
+            y,
+            marker=_mpl_marker(symbol),
+            s=np.asarray(size, dtype=float) ** 2,
+            facecolors=face,
+            edgecolors=edge,
+            zorder=3,
+        )
+    else:
+        ax.plot(
+            x,
+            y,
+            linestyle="",
+            marker=_mpl_marker(symbol),
+            markersize=size,
+            markerfacecolor=face,
+            markeredgecolor=edge,
+        )
+
+
+def _draw_scatter_item(ax, item: pg.ScatterPlotItem, x: np.ndarray, y: np.ndarray):
+    """Draw a ScatterPlotItem, whose brushes, pens and sizes may be set per point."""
+    opts = item.opts
+    face = _brush_colors(list(item.data["brush"]), default=_brush_colors(opts["brush"]))
+    edge = _pen_colors(list(item.data["pen"]), default=_pen_colors(opts["pen"]))
+    sizes = np.where(item.data["size"] < 0, opts["size"], item.data["size"])
+    ax.scatter(
+        x,
+        y,
+        marker=_mpl_marker(opts["symbol"]),
+        s=np.asarray(sizes, dtype=float) ** 2,
+        facecolors=face,
+        edgecolors=edge,
+        zorder=3,
+    )
+
+
 class PlotExportConnection(BundleConnection):
     def __init__(self, components: ToolbarComponents, target_widget):
         super().__init__()
@@ -93,9 +276,8 @@ class PlotExportConnection(BundleConnection):
         if self.target_widget:
             try:
                 import matplotlib as mpl
-                from pyqtgraph.exporters import MatplotlibExporter
 
-                MatplotlibExporter(self.target_widget.plot_item).export()
+                export_to_matplotlib(self.target_widget.plot_item)
             except ModuleNotFoundError:
                 warning_util = WarningPopupUtility()
                 warning_util.show_warning(

tests/unit_tests/test_plot_base_next_gen.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 import numpy as np
+import pyqtgraph as pg
 import pytest
 from bec_lib import messages
 from bec_lib.scan_data_container import ScanDataContainer
@@ -14,6 +15,7 @@ from bec_widgets.widgets.plots.multi_waveform.multi_waveform import MultiWavefor
 from bec_widgets.widgets.plots.plot_base import PlotBase, UIMode
 from bec_widgets.widgets.plots.plot_info_label import TextOnlyLegendSample
 from bec_widgets.widgets.plots.scatter_waveform.scatter_waveform import ScatterWaveform
+from bec_widgets.widgets.plots.toolbar_components.plot_export import _draw_plot_item
 from bec_widgets.widgets.plots.waveform.waveform import Waveform
 
 from .conftest import create_widget
@@ -802,12 +804,13 @@ def test_matplotlib_export_action_opens_exporter(qtbot, mocked_client):
 
     with (
         mock.patch.dict("sys.modules", {"matplotlib": mock.MagicMock()}),
-        mock.patch("pyqtgraph.exporters.MatplotlibExporter") as exporter,
+        mock.patch(
+            "bec_widgets.widgets.plots.toolbar_components.plot_export.export_to_matplotlib"
+        ) as export,
     ):
         matplotlib_action.trigger()
 
-    exporter.assert_called_once_with(pb.plot_item)
-    exporter.return_value.export.assert_called_once()
+    export.assert_called_once_with(pb.plot_item)
 
 
 def test_matplotlib_export_action_warns_without_matplotlib(qtbot, mocked_client):
@@ -827,3 +830,25 @@ def test_matplotlib_export_action_warns_without_matplotlib(qtbot, mocked_client)
         warning_util.return_value.show_warning.call_args.kwargs["title"]
         == "Matplotlib not installed"
     )
+
+
+def test_matplotlib_export_draws_per_point_colors(qtbot, mocked_client):
+    pb = create_widget(qtbot, PlotBase, client=mocked_client)
+    pb.title = "Export"
+    x = np.arange(5, dtype=float)
+    brushes = [pg.mkBrush(i * 50, 0, 0) for i in range(5)]
+    # PlotDataItem with one brush per point, as ScatterCurve sets it
+    pb.plot_item.addItem(pg.PlotDataItem(x=x, y=x, symbol="o", symbolBrush=brushes))
+    # ScatterPlotItem with one brush per point, as the MotorMap trace sets it
+    pb.plot_item.addItem(pg.ScatterPlotItem(x=x, y=2 * x, brush=brushes, pen=None))
+    fig = mock.MagicMock()
+    ax = fig.add_subplot.return_value
+
+    _draw_plot_item(fig, pb.plot_item)
+
+    ax.set_title.assert_called_once_with("Export")
+    assert ax.scatter.call_count == 2
+    for call in ax.scatter.call_args_list:
+        facecolors = call.kwargs["facecolors"]
+        assert len(facecolors) == 5
+        assert facecolors[1] == brushes[1].color().getRgbF()