owheatmap: Add column color annotations

Author: ales-erjavec

Orange/widgets/visualize/owheatmap.py

@@ -17,7 +17,8 @@
 from AnyQt.QtCore import Qt, QSize, QRectF, QObject
 
 from orangewidget.utils.combobox import ComboBox, ComboBoxSearch
-from Orange.data import Domain, Table, Variable, DiscreteVariable
+from Orange.data import Domain, Table, Variable, DiscreteVariable, \
+    ContinuousVariable
 from Orange.data.sql.table import SqlTable
 import Orange.distance
 
@@ -412,13 +413,36 @@ def _(idx, cb=cb):
         form.addRow("Text", self.annotation_text_cb)
         form.addRow("Color", self.row_side_color_cb)
         box.layout().addWidget(annotbox)
-        posbox = gui.vBox(box, "Column Labels Position", addSpace=False)
-        posbox.setFlat(True)
+        annotbox = QGroupBox("Column annotations", flat=True)
+        form = QFormLayout(
+            annotbox,
+            formAlignment=Qt.AlignLeft,
+            labelAlignment=Qt.AlignLeft,
+            fieldGrowthPolicy=QFormLayout.AllNonFixedFieldsGrow
+        )
+        self.col_side_color_model = DomainModel(
+            placeholder="(None)",
+            valid_types=(DiscreteVariable, ContinuousVariable),
+            parent=self
+        )
+        self.col_side_color_cb = cb = ComboBoxSearch(
+            sizeAdjustPolicy=QComboBox.AdjustToMinimumContentsLength,
+            minimumContentsLength=12
+        )
+        self.col_side_color_cb.setModel(self.col_side_color_model)
+        self.column_annotation_color_var = None
+        self.col_side_color_cb.activated.connect(self.__set_column_annotation_color_key_index)
+        # posbox = gui.vBox(box, "Column Labels Position", addSpace=False)
+        # posbox.setFlat(True)
         cb = gui.comboBox(
-            posbox, self, "column_label_pos",
+            None, self, "column_label_pos",
             callback=self.update_column_annotations)
         cb.setModel(create_list_model(ColumnLabelsPosData, parent=self))
         cb.setCurrentIndex(self.column_label_pos)
+        form.addRow("Color", self.col_side_color_cb)
+        form.addRow("Label position", cb)
+        box.layout().addWidget(annotbox)
+
         gui.checkBox(self.controlArea, self, "keep_aspect",
                      "Keep aspect ratio", box="Resize",
                      callback=self.__aspect_mode_changed)
@@ -596,7 +620,7 @@ def set_dataset(self, data=None):
             self.row_split_model.set_domain(data.domain)
             self.col_annot_data = data.transpose(data[:0].transform(Domain(data.domain.attributes)))
             self.col_split_model.set_domain(self.col_annot_data.domain)
-
+            self.col_side_color_model.set_domain(self.col_annot_data.domain)
             if data.domain.has_discrete_class:
                 self.split_by_var = data.domain.class_var
             else:
@@ -633,7 +657,7 @@ def set_split_variable(self, var):
             self.update_heatmaps()
 
     def __on_split_cols_activated(self):
-        self.set_column_split_key(self.col_split_cb.currentData(Qt.UserRole))
+        self.set_column_split_key(self.col_split_cb.currentData(Qt.EditRole))
 
     def set_column_split_key(self, key):
         if key != self.split_columns_key:
@@ -812,7 +836,9 @@ def construct_heatmaps(self, data, group_var=None, column_split_key=None) -> 'Pa
 
         self.__update_clustering_enable_state(effective_data)
 
-        parts = self._make_parts(effective_data, group_var, column_split_key)
+        parts = self._make_parts(
+            effective_data, group_var,
+            column_split_key.name if column_split_key is not None else None)
         # Restore/update the row/columns items descriptions from cache if
         # available
         rows_cache_key = (group_var,
@@ -882,9 +908,15 @@ def setup_scene(self, parts, data):
             col_names=columns,
         )
         widget.setHeatmaps(parts)
+
         side = self.row_side_colors()
         if side is not None:
             widget.setRowSideColorAnnotations(side[0], side[1], name=side[2].name)
+
+        side = self.column_side_colors()
+        if side is not None:
+            widget.setColumnSideColorAnnotations(side[0], side[1], name=side[2].name)
+
         widget.setColumnLabelsPosition(self._column_label_pos)
         widget.setAspectRatioMode(
             Qt.KeepAspectRatio if self.keep_aspect else Qt.IgnoreAspectRatio
@@ -1111,6 +1143,28 @@ def _colorize(self, var: Variable, data: np.ndarray) -> Tuple[np.ndarray, ColorM
         else:
             raise TypeError
 
+    def __set_column_annotation_color_key_index(self, index):
+        self.column_annotation_color_var = self.col_side_color_cb.itemData(index, Qt.EditRole)
+        colors = self.column_side_colors()
+        if colors is not None:
+            self.scene.widget.setColumnSideColorAnnotations(
+                colors[0], colors[1], colors[2].name,
+            )
+        else:
+            self.scene.widget.setColumnSideColorAnnotations(None)
+
+    def column_side_colors(self):
+        var = self.column_annotation_color_var
+        if var is None:
+            return None
+        table = self.col_annot_data
+        var = table.domain[var]
+        column_data = column_data_from_table(table, var)
+        data, colormap = self._colorize(var, column_data)
+        if var.is_continuous:
+            colormap.span = (np.nanmin(column_data), np.nanmax(column_data))
+        return data, colormap, var
+
     def update_column_annotations(self):
         widget = self.scene.widget
         if self.data is not None and widget is not None:

Orange/widgets/visualize/utils/heatmap.py

@@ -564,6 +564,16 @@ def setHeatmaps(self, parts: 'Parts') -> None:
             self.TopLabelsRow, RightLabelColumn - 1,
         )
 
+        col_color_annotation_header = QGraphicsSimpleTextItem("", self)
+        grid.addItem(SimpleLayoutItem(
+            col_color_annotation_header, anchor=(1, 1), anchorItem=(1, 1),
+            resizeContents=True,
+            aspectMode=Qt.KeepAspectRatio,
+            sizePolicy=QSizePolicy(QSizePolicy.Maximum, QSizePolicy.Preferred),
+        ),
+            TopAnnotationRow, 0, 1, Col0, alignment=Qt.AlignRight
+        )
+
         legend = GradientLegendWidget(
             parts.span[0], parts.span[1],
             colormap,