Source code for whitecanvas.backend.vispy.bars

from __future__ import annotations
import numpy as np
from numpy.typing import NDArray

from vispy.scene import visuals
from whitecanvas.protocols import BarProtocol, check_protocol
from whitecanvas.types import LineStyle, FacePattern
from whitecanvas.utils.normalize import as_color_array
from whitecanvas.backend import _not_implemented


[docs]@check_protocol(BarProtocol) class Bars(visuals.Compound): def __init__(self, xlow, xhigh, ylow, yhigh): center = np.stack([(xlow + xhigh) / 2, (ylow + yhigh) / 2], axis=1) rectangles: list[visuals.Rectangle] = [] for c, w, h in zip(center, xhigh - xlow, yhigh - ylow): w = abs(w) h = abs(h) if w == 0: w = h * 1e-6 elif h == 0: h = w * 1e-6 rectangles.append( visuals.Rectangle(center=c, width=w, height=h, border_width=0) ) super().__init__(rectangles) self.unfreeze() self._rectangles = rectangles ##### LayerProtocol ##### def _plt_get_visible(self) -> bool: return self.visible def _plt_set_visible(self, visible: bool): self.visible = visible ##### XYDataProtocol ##### def _plt_get_data(self): data = [] for rect in self._rectangles: xc, yc = rect.center w = rect.width h = rect.height data.append([xc - w / 2, xc + w / 2, yc - h / 2, yc + h / 2]) return tuple(np.array(data).T) def _plt_set_data(self, xlow, xhigh, ylow, yhigh): center = np.stack([(xlow + xhigh) / 2, (ylow + yhigh) / 2], axis=1) for rect, c, w, h in zip(self._rectangles, center, xhigh - xlow, yhigh - ylow): rect.center = c rect.width = w rect.height = h def _plt_get_ndata(self): return len(self._rectangles) ##### HasFace protocol ##### def _plt_get_face_color(self) -> NDArray[np.float32]: colors = [] for rect in self._rectangles: colors.append(rect.color) return np.stack(colors, axis=0) def _plt_set_face_color(self, color: NDArray[np.float32]): color = as_color_array(color, size=self._plt_get_ndata()) for rect, c in zip(self._rectangles, color): rect.color = c _plt_get_face_pattern, _plt_set_face_pattern = _not_implemented.face_patterns() ##### HasEdges protocol ##### def _plt_get_edge_color(self) -> NDArray[np.float32]: colors = [] for rect in self._rectangles: colors.append(rect.border_color) return np.stack(colors, axis=0) def _plt_set_edge_color(self, color: NDArray[np.float32]): color = as_color_array(color, size=self._plt_get_ndata()) for rect, c in zip(self._rectangles, color): rect.border_color = c def _plt_get_edge_width(self) -> NDArray[np.floating]: # NOTE: vispy forgot to add `border_width` to `Rectangle` class? widths = [] for rect in self._rectangles: widths.append(rect._border_width) return np.array(widths) def _plt_set_edge_width(self, width: float): if np.isscalar(width): width = np.full(self._plt_get_ndata(), width) for rect, w in zip(self._rectangles, width): rect._border_width = w _plt_get_edge_style, _plt_set_edge_style = _not_implemented.edge_styles()