heatmap: Add column color annotations

Author: ales-erjavec

Orange/widgets/visualize/utils/heatmap.py

@@ -278,7 +278,7 @@ class Position(enum.IntFlag):
 
     # Start row/column where the heatmap items are inserted
     # (after the titles/legends/dendrograms)
-    Row0 = 4
+    Row0 = 5
     Col0 = 3
     # The (color) legend row and column
     LegendRow, LegendCol = 0, 4
@@ -290,6 +290,8 @@ class Position(enum.IntFlag):
     DendrogramRow = 2
     # The row for top column annotation labels
     TopLabelsRow = 3
+    # Top color annotation row
+    TopAnnotationRow = 4
     # Vertical split title column
     GroupTitleColumn = 0
 
@@ -312,6 +314,7 @@ def __init__(self, parent=None, **kwargs):
         self.col_dendrograms = []  # type: List[Optional[DendrogramWidget]]
         self.row_dendrograms = []  # type: List[Optional[DendrogramWidget]]
         self.right_side_colors = []  # type: List[Optional[GraphicsPixmapWidget]]
+        self.top_side_colors = []  # type: List[Optional[GraphicsPixmapWidget]]
         self.heatmap_colormap_legend = None
         self.bottom_legend_container = None
         self.__layout = GridLayout()
@@ -343,6 +346,7 @@ def clear(self):
         self.col_dendrograms = []
         self.row_dendrograms = []
         self.right_side_colors = []
+        self.top_side_colors = []
         self.heatmap_colormap_legend = None
         self.bottom_legend_container = None
         self.parts = None
@@ -363,12 +367,14 @@ def setHeatmaps(self, parts: 'Parts') -> None:
         # The row for the horizontal dendrograms
         DendrogramRow = self.DendrogramRow
         RightLabelColumn = Col0 + 2 * M + 1
+        TopAnnotationRow = self.TopAnnotationRow
         TopLabelsRow = self.TopLabelsRow
         BottomLabelsRow = Row0 + N
         colormap = self.__colormap
         column_dendrograms: List[Optional[DendrogramWidget]] = [None] * M
         row_dendrograms: List[Optional[DendrogramWidget]] = [None] * N
         right_side_colors: List[Optional[GraphicsPixmapWidget]] = [None] * N
+        top_side_colors: List[Optional[GraphicsPixmapWidget]] = [None] * M
 
         data = parts.data
         if parts.col_names is None:
@@ -492,8 +498,6 @@ def setHeatmaps(self, parts: 'Parts') -> None:
                 objectName="row-labels-right"
             )
             labelslist.setMaximumWidth(300)
-            pm = QPixmap(1, rowitem.size)
-            pm.fill(Qt.transparent)
             rowauxsidecolor = GraphicsPixmapWidget(
                 parent=self, visible=False,
                 scaleContents=True, aspectMode=Qt.IgnoreAspectRatio,
@@ -520,10 +524,20 @@ def setHeatmaps(self, parts: 'Parts') -> None:
                 visible=self.__columnLabelPosition & Position.Top,
                 objectName="column-labels-top",
             )
+            colauxsidecolor = GraphicsPixmapWidget(
+                parent=self, visible=False,
+                scaleContents=True, aspectMode=Qt.IgnoreAspectRatio,
+                sizePolicy=QSizePolicy(QSizePolicy.Ignored,
+                                       QSizePolicy.Maximum),
+                minimumSize=QSizeF(-1, 10)
+            )
+
             grid.addItem(labelslist, TopLabelsRow, Col0 + 2 * j + 1,
                          Qt.AlignBottom | Qt.AlignLeft)
+            grid.addItem(colauxsidecolor, TopAnnotationRow, Col0 + 2 * j + 1)
             col_annotation_widgets.append(labelslist)
             col_annotation_widgets_top.append(labelslist)
+            top_side_colors[j] = colauxsidecolor
 
             # Bottom attr annotations
             labelslist = TextListWidget(
@@ -560,12 +574,33 @@ def setHeatmaps(self, parts: 'Parts') -> None:
         )
         legend.setMaximumWidth(300)
         grid.addItem(legend, self.LegendRow, self.LegendCol, 1, M * 2)
