Source code for magicclass.ext.pyqtgraph.graph_items

from __future__ import annotations
from typing import Sequence
import pyqtgraph as pg
from qtpy.QtCore import Qt
import numpy as np

from .._shared_utils import convert_color_code, to_rgba

# compatibility with matplotlib
_LINE_STYLE = {
    "-": Qt.SolidLine,
    "--": Qt.DashLine,
    ":": Qt.DotLine,
    "-.": Qt.DashDotLine,
}

_SYMBOL_MAP = {
    "*": "star",
    "D": "d",
    "^": "t1",
    "<": "t3",
    "v": "t",
    ">": "t2",
}


[docs]class LayerItem: native: pg.GraphicsItem def __repr__(self) -> str: return f"{self.__class__.__name__} '{self.name}'" @property def visible(self): """Visibility of data.""" return self.native.isVisible() @visible.setter def visible(self, value: bool): self.native.setVisible(value) @property def zorder(self) -> float: """Z-order of item. Item with larger z will be displayed on the top.""" return self.native.zValue() @zorder.setter def zorder(self, value: float): self.native.setZValue(value)
[docs]class PlotDataLayer(LayerItem): native: pg.PlotCurveItem | pg.ScatterPlotItem @property def xdata(self) -> np.ndarray: return self.native.getData()[0] @xdata.setter def xdata(self, value: Sequence[float]): self.native.setData(value, self.ydata) @property def ydata(self) -> np.ndarray: return self.native.getData()[1] @ydata.setter def ydata(self, value: Sequence[float]): self.native.setData(self.xdata, value) @property def ndata(self) -> int: return self.native.getData()[0].size @property def name(self) -> str: return self.native.opts["name"] @name.setter def name(self, value: str): value = str(value) self.native.opts["name"] = value # TODO: now name is not linked to label item
[docs] def add(self, points: np.ndarray | Sequence, **kwargs): """Add new points to the plot data item.""" points = np.atleast_2d(points) if points.shape[1] != 2: raise ValueError("Points must be of the shape (N, 2).") self.native.setData( np.concatenate([self.xdata, points[:, 0]]), np.concatenate([self.ydata, points[:, 1]]), **kwargs, ) return None
[docs] def remove(self, i: int | Sequence[int]): """Remove the i-th data.""" if isinstance(i, int): i = [i] sl = list(set(range(self.ndata)) - set(i)) xdata = self.xdata[sl] ydata = self.ydata[sl] self.native.setData(xdata, ydata) return None
@property def edge_color(self) -> np.ndarray: """Edge color of the data.""" return to_rgba(self.native.opts["pen"]) @edge_color.setter def edge_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setPen(value, width=self.lw, style=self.ls) @property def face_color(self) -> np.ndarray: """Face color of the data.""" return to_rgba(self.native.opts["brush"]) @face_color.setter def face_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setBrush(value) color = property() @color.setter def color(self, value: str | Sequence): """Set face color and edge color at the same time.""" self.face_color = value self.edge_color = value @property def lw(self): """Line width.""" return self.native.opts["pen"].width() @lw.setter def lw(self, value: float): self.native.opts["pen"].setWidth(value) linewidth = lw # alias @property def ls(self): """Line style.""" return self.native.opts["pen"].style() @ls.setter def ls(self, value: str): _ls = _LINE_STYLE[value] self.native.opts["pen"].setStyle(_ls) linestyle = ls # alias
[docs]class Scatter(PlotDataLayer): native: pg.ScatterPlotItem def __init__( self, x, y, face_color=None, edge_color=None, size: float = 7, name: str | None = None, lw: float = 1, ls: str = "-", symbol="o", ): face_color, edge_color = _set_default_colors( face_color, edge_color, "white", "white" ) pen = pg.mkPen(edge_color, width=lw, style=_LINE_STYLE[ls]) brush = pg.mkBrush(face_color) symbol = _SYMBOL_MAP.get(symbol, symbol) self.native = pg.ScatterPlotItem( x=x, y=y, pen=pen, brush=brush, size=size, symbol=symbol ) self.name = name @property def symbol(self): return self.native.opts["symbol"] @symbol.setter def symbol(self, value): value = _SYMBOL_MAP.get(value, value) self.native.setSymbol(value) @property def size(self): return self.native.opts["symbolSize"] @size.setter def size(self, size: float): self.native.setSymbolSize(size)
[docs]class Curve(PlotDataLayer): native: pg.PlotDataItem def __init__( self, x, y, face_color=None, edge_color=None, size: float = 7, name: str | None = None, lw: float = 1, ls: str = "-", symbol=None, ): face_color, edge_color = _set_default_colors( face_color, edge_color, "white", "white" ) pen = pg.mkPen(edge_color, width=lw, style=_LINE_STYLE[ls]) brush = pg.mkBrush(face_color) symbol = _SYMBOL_MAP.get(symbol, symbol) self.native = pg.PlotDataItem( x=x, y=y, pen=pen, brush=brush, symbolSize=size, symbol=symbol, symbolPen=pen, symbolBrush=brush, ) self.name = name @property def symbol(self): return self.native.opts["symbol"] @symbol.setter def symbol(self, value): value = _SYMBOL_MAP.get(value, value) self.native.setSymbol(value) @property def size(self): return self.native.opts["symbolSize"] @size.setter def size(self, size: float): self.native.setSymbolSize(size) @property def edge_color(self) -> np.ndarray: rgba = self.native.opts["pen"].color().getRgb() return np.array(rgba) / 255 @edge_color.setter def edge_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setPen(value) self.native.setSymbolPen(value) @property def face_color(self) -> np.ndarray: rgba = self.native.opts["symbolBrush"].color().getRgb() return np.array(rgba) / 255 @face_color.setter def face_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setBrush(value) self.native.setSymbolBrush(value)
[docs]class Histogram(PlotDataLayer): native: pg.ScatterPlotItem def __init__( self, data, bins: int | Sequence | str = 10, range=None, density: bool = False, face_color=None, edge_color=None, name: str | None = None, lw: float = 1, ls: str = "-", ): face_color, edge_color = _set_default_colors( face_color, edge_color, "white", "white" ) pen = pg.mkPen(edge_color, width=lw, style=_LINE_STYLE[ls]) brush = pg.mkBrush(face_color) y, x = np.histogram(data, bins=bins, range=range, density=density) self.native = pg.PlotCurveItem( x=x, y=y, pen=pen, brush=brush, stepMode="center", fillLevel=0 ) self.name = name
[docs] def set_hist(self, data, bins=10, range=None, density=False): y, x = np.histogram(data, bins=bins, range=range, density=density) self.native.setData(x=x, y=y)
[docs]class BarPlot(PlotDataLayer): native: pg.BarGraphItem def __init__( self, x, y, face_color=None, edge_color=None, width: float = 0.6, name: str | None = None, lw: float = 1, ls: str = "-", ): face_color, edge_color = _set_default_colors( face_color, edge_color, "white", "white" ) pen = pg.mkPen(edge_color, width=lw, style=_LINE_STYLE[ls]) brush = pg.mkBrush(face_color) self.native = pg.BarGraphItem(x=x, height=y, width=width, pen=pen, brush=brush) self.name = name @property def edge_color(self) -> np.ndarray: rgba = self.native.opts["pen"].color().getRgb() return np.array(rgba) / 255 @edge_color.setter def edge_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setOpts(pen=pg.mkPen(value)) @property def face_color(self) -> np.ndarray: rgba = self.native.opts["brush"].color().getRgb() return np.array(rgba) / 255 @face_color.setter def face_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setOpts(brush=pg.mkBrush(value)) @property def xdata(self) -> np.ndarray: return self.native.getData()[0] @xdata.setter def xdata(self, value: Sequence[float]): self.native.setOpts(x=value) @property def ydata(self) -> np.ndarray: return self.native.getData()[1] @ydata.setter def ydata(self, value: Sequence[float]): self.native.setOpts(height=value)
[docs]class InfLine(LayerItem): native: pg.InfiniteLine def __init__( self, pos, angle, edge_color=None, name: str | None = None, lw: float = 1, ls: str = "-", ): if edge_color is None: edge_color = "yellow" edge_color = convert_color_code(edge_color) pen = pg.mkPen(edge_color, width=lw, style=_LINE_STYLE[ls]) self.native = pg.InfiniteLine(pos, angle, pen=pen, name=name) self.name = name @property def slope(self) -> float: """Slope of the line.""" return np.tan(np.deg2rad(self.native.angle)) @slope.setter def slope(self, value: float): self.native.setAngle(np.rad2deg(np.arctan(value))) @property def intercept(self) -> float: """Y-intercept of the line.""" a = self.slope x0, y0 = self.native.getPos() return y0 - a * x0 @intercept.setter def intercept(self, value: float): value = float(value) self.native.setPos((0, value)) @property def pos(self) -> np.ndarray: return np.array(self.native.getPos()) @pos.setter def pos(self, value): self.native.setPos(value) @property def angle(self) -> float: """Angle of the line in degree.""" return self.native.angle @angle.setter def angle(self, value: float): self.native.setAngle(value) @property def edge_color(self) -> np.ndarray: return to_rgba(self.native.pen) @edge_color.setter def edge_color(self, value): value = convert_color_code(value) self.native.setPen(value, width=self.lw, style=self.ls) color = edge_color @property def name(self): return self.native._name @name.setter def name(self, value: str): self.native.setName(value) @property def lw(self): """Line width.""" return self.native.pen.width() @lw.setter def lw(self, value: float): self.native.pen.setWidth(value) linewidth = lw # alias @property def ls(self): """Line style.""" return self.native.pen.style() @ls.setter def ls(self, value: str): _ls = _LINE_STYLE[value] self.native.pen.setStyle(_ls) linestyle = ls # alias
[docs]class FillBetween(PlotDataLayer): native: pg.FillBetweenItem def __init__( self, x, y1, y2, face_color=None, edge_color=None, name: str | None = None, lw: float = 1, ls: str = "-", ): face_color, edge_color = _set_default_colors( face_color, edge_color, "white", "white" ) pen = pg.mkPen(edge_color, width=lw, style=_LINE_STYLE[ls]) brush = pg.mkBrush(face_color) curve1 = pg.PlotCurveItem(x=x, y=y1, pen=pen) curve2 = pg.PlotCurveItem(x=x, y=y2, pen=pen) self.native = pg.FillBetweenItem(curve1, curve2, brush=brush, pen=pen) self.name = name @property def edge_color(self) -> np.ndarray: rgba = self.native.curves[0].opts["pen"].color().getRgb() return np.array(rgba) / 255 @edge_color.setter def edge_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setPen(pg.mkPen(value)) @property def face_color(self) -> np.ndarray: rgba = self.native.curves[0].opts["brush"].color().getRgb() return np.array(rgba) / 255 @face_color.setter def face_color(self, value: str | Sequence): value = convert_color_code(value) self.native.setBrush(pg.mkBrush(value)) @property def name(self) -> str: return self.native.curves[0].opts["name"] @name.setter def name(self, value: str): value = str(value) self.native.curves[0].opts["name"] = value @property def lw(self): """Line width.""" return self.native.curves[0].opts["pen"].width() @lw.setter def lw(self, value: float): self.native.curves[0].opts["pen"].setWidth(value) self.native.curves[1].opts["pen"].setWidth(value) linewidth = lw # alias @property def ls(self): """Line style.""" return self.native.curves[0].opts["pen"].style() @ls.setter def ls(self, value: str): _ls = _LINE_STYLE[value] self.native.curves[0].opts["pen"].setStyle(_ls) self.native.curves[1].opts["pen"].setStyle(_ls) linestyle = ls # alias
# WIP! # How to update a subset of properties? item.text[2:5] = "new" or item[2:5].text = "new"
[docs]class TextGroup(LayerItem): def __init__( self, x: Sequence[float], y: Sequence[float], texts: Sequence[str], color=None, name: str = None, ): self.native = pg.ItemGroup() if color is None: color = "white" for x_, y_, text_ in zip(x, y, texts): item = pg.TextItem(text_, color=convert_color_code(color)) item.setPos(x_, y_) self.native.addItem(item) self.name = name @property def text_items(self) -> list[pg.TextItem]: return self.native.childItems() def __getitem__(self, key: int | slice) -> TextItemView: return TextItemView(self.text_items[key]) @property def xdata(self) -> np.ndarray: return np.array([item.pos().x() for item in self.text_items]) @property def ydata(self) -> np.ndarray: return np.array([item.pos().y() for item in self.text_items]) @property def color(self) -> np.ndarray: """Text color.""" rgba = np.stack([item.color.getRgb() for item in self.text_items]) return rgba / 255 @color.setter def color(self, value): value = convert_color_code(value) for item in self.text_items: item.setText(item.toPlainText(), value) @property def background_color(self) -> np.ndarray: """Text background color.""" return np.stack([to_rgba(item.fill) for item in self.text_items]) @background_color.setter def background_color(self, value): value = convert_color_code(value) brush = pg.mkBrush(value) for item in self.text_items: item.fill = brush item._updateView() @property def border(self) -> np.ndarray: """Border color of text bounding box.""" if isinstance(self.native, list): return np.stack([to_rgba(item.border) for item in self.native]) else: return to_rgba(self.native.border) @border.setter def border(self, value): value = convert_color_code(value) pen = pg.mkPen(value) for item in self.text_items: item.border = pen item._updateView() @property def text(self) -> str | list[str]: """Text string.""" return [item.toPlainText() for item in self.text_items] @text.setter def text(self, value: str): for item in self.text_items: item.setText(value) @property def anchor(self) -> np.ndarray: """Text anchor position.""" out = [] for item in self.text_items: anchor = item.anchor out.append([anchor.x(), anchor.y()]) return np.array(out) @anchor.setter def anchor(self, value): for item in self.text_items: item.setAnchor(value)
[docs]class TextItemView: def __init__(self, textitem: pg.TextItem | list[pg.TextItem]): self.native = textitem @property def color(self) -> np.ndarray: """Text color.""" if isinstance(self.native, list): rgba = np.stack([item.color.getRgb() for item in self.native]) arr = rgba / 255 else: rgba = self.native.color.getRgb() arr = np.array(rgba) / 255 return arr @color.setter def color(self, value): value = convert_color_code(value) if isinstance(self.native, list): for item in self.native: item.setText(item.text, value) else: self.native.setText(self.text, value) @property def background_color(self) -> np.ndarray: """Text background color.""" if isinstance(self.native, list): return np.stack([to_rgba(item.fill) for item in self.native]) else: return to_rgba(self.native.fill) @background_color.setter def background_color(self, value): value = convert_color_code(value) brush = pg.mkBrush(value) if isinstance(self.native, list): for item in self.native: item.fill = brush item._updateView() else: self.native.fill = brush self.native._updateView() @property def border(self) -> np.ndarray: """Border color of text bounding box.""" if isinstance(self.native, list): return np.stack([to_rgba(item.border) for item in self.native]) else: return to_rgba(self.native.border) @border.setter def border(self, value): value = convert_color_code(value) pen = pg.mkPen(value) if isinstance(self.native, list): for item in self.native: item.border = pen item._updateView() else: self.native.border = pen self.native._updateView() @property def text(self) -> str | list[str]: """Text string.""" if isinstance(self.native, list): return [item.toPlainText() for item in self.native] else: return self.native.toPlainText() @text.setter def text(self, value: str): if isinstance(self.native, list): for item in self.native: item.setText(value) else: self.native.setText(value) @property def anchor(self) -> np.ndarray: """Text anchor position.""" if isinstance(self.native, list): out = [] for item in self.native: anchor = item.anchor out.append([anchor.x(), anchor.y()]) return np.array(out) else: anchor = self.native.anchor return np.array([anchor.x(), anchor.y()]) @anchor.setter def anchor(self, value): if isinstance(self.native, list): for item in self.native: item.setAnchor(value) else: self.native.setAnchor(value)
def _set_default_colors(face_color, edge_color, default_f, default_e): if face_color is None: face_color = default_f else: face_color = convert_color_code(face_color) if edge_color is None: edge_color = default_e else: edge_color = convert_color_code(edge_color) return face_color, edge_color