import numpy as np from pyqtgraph import ( InfiniteLine, ErrorBarItem, LinearRegionItem, mkBrush, mkColor, mkPen, PlotDataItem, LegendItem, ) from nmreval.lib.colors import BaseColor, Colors from nmreval.lib.lines import LineStyle from nmreval.lib.symbols import SymbolStyle from ..Qt import QtCore, QtGui """ Subclasses of pyqtgraph items, mostly to take care of log-scaling. pyqtgraph looks for function "setLogMode" for logarithmic axes, so needs to be implemented. """ class LogInfiniteLine(InfiniteLine): def __init__(self, **kwargs): super().__init__(**kwargs) self.logmode = [False, False] def setLogMode(self, xmode, ymode): """ Does only work for vertical and horizontal lines """ if self.logmode == [xmode, ymode]: return new_p = list(self.p[:]) if (self.angle == 90) and (self.logmode[0] != xmode): if xmode: new_p[0] = np.log10(new_p[0]+np.finfo(float).eps) else: new_p[0] = 10**new_p[0] if (self.angle == 0) and (self.logmode[1] != ymode): if ymode: new_p[1] = np.log10(new_p[1]+np.finfo(float).eps) else: new_p[1] = 10**new_p[1] self.logmode = [xmode, ymode] if np.all(np.isfinite(new_p)): self.setPos(new_p) else: self.setPos(self.p) self.sigPositionChanged.emit(self) def setValue(self, v): if isinstance(v, QtCore.QPointF): v = [v.x(), v.y()] with np.errstate(divide='ignore'): if isinstance(v, (list, tuple)): for i in [0, 1]: if self.logmode[i]: v[i] = np.log10(v[i]+np.finfo(float).eps) else: if self.angle == 90: if self.logmode[0]: v = [np.log10(v+np.finfo(float).eps), 0] else: v = [v, 0] elif self.angle == 0: if self.logmode[1]: v = [0, np.log10(v+np.finfo(float).eps)] else: v = [0, v] else: raise ValueError('LogInfiniteLine: Diagonal lines need two values') self.setPos(v) def value(self): p = self.getPos() if self.angle == 0: return 10**p[1] if self.logmode[1] else p[1] elif self.angle == 90: return 10**p[0] if self.logmode[0] else p[0] else: if self.logmode[0]: p[0] = 10**p[0] if self.logmode[1]: p[1] = 10**p[1] return p class ErrorBars(ErrorBarItem): def __init__(self, **opts): self.log = [False, False] opts['xData'] = opts.get('x', None) opts['yData'] = opts.get('y', None) opts['topData'] = opts.get('top', None) opts['bottomData'] = opts.get('bottom', None) super().__init__(**opts) def setLogMode(self, x_mode, y_mode): if self.log == [x_mode, y_mode]: return self._make_log_scale(x_mode, y_mode) self.log[0] = x_mode self.log[1] = y_mode super().setData() def setData(self, **opts): self.opts.update(opts) self.opts['xData'] = opts.get('x', self.opts['xData']) self.opts['yData'] = opts.get('y', self.opts['yData']) self.opts['topData'] = opts.get('top', self.opts['topData']) self.opts['bottomData'] = opts.get('bottom', self.opts['bottomData']) if any(self.log): self._make_log_scale(*self.log) super().setData() def _make_log_scale(self, x_mode, y_mode): _x = self.opts['xData'] _xmask = np.logical_not(np.isnan(_x)) if x_mode: with np.errstate(all='ignore'): _x = np.log10(_x) _xmask = np.logical_not(np.isnan(_x)) _y = self.opts['yData'] _ymask = np.ones(_y.size, dtype=bool) _top = self.opts['topData'] _bottom = self.opts['bottomData'] if y_mode: with np.errstate(all='ignore'): logtop = np.log10(self.opts['topData']+_y) logbottom = np.log10(_y-self.opts['bottomData']) _y = np.log10(_y) _ymask = np.logical_not(np.isnan(_y)) logbottom[logbottom == -np.inf] = _y[logbottom == -np.inf] _bottom = np.nan_to_num(np.maximum(_y-logbottom, 0)) logtop[logtop == -np.inf] = _y[logtop == -np.inf] _top = np.nan_to_num(np.maximum(logtop-_y, 0)) _mask = np.logical_and(_xmask, _ymask) self.opts['x'] = _x[_mask] self.opts['y'] = _y[_mask] self.opts['top'] = _top[_mask] self.opts['bottom'] = _bottom[_mask] class PlotItem(PlotDataItem): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.opts['linecolor'] = (0, 0, 0) self.opts['symbolcolor'] = (0, 0, 0) if self.opts['pen'] is not None: pen = self.opts['pen'] if isinstance(pen, tuple): self.opts['linecolor'] = pen else: c = pen.color() self.opts['linecolor'] = c.red(), c.green(), c.blue() if self.symbol != SymbolStyle.No: c = self.opts['symbolBrush'].color() self.opts['symbolcolor'] = c.red(), c.green(), c.blue() def __getitem__(self, item): return self.opts.get(item, None) @property def symbol(self): return SymbolStyle.from_str(self.opts['symbol']) @property def symbolcolor(self): sc = self.opts['symbolcolor'] if isinstance(sc, tuple): return Colors(sc) elif isinstance(sc, str): return Colors.from_str(sc) else: return sc @property def symbolsize(self): return self.opts['symbolSize'] @property def linestyle(self) -> LineStyle: pen = self.opts['pen'] if pen is None: return LineStyle.No else: return LineStyle(pen.style()) @property def linewidth(self) -> float: pen = self.opts['pen'] if pen is None: return 1. else: return pen.widthF() @property def linecolor(self) -> Colors: lc = self.opts['linecolor'] if isinstance(lc, tuple): return Colors(lc) elif isinstance(lc, str): return Colors.from_str(lc) else: return lc def updateItems(self): """ We override this function so that curves with nan/inf values can be displayed. Newer versions close this bug differently (https://github.com/pyqtgraph/pyqtgraph/pull/1058) but this works somewhat. """ curveArgs = {} for k, v in [('pen', 'pen'), ('shadowPen', 'shadowPen'), ('fillLevel', 'fillLevel'), ('fillOutline', 'fillOutline'), ('fillBrush', 'brush'), ('antialias', 'antialias'), ('connect', 'connect'), ('stepMode', 'stepMode')]: curveArgs[v] = self.opts[k] scatterArgs = {} for k, v in [('symbolPen', 'pen'), ('symbolBrush', 'brush'), ('symbol', 'symbol'), ('symbolSize', 'size'), ('data', 'data'), ('pxMode', 'pxMode'), ('antialias', 'antialias')]: if k in self.opts: scatterArgs[v] = self.opts[k] x, y = self.getData() if x is None: x = [] if y is None: y = [] if curveArgs['pen'] is not None or (curveArgs['brush'] is not None and curveArgs['fillLevel'] is not None): is_finite = np.isfinite(x) & np.isfinite(y) all_finite = np.all(is_finite) if not all_finite: # remove all bad values x = x[is_finite] y = y[is_finite] self.curve.setData(x=x, y=y, **curveArgs) self.curve.show() else: self.curve.hide() if scatterArgs['symbol'] is not None: if self.opts.get('stepMode', False) is True: x = 0.5 * (x[:-1] + x[1:]) self.scatter.setData(x=x, y=y, **scatterArgs) self.scatter.show() else: self.scatter.hide() def set_symbol(self, symbol=None, size=None, color=None): if symbol is not None: if isinstance(symbol, int): self.setSymbol(SymbolStyle(symbol).to_str()) elif isinstance(symbol, SymbolStyle): self.setSymbol(symbol.to_str()) else: self.setSymbol(symbol) if color is not None: self.set_color(color, symbol=True) if size is not None: self.setSymbolSize(size) def set_color(self, color, symbol=False, line=False): if isinstance(color, BaseColor): color = color.rgb() elif isinstance(color, QtGui.QColor): color = color.getRgb()[:3] if symbol: self.setSymbolBrush(mkBrush(color)) self.setSymbolPen(mkPen(color=color)) self.opts['symbolcolor'] = color if line: pen = self.opts['pen'] self.opts['linecolor'] = color if pen is not None: pen.setColor(mkColor(color)) self.opts['pen'] = pen self.updateItems() def set_line(self, style=None, width=None, color=None): pen = self.opts['pen'] if pen is None: pen = mkPen(style=QtCore.Qt.NoPen) if width is not None: pen.setWidthF(width) if style is not None: if isinstance(style, LineStyle): style = style.value pen.setStyle(style) self.opts['pen'] = pen self.updateItems() if color is not None: self.set_color(color, symbol=False, line=True) def get_data_opts(self) -> dict: x, y = self.xData, self.yData if (x is None) or (len(x) == 0): return {} opts = self.opts item_dic = { 'x': x, 'y': y, 'name': opts.get('name', ''), 'symbolsize': opts['symbolSize'], } if opts['symbol'] is None: item_dic['symbol'] = SymbolStyle.No item_dic['symbolcolor'] = None else: item_dic['symbol'] = SymbolStyle.from_str(opts['symbol']) item_dic['symbolcolor'] = opts['symbolcolor'] pen = opts['pen'] if pen is not None: item_dic['linestyle'] = LineStyle(pen.style()) item_dic['linecolor'] = opts['linecolor'] item_dic['linewidth'] = pen.widthF() else: item_dic['linestyle'] = LineStyle.No item_dic['linecolor'] = None item_dic['linewidth'] = 0.0 if item_dic['linecolor'] is None and item_dic['symbolcolor'] is None: item_dic['symbolcolor'] = Colors.Black.rgb() elif item_dic['linecolor'] is None: item_dic['linecolor'] = item_dic['symbolcolor'] elif item_dic['symbolcolor'] is None: item_dic['symbolcolor'] = item_dic['linecolor'] return item_dic class RegionItem(LinearRegionItem): def __init__(self, *args, **kwargs): self.mode = kwargs.pop('mode', 'half') super().__init__(*args, **kwargs) self.logmode = False self.first = True def setLogMode(self, xmode, _): if self.logmode == xmode: return if xmode: new_region = [np.log10(self.lines[0].value()), np.log10(self.lines[1].value())] if np.isnan(new_region[1]): new_region[1] = self.lines[1].value() if np.isnan(new_region[0]): new_region[0] = new_region[1]/10. else: new_region = [10**self.lines[0].value(), 10**self.lines[1].value()] self.logmode = xmode self.setRegion(new_region) def dataBounds(self, axis, frac=1.0, orthoRange=None): if axis == self._orientation_axis[self.orientation]: r = self.getRegion() if self.logmode: r = np.log10(r[0]), np.log10(r[1]) return r else: return None def getRegion(self): region = super().getRegion() if self.logmode: return 10**region[0], 10**region[1] else: return region def boundingRect(self): # overwrite to draw correct rect in logmode br = self.viewRect() # bounds of containing ViewBox mapped to local coords. rng = self.getRegion() if self.logmode: rng = np.log10(rng[0]), np.log10(rng[1]) if self.orientation in ('vertical', LinearRegionItem.Vertical): br.setLeft(rng[0]) br.setRight(rng[1]) length = br.height() br.setBottom(br.top() + length * self.span[1]) br.setTop(br.top() + length * self.span[0]) else: br.setTop(rng[0]) br.setBottom(rng[1]) length = br.width() br.setRight(br.left() + length * self.span[1]) br.setLeft(br.left() + length * self.span[0]) br = br.normalized() if self._bounds != br: self._bounds = br self.prepareGeometryChange() return br class LegendItemBlock(LegendItem): """ Simple subclass that stops dragging legend outside of view """ def __init__(self, **kwargs): super().__init__(**kwargs) self.layout.setContentsMargins(1, 1, 1, 1) def mouseDragEvent(self, ev): if ev.button() == QtCore.Qt.LeftButton: ev.accept() dpos = ev.pos() - ev.lastPos() vb_rect = self.parentItem().rect() pos = self.pos() # upper left corner and a point a little more to the bottom right must be inside if vb_rect.contains(pos+dpos) and vb_rect.contains(pos+dpos+QtCore.QPointF(20., 20.)): self.autoAnchor(pos + dpos) else: self.autoAnchor(pos)