class ViolinItem(QGraphicsWidget):
HEIGHT = 50
POINT_R = 6
SCALE_FACTOR = 0.5
selection_changed = Signal(float, float, str)
class SelectionRect(QGraphicsRectItem):
COLOR = [255, 255, 0]
def __init__(self, parent, width: int):
super().__init__(parent)
self.parent_width = width
color = QColor(*self.COLOR)
color.setAlpha(100)
self.setBrush(color)
color = QColor(*self.COLOR)
self.setPen(color)
def __init__(self, parent, attr_name: str, x_range: Tuple[float],
width: int):
super().__init__(parent)
assert x_range[0] == -x_range[1]
self.__attr_name = attr_name
self.__width = width
self.__range = x_range[1] if x_range[1] else 1
self.__group = None # type: Optional[QGraphicsItemGroup]
self.__selection_rect = None # type: Optional[QGraphicsRectItem]
self.x_data = None # type: Optional[np.ndarray]
parent.selection_cleared.connect(self.__remove_selection_rect)
@property
def attr_name(self):
return self.__attr_name
def set_data(self, x_data: np.ndarray, color_data: np.ndarray):
def put_point(_x, _y):
item = QGraphicsEllipseItem()
item.setX(_x)
item.setY(_y)
item.setRect(0, 0, self.POINT_R, self.POINT_R)
color = QColor(*colors.pop().astype(int))
item.setPen(QPen(color))
item.setBrush(QBrush(color))
self.__group.addToGroup(item)
self.x_data = x_data
self.__group = QGraphicsItemGroup(self)
x_data_unique, dist, x_data = self.prepare_data()
for x, d in zip(x_data_unique, dist):
colors = color_data[x_data == np.round(x, 3)]
colors = list(colors[np.random.choice(len(colors), 11)])
y = self.HEIGHT / 2 - self.POINT_R / 2
self.plot_data(put_point, x, y, d)
def prepare_data(self):
x_data = self._values_to_pixels(self.x_data)
x_data = x_data[~np.isnan(x_data)]
if len(x_data) == 0:
return
x_data = np.round(x_data - self.POINT_R / 2, 3)
# remove duplicates and get counts (distribution) to set y
x_data_unique, counts = np.unique(x_data, return_counts=True)
min_count, max_count = np.min(counts), np.max(counts)
dist = (counts - min_count) / (max_count - min_count)
if min_count == max_count:
dist[:] = 1
dist = dist ** 0.7
# plot rarest values first
indices = np.argsort(counts)
return x_data_unique[indices], dist[indices], x_data
@staticmethod
def plot_data(func: Callable, x: float, y: float, d: float):
func(x, y) # y = 0
if d > 0:
offset = d * 10
func(x, y + offset) # y = (0, 10]
func(x, y - offset) # y = (-10, 0]
for i in range(2, int(offset), 2):
func(x, y + i) # y = [2, 8]
func(x, y - i) # y = [-8, -2]
def _values_to_pixels(self, x: np.ndarray) -> np.ndarray:
# scale data to [-0.5, 0.5]
x = x / self.__range * self.SCALE_FACTOR
# round data to 3. decimal for sampling
x = np.round(x, 3)
# convert to pixels
return x * self.__width + self.__width / 2
def _values_from_pixels(self, p: np.ndarray) -> np.ndarray:
# convert from pixels
x = (p - self.__width / 2) / self.__width
# rescale data from [-0.5, 0.5]
return np.round(x * self.__range / self.SCALE_FACTOR, 3)
def rescale(self, width):
def move_point(_x, *_):
item = next(points)
item.setX(_x)
self.__width = width
self.updateGeometry()
points = (item for item in self.__group.childItems())
x_data_unique, dist, x_data = self.prepare_data()
for x, d in zip(x_data_unique, dist):
self.plot_data(move_point, x, 0, d)
if self.__selection_rect is not None:
old_width = self.__selection_rect.parent_width
rect = self.__selection_rect.rect()
x1 = self.__width * rect.x() / old_width
x2 = self.__width * (rect.x() + rect.width()) / old_width
rect = QRectF(x1, rect.y(), x2 - x1, rect.height())
self.__selection_rect.setRect(rect)
self.__selection_rect.parent_width = self.__width
def sizeHint(self, *_):
return QSizeF(self.__width, self.HEIGHT)
def __remove_selection_rect(self):
if self.__selection_rect is not None:
self.__selection_rect.setParentItem(None)
if self.scene() is not None:
self.scene().removeItem(self.__selection_rect)
self.__selection_rect = None
def add_selection_rect(self, x1, x2):
x1, x2 = self._values_to_pixels(np.array([x1, x2]))
rect = QRectF(x1, 0, x2 - x1, self.HEIGHT)
self.__selection_rect = ViolinItem.SelectionRect(self, self.__width)
self.__selection_rect.setRect(rect)
def mousePressEvent(self, event: QGraphicsSceneMouseEvent):
event.accept()
def mouseMoveEvent(self, event: QGraphicsSceneMouseEvent):
if event.buttons() & Qt.LeftButton:
if self.__selection_rect is None:
self.__selection_rect = ViolinItem.SelectionRect(
self, self.__width)
x = event.buttonDownPos(Qt.LeftButton).x()
rect = QRectF(x, 0, event.pos().x() - x, self.HEIGHT).normalized()
rect = rect.intersected(self.contentsRect())
self.__selection_rect.setRect(rect)
event.accept()
def mouseReleaseEvent(self, event: QGraphicsSceneMouseEvent):
x1 = event.buttonDownPos(Qt.LeftButton).x()
x2 = event.pos().x()
if x1 > x2:
x2, x1 = x1, x2
x1, x2 = self._values_from_pixels(np.array([x1, x2]))
self.selection_changed.emit(x1, x2, self.__attr_name)
event.accept()
