Violin Plot: Support scaling

Author: VesnaT

Orange/widgets/visualize/owviolinplot.py

@@ -1,3 +1,5 @@
+# pylint: disable=too-many-lines
+from collections import namedtuple
 from itertools import chain, count
 from typing import List, Optional, Tuple, Set, Sequence
 
@@ -8,7 +10,7 @@
 from AnyQt.QtCore import QItemSelection, QPointF, QRectF, QSize, Qt, Signal
 from AnyQt.QtGui import QBrush, QColor, QPainter, QPainterPath, QPolygonF
 from AnyQt.QtWidgets import QCheckBox, QSizePolicy, QGraphicsRectItem, \
-    QGraphicsSceneMouseEvent, QApplication, QWidget
+    QGraphicsSceneMouseEvent, QApplication, QWidget, QComboBox
 
 import pyqtgraph as pg
 
@@ -124,8 +126,12 @@ def fit_kernel(data: np.ndarray, kernel: str) -> \
 
 
 class ViolinItem(pg.GraphicsObject):
+    RugPlot = namedtuple("RugPlot", "support, density")
+    AREA, COUNT, WIDTH = range(3)
+
     def __init__(self, data: np.ndarray, color: QColor, kernel: str,
-                 show_rug: bool, orientation: Qt.Orientations):
+                 scale: int, show_rug: bool, orientation: Qt.Orientations):
+        self.__scale = scale
         self.__show_rug_plot = show_rug
         self.__orientation = orientation
 
