From a75b07c5fab0861de087820cc04f4714aed3995f Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Fri, 16 Oct 2026 14:33:02 +0000 Subject: [PATCH] fix(plots): export per-point colors and the plot title to matplotlib --- .../plots/toolbar_components/plot_export.py | 186 +++++++++++++++++- tests/unit_tests/test_plot_base_next_gen.py | 31 ++- 2 files changed, 212 insertions(+), 5 deletions(-) diff --git a/bec_widgets/widgets/plots/toolbar_components/plot_export.py b/bec_widgets/widgets/plots/toolbar_components/plot_export.py index 5e1ecec1..cc44d01a 100644 --- a/bec_widgets/widgets/plots/toolbar_components/plot_export.py +++ b/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( diff --git a/tests/unit_tests/test_plot_base_next_gen.py b/tests/unit_tests/test_plot_base_next_gen.py index 6b8a57f1..ce28d871 100644 --- a/tests/unit_tests/test_plot_base_next_gen.py +++ b/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()