Merge pull request #4504 from janezd/scatterplot-z-order

[ENH] Impose a sensible z-order to points in projections

Author: VesnaT

Orange/widgets/visualize/owscatterplotgraph.py

@@ -128,7 +128,8 @@ def restoreAnchor(self, anchors):
         anchor, parentanchor = anchors
         self.anchor(*bound_anchor_pos(anchor, parentanchor))
 
-    def paint(self, painter, option, widget=None):
+    # pylint: disable=arguments-differ
+    def paint(self, painter, _option, _widget=None):
         painter.setPen(self.__pen)
         painter.setBrush(self.__brush)
         rect = self.contentsRect()
@@ -244,22 +245,60 @@ def __init__(self, scatter_widget, parent=None, _="None", view_box=InteractiveVi
 
 
 class ScatterPlotItem(pg.ScatterPlotItem):
-    """PyQtGraph's ScatterPlotItem calls updateSpots at any change of sizes/colors/symbols,
-    which then rebuilds the stored pixmaps for each symbol. Because Orange calls
-    set* function in succession, we postpone updateSpots() to paint()."""
+    """
+    Modifies the behaviour of ScatterPlotItem as follows:
+
+    - Add z-index. ScatterPlotItem paints points in order of appearance in
+      self.data. Plotting by z-index is achieved by sorting before calling
+      super().paint() and re-sorting afterwards. Re-sorting (instead of
+      storing the original data) is needed because the inherited paint
+      may modify the data.
+
+    - Prevent multiple calls to updateSpots. ScatterPlotItem calls updateSpots
+      at any change of sizes/colors/symbols, which then rebuilds the stored
+      pixmaps for each symbol. Orange calls set* functions in succession,
+      so we postpone updateSpots() to paint()."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._update_spots_in_paint = False
+        self._z_mapping = None
+        self._inv_mapping = None
+
+    def setZ(self, z):
+        """
+        Set z values for all points.
+
+        Points with higher values are plotted on top of those with lower.
 
-    _update_spots_in_paint = False
+        Args:
+            z (np.ndarray or None): a vector of z values
+        """
+        if z is None:
+            self._z_mapping = self._inv_mapping = None
+        else:
+            assert len(z) == len(self.data)
+            self._z_mapping = np.argsort(z)
+            self._inv_mapping = np.argsort(self._z_mapping)
 
     def updateSpots(self, dataSet=None):  # pylint: disable=unused-argument
         self._update_spots_in_paint = True
         self.update()
 
+    # pylint: disable=arguments-differ
     def paint(self, painter, option, widget=None):
-        if self._update_spots_in_paint:
-            self._update_spots_in_paint = False
-            super().updateSpots()
-        painter.setRenderHint(QPainter.SmoothPixmapTransform, True)
-        super().paint(painter, option, widget)
+        try:
+            if self._z_mapping is not None:
+                assert len(self._z_mapping) == len(self.data)
+                self.data = self.data[self._z_mapping]
+            if self._update_spots_in_paint:
+                self._update_spots_in_paint = False
+                super().updateSpots()
+            painter.setRenderHint(QPainter.SmoothPixmapTransform, True)
+            super().paint(painter, option, widget)
+        finally:
+            if self._inv_mapping is not None:
+                self.data = self.data[self._inv_mapping]
 
 
 def _define_symbols():
@@ -581,7 +620,7 @@ def _get_jittering_tooltip(self):
     def update_jittering(self):
         self.update_tooltip()
         x, y = self.get_coordinates()
-        if x is None or not len(x) or self.scatterplot_item is None:
+        if x is None or len(x) == 0 or self.scatterplot_item is None:
             return
         self._update_plot_coordinates(self.scatterplot_item, x, y)
         self._update_plot_coordinates(self.scatterplot_item_sel, x, y)
@@ -766,7 +805,9 @@ def _jitter_data(self, x, y, span_x=None, span_y=None):
 
     def _update_plot_coordinates(self, plot, x, y):
         """
-        Change the coordinates of points while keeping other properites
+        Change the coordinates of points while keeping other properites.
+
+        Asserts that the number of points stays the same.
 
         Note. Pyqtgraph does not offer a method for this: setting coordinates
         invalidates other data. We therefore retrieve the data to set it
@@ -776,6 +817,7 @@ def _update_plot_coordinates(self, plot, x, y):
         update for every property would essentially reset the graph, which
         can be time consuming.
         """
+        assert self.n_shown == len(x) == len(y)
         data = dict(x=x, y=y)
         for prop in ('pen', 'brush', 'size', 'symbol', 'data',
                      'sourceRect', 'targetRect'):
@@ -793,7 +835,7 @@ def update_coordinates(self):
         the complete update by calling `reset_graph` instead of this method.
         """
         x, y = self.get_coordinates()
-        if x is None or not len(x):
+        if x is None or len(x) == 0:
             return
         if self.scatterplot_item is None:
             if self.sample_indices is None:
@@ -1012,9 +1054,9 @@ def _get_continuous_colors(self, c_data, subset):
         # Reuse pens and brushes with the same colors because PyQtGraph then
         # builds smaller pixmap atlas, which makes the drawing faster
 
-        def reuse(cache, fn, *args):
+        def reuse(cache, fun, *args):
             if args not in cache:
-                cache[args] = fn(args)
+                cache[args] = fun(args)
             return cache[args]
 
         def create_pen(col):
@@ -1072,7 +1114,7 @@ def _get_discrete_colors(self, c_data, subset):
 
     def update_colors(self):
         """
-        Trigger an update of point sizes
+        Trigger an update of point colors
 
         The method calls `self.get_colors`, which in turn calls the widget's
         `get_color_data` to get the indices in the pallette. `get_colors`
@@ -1084,6 +1126,7 @@ def update_colors(self):
             pen_data, brush_data = self.get_colors()
             self.scatterplot_item.setPen(pen_data, update=False, mask=None)
             self.scatterplot_item.setBrush(brush_data, mask=None)
+        self.update_z_values()
         self.update_legends()
         self.update_density()
 
@@ -1128,6 +1171,7 @@ def update_selection_colors(self):
         pen, brush = self.get_colors_sel()
         self.scatterplot_item_sel.setPen(pen, update=False, mask=None)
         self.scatterplot_item_sel.setBrush(brush, mask=None)
+        self.update_z_values()
 
     def get_colors_sel(self):
         """
@@ -1292,6 +1336,52 @@ def update_shapes(self):
             self.scatterplot_item.setSymbol(shape_data)
         self.update_legends()
 
+    def update_z_values(self):
+        """
+        Set z-values for point in the plot
+
+        The order is as follows:
+        - selected points that are also in the subset on top,
+        - followed by selected points,
+        - followed by points from the subset,
+        - followed by the rest.
+        Within each of these four groups, points are ordered by their colors.
+
+        Points with less frequent colors are above those with more frequent.
+        The points for which the value for the color is missing are at the
+        bottom of their respective group.
+        """
+        if not self.scatterplot_item:
+            return
+
+        subset = self.master.get_subset_mask()
+        c_data = self.master.get_color_data()
+        if subset is None and self.selection is None and c_data is None:
+            self.scatterplot_item.setZ(None)
+            return
+
+        z = np.zeros(self.n_shown)
+
+        if subset is not None:
+            subset = self._filter_visible(subset)
+            z[subset] += 1000
+
+        if self.selection is not None:
+            z[self._filter_visible(self.selection) != 0] += 2000
+
+        if c_data is not None:
+            c_nan = np.isnan(c_data)
+            vis_data = self._filter_visible(c_data)
+            vis_nan = np.isnan(vis_data)
+            z[vis_nan] -= 999
+            if not self.master.is_continuous_color():
+                dist = np.bincount(c_data[~c_nan].astype(int))
+                vis_knowns = vis_data[~vis_nan].astype(int)
+                argdist = np.argsort(dist)
+                z[~vis_nan] -= argdist[vis_knowns]
+
+        self.scatterplot_item.setZ(z)
+
     def update_grid_visibility(self):
         """Show or hide the grid"""
         self.plot_widget.showGrid(x=self.show_grid, y=self.show_grid)