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
 
@@ -26,7 +27,7 @@
 from Orange.widgets.utils.itemmodels import DomainModel
 from Orange.widgets.utils.stickygraphicsview import StickyGraphicsView
 from Orange.widgets.utils.graphicsview import GraphicsWidgetView
-from Orange.widgets.utils.colorpalettes import DiscretePalette, Palette
+from Orange.widgets.utils.colorpalettes import Palette
 
 from Orange.widgets.utils.annotated_data import (create_annotated_table,
                                                  ANNOTATED_DATA_SIGNAL_NAME)
@@ -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
@@ -1065,7 +1097,7 @@ def row_side_colors(self):
         merges = self._merge_row_indices()
         if merges is not None:
             column_data = aggregate(var, column_data, merges)
-        data, colormap = self._colorize(var, column_data)
+        data, colormap = colorize(var, column_data)
         if var.is_continuous:
             span = (np.nanmin(column_data), np.nanmax(column_data))
             if np.any(np.isnan(span)):
@@ -1091,27 +1123,27 @@ def update_row_side_colors(self):
         else:
             widget.setRowSideColorAnnotations(colors[0], colors[1], colors[2].name)
 
-    def _colorize(self, var: Variable, data: np.ndarray) -> Tuple[np.ndarray, ColorMap]:
-        palette = var.palette  # type: Palette
-        colors = np.array(
-            [[c.red(), c.green(), c.blue()] for c in palette.qcolors_w_nan],
-            dtype=np.uint8,
-        )
-        if var.is_discrete:
-            mask = np.isnan(data)
-            data[mask] = -1
-            data = data.astype(int)
-            if mask.any():
-                values = (*var.values, "N/A")
+    def __set_column_annotation_color_key_index(self, index: int):
+        key = self.col_side_color_cb.itemData(index, Qt.EditRole)
+        self.set_column_annotation_color_key(key)
+
+    def set_column_annotation_color_key(self, key):
+        if self.col_side_color_model != key:
+            self.column_annotation_color_var = key
+            colors = self.column_side_colors()
+            if colors is not None:
+                self.scene.widget.setColumnSideColorAnnotations(
+                    colors[0], colors[1], colors[2].name,
+                )
             else:
-                values = var.values
-                colors = colors[: -1]
-            return data, CategoricalColorMap(colors, values)
-        elif var.is_continuous:
-            cmap = GradientColorMap(colors[:-1])
-            return data, cmap
-        else:
-            raise TypeError
+                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
+        return color_annotation_data(table, var)
 
     def update_column_annotations(self):
         widget = self.scene.widget
@@ -1320,6 +1352,40 @@ def column_data_from_table(
     return data
 
 
+def color_annotation_data(
+        table: Table, var: Union[int, str, Variable]
+) -> Tuple[np.ndarray, ColorMap, Variable]:
+    var = table.domain[var]
+    column_data = column_data_from_table(table, var)
+    data, colormap = colorize(var, column_data)
+    return data, colormap, var
+
+
+def colorize(var: Variable, data: np.ndarray) -> Tuple[np.ndarray, ColorMap]:
+    palette = var.palette  # type: Palette
+    colors = np.array(
+        [[c.red(), c.green(), c.blue()] for c in palette.qcolors_w_nan],
+        dtype=np.uint8,
+    )
+    if var.is_discrete:
+        mask = np.isnan(data)
+        data = data.astype(int)
+        data[mask] = -1
+        if mask.any():
+            values = (*var.values, "N/A")
+        else:
+            values = var.values
+            colors = colors[: -1]
+        return data, CategoricalColorMap(colors, values)
+    elif var.is_continuous:
+        span = np.nanmin(data), np.nanmax(data)
+        if np.any(np.isnan(span)):
+            span = 0, 1.
+        return data, GradientColorMap(colors[:-1], span=span)
+    else:
+        raise TypeError
+
+
 def aggregate(
         var: Variable, data: np.ndarray, groupindices: Sequence[Sequence[int]],
 ) -> np.ndarray:

Orange/widgets/visualize/tests/test_owheatmap.py

@@ -290,6 +290,7 @@ def test_row_color_annotations_with_na(self):
         data = data.transform(
             Domain(data.domain.attributes[:10], data.domain.class_vars,
                    data.domain.metas + (data.domain["diau g"], )))
+        data.ensure_copy()
         data.Y[:3] = np.nan
         data.metas[:3, -1] = np.nan
         self.send_signal(widget.Inputs.data, data, widget=widget)
@@ -304,6 +305,46 @@ def test_row_color_annotations_with_na(self):
         widget.set_annotation_color_var(None)
         self.assertFalse(widget.scene.widget.right_side_colors[0].isVisible())
 
+    def test_col_color_annotations(self):
+        widget = self.widget
+        data = self.brown_selected[::5]
+        data = data.transform(
+            Domain(data.domain.attributes[:10], data.domain.class_vars,
+                   data.domain.metas + (data.domain["diau g"],)))
+        data.ensure_copy()
+        data_t = data.transpose(data)
+        self.send_signal(widget.Inputs.data, data_t, widget=widget)
+        # discrete
+        widget.set_column_annotation_color_key(data.domain["function"])
+        self.assertTrue(widget.scene.widget.top_side_colors[0].isVisible())
+        # continuous
+        widget.set_column_annotation_color_key(data.domain["diau g"])
+        widget.set_column_annotation_color_key(None)
+        self.assertFalse(widget.scene.widget.top_side_colors[0].isVisible())
+
+    def test_col_color_annotations_with_na(self):
+        widget = self.widget
+        data = self.brown_selected[::5]
+        data = data.transform(
+            Domain(data.domain.attributes[:10], data.domain.class_vars,
+                   data.domain.metas + (data.domain["diau g"], )))
+        data.ensure_copy()
+        data.Y[:3] = np.nan
+        data.metas[:3, -1] = np.nan
+        data_t = data.transpose(data)
+        self.send_signal(widget.Inputs.data, data_t, widget=widget)
+        widget.set_column_annotation_color_key(data.domain["function"])
+        self.assertTrue(widget.scene.widget.top_side_colors[0].isVisible())
+        widget.set_column_annotation_color_key(data.domain["diau g"])
+        data.Y[:] = np.nan
+        data.metas[:, -1] = np.nan
+        data_t = data.transpose(data)
+        self.send_signal(widget.Inputs.data, data_t, widget=widget)
+        widget.set_column_annotation_color_key(data.domain["function"])
+        widget.set_column_annotation_color_key(data.domain["diau g"])
+        widget.set_column_annotation_color_key(None)
+        self.assertFalse(widget.scene.widget.top_side_colors[0].isVisible())
+
     def test_summary(self):
         """Check if status bar is updated when data is received"""
         info = self.widget.info