-        legend_container = QGraphicsWidget(
+
+        def container(orientation=Qt.Horizontal, **kwargs):
+            container = QGraphicsWidget(**kwargs)
+            layout = QGraphicsLinearLayout(orientation)
+            layout.setContentsMargins(0, 0, 0, 0)
+            container.setLayout(layout)
+            return container
+        # Container for color annotation legends
+        legend_container = container(
             visible=False,
-            sizePolicy=QSizePolicy(QSizePolicy.Maximum, QSizePolicy.Fixed)
+            sizePolicy=QSizePolicy(QSizePolicy.Maximum, QSizePolicy.Fixed),
+            objectName="annotation-legend-container"
         )
-        legend_container.setLayout(QGraphicsLinearLayout())
-        legend_container.layout().setContentsMargins(0, 0, 0, 0)
+        legend_container.layout().setSpacing(3)
+        legend_container_rows = container(
+            sizePolicy=QSizePolicy(QSizePolicy.Maximum, QSizePolicy.Fixed),
+            visible=False, objectName="row-annotation-legend-container"
+        )
+        legend_container_cols = container(
+            sizePolicy=QSizePolicy(QSizePolicy.Maximum, QSizePolicy.Fixed),
+            visible=False, objectName="row-annotation-legend-container"
+        )
+        # ? segfault on scene.clear() ?
+        legend_container._keep_ref = (legend_container_cols, legend_container_rows)
+        legend_container.layout().addItem(legend_container_rows)
+        legend_container.layout().addItem(legend_container_cols)
+
         grid.addItem(legend_container, BottomLabelsRow + 1, Col0 + 1, 1, M * 2 - 1,
                      alignment=Qt.AlignRight)
 
@@ -577,6 +612,7 @@ def setHeatmaps(self, parts: 'Parts') -> None:
         self.col_dendrograms = column_dendrograms
         self.row_dendrograms = row_dendrograms
         self.right_side_colors = right_side_colors
+        self.top_side_colors = top_side_colors
         self.heatmap_colormap_legend = legend
         self.bottom_legend_container = legend_container
         self.parts = parts
@@ -679,50 +715,113 @@ def setRowSideColorAnnotations(
         name: str
             Name/title for the annotation column.
         """
-        items = self.right_side_colors
         col = self.Col0 + 2 * len(self.parts.columns)
+        legend_layout = self.bottom_legend_container.layout()
+        legend_container = legend_layout.itemAt(1)
+        self.__setColorAnnotationsHelper(
+            data, colormap, name, self.right_side_colors, col, Qt.Vertical,
+            legend_container
+        )
+        legend_container.setVisible(True)
+
+    def setColumnSideColorAnnotations(
+            self, data: np.ndarray, colormap: ColorMap=None, name=""
+    ):
+        """
+        Set an optional column color annotations.
+
+        Parameters
+        ----------
+        data: Optional[np.ndarray]
+            A sequence such that it is accepted by `colormap.apply`. If None
+            then the color annotations are cleared.
+        colormap: ColorMap
+        name: str
+            Name/title for the annotation column.
+        """
+        row = self.TopAnnotationRow
+        legend_layout = self.bottom_legend_container.layout()
+        legend_container = legend_layout.itemAt(0)
+        self.__setColorAnnotationsHelper(
+            data, colormap, name, self.top_side_colors, row, Qt.Horizontal,
+            legend_container)
+        legend_container.setVisible(True)
+
+    def __setColorAnnotationsHelper(
+            self, data: np.ndarray, colormap: ColorMap, name: str,
+            items: List[GraphicsPixmapWidget], position: int,
+            orientation: Qt.Orientation, legend_container: QGraphicsWidget):
         layout = self.__layout
-        nameitem = layout.itemAt(self.TopLabelsRow, col)
-        width = QFontMetrics(self.font()).lineSpacing()
-        legend_container = self.bottom_legend_container
+        if orientation == Qt.Horizontal:
+            nameitem = layout.itemAt(position, 0)
+        else:
+            nameitem = layout.itemAt(self.TopLabelsRow, position)
+        size = QFontMetrics(self.font()).lineSpacing()
         layout_clear(legend_container.layout())
 
+        def grid_set_maximum_size(position: int, size: float):
+            if orientation == Qt.Horizontal:
+                layout.setRowMaximumHeight(position, size)
+            else:
+                layout.setColumnMaximumWidth(position, size)
+
+        def set_minimum_size(item: QGraphicsLayoutItem, size: float):
+            if orientation == Qt.Horizontal:
+                item.setMinimumHeight(size)
+            else:
+                item.setMinimumWidth(size)
+            item.updateGeometry()
+
+        def reset_minimum_size(item: QGraphicsLayoutItem):
+            set_minimum_size(item, -1)
+
         def set_hidden(item: GraphicsPixmapWidget):
             item.setVisible(False)
-            item.setMinimumWidth(-1)
-            item.updateGeometry()
+            reset_minimum_size(item,)
 
         def set_visible(item: GraphicsPixmapWidget):
             item.setVisible(True)
-            item.setMinimumWidth(10)
+            set_minimum_size(item, 10)
+
+        def set_preferred_size(item, size):
+            if orientation == Qt.Horizontal:
+                item.setPreferredHeight(size)
+            else:
+                item.setPreferredWidth(size)
             item.updateGeometry()
 
         if data is None:
             apply_all(filter(None, items), set_hidden)
-            layout.setColumnMaximumWidth(col, 0)
+            grid_set_maximum_size(position, 0)
+
             nameitem.item.setVisible(False)
             nameitem.updateGeometry()
             legend_container.setVisible(False)
             return
         else:
             apply_all(filter(None, items), set_visible)
-            layout.setColumnMaximumWidth(col, FLT_MAX)
+            grid_set_maximum_size(position, FLT_MAX)
             legend_container.setVisible(True)
 
-        parts = self.parts.rows
+        if orientation == Qt.Horizontal:
+            parts = self.parts.columns
+        else:
+            parts = self.parts.rows
         for p, item in zip(parts, items):
             if item is not None:
                 subset = data[p.normalized_indices]
                 subset = colormap.apply(subset)
-                img = qimage_from_array(subset.reshape((-1, 1, subset.shape[-1])))
+                rgbdata = subset.reshape((-1, 1, subset.shape[-1]))
+                if orientation == Qt.Horizontal:
+                    rgbdata = rgbdata.reshape((1, -1, rgbdata.shape[-1]))
+                img = qimage_from_array(rgbdata)
                 item.setPixmap(img)
                 item.setVisible(True)
-                item.setPreferredWidth(width)
+                set_preferred_size(item, size)
 
         nameitem.item.setText(name)
         nameitem.item.setVisible(True)
-        nameitem.setPreferredWidth(width)
-        nameitem.updateGeometry()
+        set_preferred_size(nameitem, size)
 
         container = legend_container.layout()
         if isinstance(colormap, CategoricalColorMap):