@@ -136,13 +142,13 @@ def __init__(self, data: np.ndarray, color: QColor, kernel: str,
         self.__violin_path: QPainterPath = self._create_violin(data)
         self.__violin_brush: QBrush = QBrush(color)
 
-        self.__rug_plot_data: Tuple = self._create_rug_plot(data)
+        self.__rug_plot_data: ViolinItem.RugPlot = self._create_rug_plot(data)
 
         super().__init__()
 
     @property
-    def kde(self) -> KernelDensity:
-        return self.__kde
+    def density(self) -> np.ndarray:
+        return self.__rug_plot_data.density
 
     @property
     def violin_width(self) -> float:
@@ -180,6 +186,7 @@ def _create_violin(self, data: np.ndarray) -> QPainterPath:
             x = np.linspace(data.min() - self.__bandwidth * 2,
                             data.max() + self.__bandwidth * 2, 1000)
             p = np.exp(self.__kde.score_samples(x.reshape(-1, 1)))
+            p = self.__apply_scaling(p, len(data))
 
         if self.__orientation == Qt.Vertical:
             pts = [QPointF(pi, xi) for xi, pi in zip(x, p)]
@@ -194,13 +201,23 @@ def _create_violin(self, data: np.ndarray) -> QPainterPath:
         path.addPolygon(polygon)
         return path
 
+    def __apply_scaling(self, density: np.ndarray, n_data: int) -> np.ndarray:
+        if self.__scale == self.AREA:
+            return density
+        elif self.__scale == self.COUNT:
+            return density * n_data / density.max()
+        elif self.__scale == self.WIDTH:
+            return density / density.max()
+        else:
+            raise NotImplementedError
+
     def _create_rug_plot(self, data: np.ndarray) -> Tuple:
-        unique_data = np.unique(data)
         if self.__kde is None:
-            return unique_data, np.zeros(unique_data.size)
+            return self.RugPlot(data, np.zeros(data.size))
 
-        density = np.exp(self.__kde.score_samples(unique_data.reshape(-1, 1)))
-        return unique_data, density
+        density = np.exp(self.__kde.score_samples(data.reshape(-1, 1)))
+        density = self.__apply_scaling(density, len(data))
+        return self.RugPlot(data, density)
 
 
 class BoxItem(pg.GraphicsObject):
@@ -216,7 +233,7 @@ def __init__(self, data: np.ndarray, rect: QRectF,
     def boundingRect(self) -> QRectF:
         return self.__bounding_rect
 
-    def paint(self, painter: QPainter, _, widget: QWidget):
+    def paint(self, painter: QPainter, _, widget: Optional[QWidget]):
         painter.save()
 
         q0, q25, q75, q100 = self.__box_plot_data
@@ -227,7 +244,7 @@ def paint(self, painter: QPainter, _, widget: QWidget):
             quartile1 = QPointF(q0, 0), QPointF(q100, 0)
             quartile2 = QPointF(q25, 0), QPointF(q75, 0)
 
-        factor = widget.devicePixelRatio()
+        factor = 1 if widget is None else widget.devicePixelRatio()
         painter.setPen(pg.mkPen(QColor(Qt.black), width=2 * factor))
         painter.drawLine(*quartile1)
         painter.setPen(pg.mkPen(QColor(Qt.black), width=6 * factor))
@@ -261,13 +278,9 @@ def value(self) -> float:
 
 
 class StripItem(pg.ScatterPlotItem):
-    def __init__(self, data: np.ndarray, kde: KernelDensity,
+    def __init__(self, data: np.ndarray, density: np.ndarray,
                  color: QColor, orientation: Qt.Orientations):
-        if kde is not None:
-            lim = np.exp(kde.score_samples(data.reshape(-1, 1)))
-        else:
-            lim = np.zeros(data.size)
-        x = np.random.RandomState(0).uniform(-lim, lim)
+        x = np.random.RandomState(0).uniform(-density, density)
         x, y = (x, data) if orientation == Qt.Vertical else (data, x)
         color = color.lighter(150)
         super().__init__(x=x, y=y, size=5, brush=pg.mkBrush(color))
@@ -317,7 +330,7 @@ class ViolinPlot(pg.PlotWidget):
     VIOLIN_PADDING_FACTOR = 1.25
     selection_changed = Signal(list, list)
 
-    def __init__(self, parent: OWWidget, kernel: str,
+    def __init__(self, parent: OWWidget, kernel: str, scale: int,
                  orientation: Qt.Orientations, show_box_plot: bool,
                  show_strip_plot: bool, show_rug_plot: bool, sort_items: bool):
 
@@ -329,6 +342,7 @@ def __init__(self, parent: OWWidget, kernel: str,
 
         # settings
         self.__kernel = kernel
+        self.__scale = scale
         self.__orientation = orientation
         self.__show_box_plot = show_box_plot
         self.__show_strip_plot = show_strip_plot
@@ -396,6 +410,11 @@ def set_kernel(self, kernel: str):
             self.__kernel = kernel
             self._plot_data()
 
+    def set_scale(self, scale: int):
+        if self.__scale != scale:
+            self.__scale = scale
+            self._plot_data()
+
     def set_orientation(self, orientation: Qt.Orientations):
         if self.__orientation != orientation:
             self.__orientation = orientation
@@ -491,6 +510,9 @@ def _set_axes(self):
         self.getAxis("left" if vertical else "bottom").setLabel(value_title)
         self.getAxis("bottom" if vertical else "left").setLabel(group_title)
 
+        if self.__group_var is None:
+            self.getAxis("bottom" if vertical else "left").setTicks([])
+
     def _plot_data(self):
         # save selection ranges
         ranges = self._selection_ranges
@@ -516,7 +538,7 @@ def _plot_data(self):
     def _set_violin_item(self, values: np.ndarray, color: QColor):
         values = values[~np.isnan(values)]
 
-        violin = ViolinItem(values, color, self.__kernel,
+        violin = ViolinItem(values, color, self.__kernel, self.__scale,
                             self.__show_rug_plot, self.__orientation)
         self.addItem(violin)
         self.__violin_items.append(violin)
@@ -531,7 +553,7 @@ def _set_violin_item(self, values: np.ndarray, color: QColor):
         self.addItem(median)
         self.__median_items.append(median)
 
-        strip = StripItem(values, violin.kde, color, self.__orientation)
+        strip = StripItem(values, violin.density, color, self.__orientation)
         strip.setVisible(self.__show_strip_plot)
         self.addItem(strip)
         self.__strip_items.append(strip)
@@ -658,7 +680,8 @@ class Error(OWWidget.Error):
         not_enough_instances = Msg("Plotting requires at least two instances.")
 
     KERNELS = ["gaussian", "epanechnikov", "linear"]
-    KERNEL_LABELS = ["Normal kernel", "Epanechnikov kernel", "Linear kernel"]
+    KERNEL_LABELS = ["Normal", "Epanechnikov", "Linear"]
+    SCALE_LABELS = ["Area", "Count", "Width"]
 
     settingsHandler = DomainContextHandler()
     value_var = ContextSetting(None)
@@ -671,6 +694,7 @@ class Error(OWWidget.Error):
     order_violins = Setting(False)
     orientation_index = Setting(1)  # Vertical
     kernel_index = Setting(0)  # Normal kernel
+    scale_index = Setting(0)  # Area
     selection_ranges = Setting([], schema_only=True)
     visual_settings = Setting({}, schema_only=True)
 
@@ -686,6 +710,7 @@ def __init__(self):
         self._value_var_view: ListViewSearch = None
         self._group_var_view: ListViewSearch = None
         self._order_violins_cb: QCheckBox = None
+        self._scale_combo: QComboBox = None
         self.selection = []
         self.__pending_selection: List = self.selection_ranges
 
@@ -700,7 +725,8 @@ def setup_gui(self):
 
     def _add_graph(self):
         box = gui.vBox(self.mainArea)
-        self.graph = ViolinPlot(self, self.kernel, self.orientation,
+        self.graph = ViolinPlot(self, self.kernel,
+                                self.scale_index, self.orientation,
                                 self.show_box_plot, self.show_strip_plot,
                                 self.show_rug_plot, self.order_violins)
         self.graph.selection_changed.connect(self.__selection_changed)
@@ -754,8 +780,7 @@ def _add_controls(self):
                      callback=self.apply_group_var_sorting)
 
         box = gui.vBox(self.controlArea, "Display",
-                       sizePolicy=(QSizePolicy.Minimum, QSizePolicy.Maximum),
-                       addSpace=False)
+                       sizePolicy=(QSizePolicy.Minimum, QSizePolicy.Maximum))
         gui.checkBox(box, self, "show_box_plot", "Box plot",
                      callback=self.__show_box_plot_changed)
         gui.checkBox(box, self, "show_strip_plot", "Strip plot",
@@ -772,13 +797,15 @@ def _add_controls(self):
                          callback=self.__orientation_changed)
 
         box = gui.vBox(self.controlArea, "Density Estimation",
-                       sizePolicy=(QSizePolicy.Minimum, QSizePolicy.Maximum),
-                       addSpace=False)
+                       sizePolicy=(QSizePolicy.Minimum, QSizePolicy.Maximum))
         gui.comboBox(box, self, "kernel_index", items=self.KERNEL_LABELS,
+                     label="Kernel:", labelWidth=60, orientation=Qt.Horizontal,
                      callback=self.__kernel_changed)
-
-        # stretch over buttonsArea
-        self.left_side.layout().setStretch(0, 9999999)
+        self._scale_combo = gui.comboBox(
+            box, self, "scale_index", items=self.SCALE_LABELS,
+            label="Scale:", labelWidth=60, orientation=Qt.Horizontal,
+            callback=self.__scale_changed
+        )
 
         self._set_input_summary(None)
         self._set_output_summary(None)
@@ -796,7 +823,7 @@ def __group_var_changed(self, selection: QItemSelection):
             return
         self.group_var = selection.indexes()[0].data(gui.TableVariable)
         self.apply_value_var_sorting()
-        self.enable_order_violins_cb()
+        self.enable_controls()
         self.setup_plot()
         self.__selection_changed([], [])
 
@@ -818,6 +845,9 @@ def __orientation_changed(self):
     def __kernel_changed(self):
         self.graph.set_kernel(self.kernel)
 
+    def __scale_changed(self):
+        self.graph.set_scale(self.scale_index)
+
     @property
     def kernel(self) -> str:
         # pylint: disable=invalid-sequence-index
@@ -841,7 +871,7 @@ def set_data(self, data: Optional[Table]):
         self.set_list_view_selection()
         self.apply_value_var_sorting()
         self.apply_group_var_sorting()
-        self.enable_order_violins_cb()
+        self.enable_controls()
         self.setup_plot()
         self.apply_selection()
 
@@ -958,8 +988,10 @@ def _ensure_selection_visible(view):
         if len(selection) == 1:
             view.scrollTo(selection[0])
 
-    def enable_order_violins_cb(self):
-        self._order_violins_cb.setEnabled(self.group_var is not None)
+    def enable_controls(self):
+        enable = self.group_var is not None or not self.data
+        self._order_violins_cb.setEnabled(enable)
+        self._scale_combo.setEnabled(enable)
 
     def setup_plot(self):
         self.graph.clear_plot()