Repository: biolab/orange3

Issue:
Predictions widget sorts numerical columns as text
The table view of the _**Predictions**_ widget - when clicked on column header - sorts numerical values like text instead of number in the _**target**_ and _**prediction**_ columns but NOT in the rest of the data columns. (e.g. "100.0" < "58.0" because it starts with a '1')

![image](https://user-images.githubusercontent.com/5299789/115367050-20c04480-a1c6-11eb-8a0d-bf7837237858.png)



Discussion:
Thanks for reporting this.

This widget is, well, complicated.

To whomever would work on this: I think the widget should be migrated to `orangewidget.utils.itemmodels.AbstractSortTableModel`. This would improve its performance (see https://github.com/biolab/orange3/blob/master/Orange/widgets/evaluate/owpredictions.py#L404), as well as simplify obtaining the keys for sorting. The problem is that the current code expects various `QSortFilterProxyModel`'s method at several places.

Resolving patch:
diff --git a/Orange/widgets/evaluate/owpredictions.py b/Orange/widgets/evaluate/owpredictions.py
--- a/Orange/widgets/evaluate/owpredictions.py
+++ b/Orange/widgets/evaluate/owpredictions.py
@@ -12,15 +12,16 @@
 from AnyQt.QtGui import QPainter, QStandardItem, QPen, QColor
 from AnyQt.QtCore import (
     Qt, QSize, QRect, QRectF, QPoint, QLocale,
-    QModelIndex, QAbstractTableModel, QSortFilterProxyModel, pyqtSignal, QTimer,
+    QModelIndex, pyqtSignal, QTimer,
     QItemSelectionModel, QItemSelection)
 
 from orangewidget.report import plural
+from orangewidget.utils.itemmodels import AbstractSortTableModel
 
 import Orange
 from Orange.evaluation import Results
 from Orange.base import Model
-from Orange.data import ContinuousVariable, DiscreteVariable, Value, Domain
+from Orange.data import ContinuousVariable, DiscreteVariable, Domain
 from Orange.data.table import DomainTransformationError
 from Orange.data.util import get_unique_names
 from Orange.widgets import gui, settings
@@ -137,11 +138,11 @@ def __init__(self):
         self.vsplitter.layout().addWidget(self.splitter)
         self.vsplitter.layout().addWidget(self.score_table.view)
 
-    def get_selection_store(self, proxy):
-        # Both proxies map the same, so it doesn't matter which one is used
+    def get_selection_store(self, model):
+        # Both models map the same, so it doesn't matter which one is used
         # to initialize SharedSelectionStore
         if self.selection_store is None:
-            self.selection_store = SharedSelectionStore(proxy)
+            self.selection_store = SharedSelectionStore(model)
         return self.selection_store
 
     @Inputs.data
@@ -156,12 +157,10 @@ def set_data(self, data):
         else:
             # force full reset of the view's HeaderView state
             self.dataview.setModel(None)
-            model = TableModel(data, parent=None)
-            modelproxy = SortProxyModel()
-            modelproxy.setSourceModel(model)
-            self.dataview.setModel(modelproxy)
+            model = DataModel(data, parent=None)
+            self.dataview.setModel(model)
             sel_model = SharedSelectionModel(
-                self.get_selection_store(modelproxy), modelproxy, self.dataview)
+                self.get_selection_store(model), model, self.dataview)
             self.dataview.setSelectionModel(sel_model)
             if self.__pending_selection is not None:
                 self.selection = self.__pending_selection
@@ -221,6 +220,9 @@ def _set_class_values(self):
             self.selected_classes = []
 
     def handleNewSignals(self):
+        # Disconnect the model: the model and the delegate will be inconsistent
+        # between _set_class_values and update_predictions_model.
+        self.predictionsview.setModel(None)
         self._set_class_values()
         self._call_predictors()
         self._update_scores()
@@ -296,6 +298,7 @@ def _update_scores(self):
                     item.setText(f"{score:.3f}")
                 except Exception as exc:  # pylint: disable=broad-except
                     item.setToolTip(str(exc))
+                    # false pos.; pylint: disable=unsupported-membership-test
                     if scorer.name in self.score_table.shown_scores:
                         errors.append(str(exc))
                 row.append(item)
@@ -369,23 +372,28 @@ def _reordered_probabilities(self, prediction):
         return new_probs
 
     def _update_predictions_model(self):
-        results = []
         headers = []
+        all_values = []
+        all_probs = []
         for p in self._non_errored_predictors():
             values = p.results.unmapped_predicted
             target = p.predictor.domain.class_var
             if target.is_discrete:
                 # order probabilities in order from Show prob. for
                 prob = self._reordered_probabilities(p)
-                values = [Value(target, v) for v in values]
+                values = numpy.array(target.values)[values.astype(int)]
             else:
                 prob = numpy.zeros((len(values), 0))
-            results.append((values, prob))
+            all_values.append(values)
+            all_probs.append(prob)
             headers.append(p.predictor.name)
 
-        if results:
-            results = list(zip(*(zip(*res) for res in results)))
-            model = PredictionsModel(results, headers)
+        if all_values:
+            model = PredictionsModel(all_values, all_probs, headers)
+            model.list_sorted.connect(
+                partial(
+                    self._update_data_sort_order, self.predictionsview,
+                    self.dataview))
         else:
             model = None
 
@@ -393,49 +401,27 @@ def _update_predictions_model(self):
             self.selection_store.unregister(
                 self.predictionsview.selectionModel())
 
-        predmodel = PredictionsSortProxyModel()
-        predmodel.setSourceModel(model)
-        predmodel.setDynamicSortFilter(True)
-        self.predictionsview.setModel(predmodel)
-
-        self.predictionsview.setSelectionModel(
-            SharedSelectionModel(self.get_selection_store(predmodel),
-                                 predmodel, self.predictionsview))
+        self.predictionsview.setModel(model)
+        if model is not None:
+            self.predictionsview.setSelectionModel(
+                SharedSelectionModel(self.get_selection_store(model),
+                                     model, self.predictionsview))
 
         hheader = self.predictionsview.horizontalHeader()
         hheader.setSortIndicatorShown(False)
-        # SortFilterProxyModel is slow due to large abstraction overhead
-        # (every comparison triggers multiple `model.index(...)`,
-        # model.rowCount(...), `model.parent`, ... calls)
-        hheader.setSectionsClickable(predmodel.rowCount() < 20000)
-
-        self.predictionsview.model().list_sorted.connect(
-            partial(
-                self._update_data_sort_order, self.predictionsview,
-                self.dataview))
+        hheader.setSectionsClickable(True)
 
     def _update_data_sort_order(self, sort_source_view, sort_dest_view):
+        sort_source_view.horizontalHeader().setSortIndicatorShown(True)
+        sort_dest_view.horizontalHeader().setSortIndicatorShown(False)
+
         sort_dest = sort_dest_view.model()
         sort_source = sort_source_view.model()
-        sortindicatorshown = False
         if sort_dest is not None:
-            assert isinstance(sort_dest, QSortFilterProxyModel)
-            n = sort_dest.rowCount()
             if sort_source is not None and sort_source.sortColumn() >= 0:
-                sortind = numpy.argsort(
-                    [sort_source.mapToSource(sort_source.index(i, 0)).row()
-                     for i in range(n)])
-                sortind = numpy.array(sortind, int)
-                sortindicatorshown = True
+                sort_dest.setSortIndices(sort_source.mapToSourceRows(...))
             else:
-                sortind = None
-
-            sort_dest.setSortIndices(sortind)
-
-        sort_dest_view.horizontalHeader().setSortIndicatorShown(
-            False)
-        sort_source_view.horizontalHeader().setSortIndicatorShown(
-            sortindicatorshown)
+                sort_dest.setSortIndices(None)
         self.commit()
 
     def _reset_order(self):
@@ -601,20 +587,20 @@ def _commit_predictions(self):
                 [numpy.atleast_2d(cols) for cols in newcolumns])
             predictions.metas[:, -newcolumns.shape[1]:] = newcolumns
 
-        index = self.dataview.model().index
-        map_to = self.dataview.model().mapToSource
+        datamodel = self.dataview.model()
+        predmodel = self.predictionsview.model()
+        assert datamodel is not None  # because we have data
         assert self.selection_store is not None
-        rows = None
-        if self.selection_store.rows:
-            rows = [ind.row()
-                    for ind in self.dataview.selectionModel().selectedRows(0)]
-            rows.sort()
-        elif self.dataview.model().isSorted() \
-                or self.predictionsview.model().isSorted():
-            rows = list(range(len(self.data)))
-        if rows:
-            source_rows = [map_to(index(row, 0)).row() for row in rows]
-            predictions = predictions[source_rows]
+        rows = numpy.array(list(self.selection_store.rows))
+        if rows.size:
+            # Reorder rows as they are ordered in view
+            shown_rows = datamodel.mapFromSourceRows(rows)
+            rows = rows[numpy.argsort(shown_rows)]
+            predictions = predictions[rows]
+        elif datamodel.sortColumn() >= 0 \
+                or predmodel is not None and predmodel.sortColumn() > 0:
+            # No selection: output all, but in the shown order
+            predictions = predictions[datamodel.mapToSourceRows(...)]
         self.Outputs.predictions.send(predictions)
 
     @staticmethod
@@ -663,7 +649,8 @@ def merge_data_with_predictions():
                           predictions_model.data(predictions_model.index(i, j)),
                           QLocale())
                        for j, delegate in enumerate(delegates)] + \
-                      [data_model.data(data_model.index(i, j))
+                      [data_model.data(data_model.index(i, j),
+                                       role=Qt.DisplayRole)
                        for j in iter_data_cols]
 
         if self.data:
@@ -737,7 +724,7 @@ def setFormat(self, shown_probabilities=()):
             )
             self.tooltip = f"p({val})"
 
-    def displayText(self, value, _locale):
+    def displayText(self, value, _):
         try:
             value, dist = value
         except ValueError:
@@ -847,138 +834,85 @@ def drawDistBar(self, painter, rect, distribution):
         painter.restore()
 
 
-class SortProxyModel(QSortFilterProxyModel):
-    """
-    QSortFilter model used in both TableView and PredictionsView
-    """
+class PredictionsModel(AbstractSortTableModel):
     list_sorted = pyqtSignal()
 
-    def __init__(self, parent=None):
-        super().__init__(parent)
-        self.__sortInd = None
-
-    def setSortIndices(self, indices):
-        if indices is not None:
-            indices = numpy.array(indices, dtype=int)
-            if indices.shape != (self.rowCount(),):
-                raise ValueError("indices.shape != (self.rowCount(),)")
-            indices.flags.writeable = False
-
-        self.__sortInd = indices
-
-        if self.__sortInd is not None:
-            self.custom_sort(0)  # need valid order to call lessThan
-
-    def lessThan(self, left, right):
-        if self.__sortInd is None:
-            return super().lessThan(left, right)
-
-        assert not (left.parent().isValid() or right.parent().isValid()), \
-            "Not a table model"
-
-        rleft, rright = left.row(), right.row()
-        try:
-            ileft, iright = self.__sortInd[rleft], self.__sortInd[rright]
-        except IndexError:
-            return False
-        else:
-            return ileft < iright
-
-    def isSorted(self):
-        return self.__sortInd is not None
-
-    def sort(self, n, order=Qt.AscendingOrder):
-        """
-        This sort is called only on click, in other cases we manually call
-        custom_sort
-        """
-        # reset sort - otherwise when same parameters set by lessThan function
-        # clicking on header would not trigger resort
-        self.__sortInd = None
-        self.custom_sort(n, order=order)
-        self.list_sorted.emit()
-
-    def custom_sort(self, n, order=Qt.AscendingOrder):
-        """
-        When sorting is dmanded because of sort change in the other view
-        this sorting is called. It will not reset __sortInd to None.
-        """
-        if self.sortColumn() == n and self.sortOrder() == order:
-            self.invalidate()
-        else:
-            super().sort(n, order)  # need some valid sort column
-
-
-class PredictionsSortProxyModel(SortProxyModel):
-    def __init__(self, parent=None):
+    def __init__(self, values=None, probs=None, headers=None, parent=None):
         super().__init__(parent)
+        self._values = values
+        self._probs = probs
         self.__probInd = None
-
-    def setProbInd(self, indices):
-        self.__probInd = indices
-        self.invalidate()
-        self.list_sorted.emit()
-
-    def lessThan(self, left, right):
-        if self.isSorted():
-            return super().lessThan(left, right)
-
-        role = self.sortRole()
-        left_data = self.sourceModel().data(left, role)
-        right_data = self.sourceModel().data(right, role)
-
-        return self._key(left_data) < self._key(right_data)
-
-    def _key(self, prediction):
-        value, probs = prediction
-        if probs is not None:
-            if self.__probInd is not None:
-                probs = probs[self.__probInd]
-            probs = tuple(probs)
-
-        return probs, value
-
-
-class PredictionsModel(QAbstractTableModel):
-    def __init__(self, table=None, headers=None, parent=None):
-        super().__init__(parent)
-        self._table = [[]] if table is None else table
-        if headers is None:
-            headers = [None] * len(self._table)
+        if values is not None:
+            assert len(self._values) == len(self._probs) != 0
+            sizes = {len(x) for c in (values, probs) for x in c}
+            assert len(sizes) == 1
+            self.__columnCount = len(values)
+            self.__rowCount = sizes.pop()
+            if headers is None:
+                headers = [None] * self.__columnCount
+        else:
+            assert probs is None
+            assert headers is None
+            self.__columnCount = self.__rowCount = 0
         self._header = headers
-        self.__columnCount = max([len(row) for row in self._table] or [0])
 
     def rowCount(self, parent=QModelIndex()):
-        return 0 if parent.isValid() else len(self._table)
+        return 0 if parent.isValid() else self.__rowCount
 
     def columnCount(self, parent=QModelIndex()):
         return 0 if parent.isValid() else self.__columnCount
 
-    def _value(self, index):
-        return self._table[index.row()][index.column()]
-
     def data(self, index, role=Qt.DisplayRole):
         if role in (Qt.DisplayRole, Qt.EditRole):
-            return self._value(index)
+            column = index.column()
+            row = self.mapToSourceRows(index.row())
+            return self._values[column][row], self._probs[column][row]
         return None
 
     def headerData(self, section, orientation, role=Qt.DisplayRole):
-        if orientation == Qt.Vertical and role == Qt.DisplayRole:
-            return str(section + 1)
-        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
-            return (self._header[section] if section < len(self._header)
-                    else str(section))
+        if role == Qt.DisplayRole:
+            if orientation == Qt.Vertical:
+                return str(section + 1)
+            elif self._header is not None and section < len(self._header):
+                return self._header[section]
         return None
 
+    def setProbInd(self, indices):
+        self.__probInd = indices
+        self.sort(self.sortColumn())
+
+    def sortColumnData(self, column):
+        values = self._values[column]
+        probs = self._probs[column]
+        # Let us assume that probs can be None, numpy array or list of arrays
+        if probs is not None and len(probs) and len(probs[0]) \
+                and self.__probInd is not None and len(self.__probInd):
+            return probs[:, self.__probInd]
+        else:
+            return values
+
+    def sort(self, column, order=Qt.AscendingOrder):
+        super().sort(column, order)
+        self.list_sorted.emit()
+
+
+# PredictionsModel and DataModel have the same signal and sort method, but
+# extracting it into a common base class would require diamond inheritance
+# because they're derived from different classes. It's not worth it.
+class DataModel(TableModel):
+    list_sorted = pyqtSignal()
+
+    def sort(self, column, order=Qt.AscendingOrder):
+        super().sort(column, order)
+        self.list_sorted.emit()
+
 
 class SharedSelectionStore:
     """
     An object shared between multiple selection models
 
-    The object assumes that the underlying models are proxies.
-    - Method `select` and emit refer to indices in proxy.
-    - Internally, the object stores indices into source model (as int). Method
-      `select_rows` also uses internal, source-model indices.
+    The object assumes that the underlying models are AbstractSortTableModel.
+    Internally, the object stores indices of unmapped, source rows (as int).
 
     The class implements method `select` with the same signature as
     QItemSelectionModel.select. Selection models that share this object
@@ -986,10 +920,10 @@ class SharedSelectionStore:
     call `emit_selection_rows_changed` of all selection models, so they
     can emit the signal selectionChanged.
     """
-    def __init__(self, proxy):
-        # indices of selected rows in the original model, not in the proxy
+    def __init__(self, model):
+        # unmapped indices of selected rows
         self._rows: Set[int] = set()
-        self.proxy: SortProxyModel = proxy
+        self.model: AbstractSortTableModel = model
         self._selection_models: List[SharedSelectionModel] = []
 
     @property
@@ -1022,22 +956,19 @@ def select(self, selection: Union[QModelIndex, QItemSelection], flags: int):
 
         Args:
             selection (QModelIndex or QItemSelection):
-                rows to select; indices refer to the proxy model, not the source
+                rows to select; indices are mapped to rows in the view
             flags (QItemSelectionModel.SelectionFlags):
                 flags that tell whether to Clear, Select, Deselect or Toggle
         """
+        rows = set()
         if isinstance(selection, QModelIndex):
             if selection.model() is not None:
-                rows = {selection.model().mapToSource(selection).row()}
-            else:
-                rows = set()
+                rows = {selection.model().mapToSourceRows(selection.row())}
         else:
             indices = selection.indexes()
             if indices:
-                selection = indices[0].model().mapSelectionToSource(selection)
-                rows = {index.row() for index in selection.indexes()}
-            else:
-                rows = set()
+                map_to = indices[0].model().mapToSourceRows
+                rows = set(map_to([index.row() for index in indices]))
         self.select_rows(rows, flags)
 
     def select_rows(self, rows: Set[int], flags):
@@ -1046,7 +977,7 @@ def select_rows(self, rows: Set[int], flags):
 
         Args:
             selection (set of int):
-                rows to select; indices refer to the source model.
+                rows to select; indices refer to unmapped rows in model, not view
             flags (QItemSelectionModel.SelectionFlags):
                 flags that tell whether to Clear, Select, Deselect or Toggle
         """
@@ -1076,18 +1007,16 @@ def _emit_changed(self):
         changing a selection.
         """
         def map_from_source(rows):
-            from_src = self.proxy.mapFromSource
-            index = self.proxy.sourceModel().index
-            return {from_src(index(row, 0)).row() for row in rows}
+            return self.model.mapFromSourceRows(list(rows))
 
         old_rows = self._rows.copy()
         try:
             yield
         finally:
-            if self.proxy.sourceModel() is not None:
+            if self.model.rowCount() != 0:
                 deselected = map_from_source(old_rows - self._rows)
                 selected = map_from_source(self._rows - old_rows)
-                if selected or deselected:
+                if len(selected) != 0 or len(deselected) != 0:
                     for model in self._selection_models:
                         model.emit_selection_rows_changed(selected, deselected)
 
@@ -1096,28 +1025,26 @@ class SharedSelectionModel(QItemSelectionModel):
     """
     A selection model that shares the selection with its peers.
 
-    It assumes that the underlying model is a proxy.
+    It assumes that the underlying model is a AbstractTableModel.
     """
-    def __init__(self, shared_store, proxy, parent):
-        super().__init__(proxy, parent)
+    def __init__(self, shared_store, model, parent):
+        super().__init__(model, parent)
         self.store: SharedSelectionStore = shared_store
         self.store.register(self)
 
     def select(self, selection, flags):
         self.store.select(selection, flags)
 
-    def selection_from_rows(self, rows: Sequence[int],
-                            model=None) -> QItemSelection:
+    def selection_from_rows(self, rows: Sequence[int]) -> QItemSelection:
         """
         Return selection across all columns for given row indices (as ints)
 
         Args:
-            rows (sequence of int): row indices (in proxy model)
+            rows (sequence of int): row indices, as shown in the view, not model
 
         Returns: QItemSelection
         """
-        if model is None:
-            model = self.model()
+        model = self.model()
         index = model.index
         last_col = model.columnCount() - 1
         sel = QItemSelection()
@@ -1129,7 +1056,7 @@ def emit_selection_rows_changed(
             self, selected: Sequence[int], deselected: Sequence[int]):
         """
         Given a sequence of indices of selected and deselected rows,
-        emit a selectionChanged signal. Indices refer to proxy model.
+        emit a selectionChanged signal.
 
         Args:
             selected (Sequence[int]): indices of selected rows
@@ -1140,9 +1067,8 @@ def emit_selection_rows_changed(
             self.selection_from_rows(deselected))
 
     def selection(self):
-        src_sel = self.selection_from_rows(self.store.rows,
-                                           model=self.model().sourceModel())
-        return self.model().mapSelectionFromSource(src_sel)
+        rows = self.model().mapFromSourceRows(list(self.store.rows))
+        return self.selection_from_rows(rows)
 
     def hasSelection(self) -> bool:
         return bool(self.store.rows)
@@ -1151,12 +1077,13 @@ def isColumnSelected(self, *_) -> bool:
         return len(self.store.rows) == self.model().rowCount()
 
     def isRowSelected(self, row, _parent=None) -> bool:
-        return self.isSelected(self.model().index(row, 0))
+        mapped_row = self.model().mapToSourceRows(row)
+        return mapped_row in self.store.rows
 
     rowIntersectsSelection = isRowSelected
 
     def isSelected(self, index) -> bool:
-        return self.model().mapToSource(index).row() in self.store.rows
+        return self.model().mapToSourceRows(index.row()) in self.store.rows
 
     def selectedColumns(self, row: int):
         if self.isColumnSelected():
@@ -1166,14 +1093,17 @@ def selectedColumns(self, row: int):
         else:
             return []
 
+    def _selected_rows_arr(self):
+        return numpy.fromiter(self.store.rows, int, len(self.store.rows))
+
     def selectedRows(self, col: int):
-        index = self.model().sourceModel().index
-        map_from = self.model().mapFromSource
-        return [map_from(index(row, col)) for row in self.store.rows]
+        index = self.model().index
+        rows = self.model().mapFromSourceRows(self._selected_rows_arr())
+        return [index(row, col) for row in rows]
 
     def selectedIndexes(self):
         index = self.model().index
-        rows = [index.row() for index in self.selectedRows(0)]
+        rows = self.model().mapFromSourceRows(self._selected_rows_arr())
         return [index(row, col)
                 for col in range(self.model().columnCount())
                 for row in rows]
@@ -1246,12 +1176,7 @@ def tool_tip(value):
 if __name__ == "__main__":  # pragma: no cover
     filename = "iris.tab"
     iris = Orange.data.Table(filename)
-    idom = iris.domain
-    dom = Domain(
-        idom.attributes,
-        DiscreteVariable(idom.class_var.name, idom.class_var.values[1::-1])
-    )
-    iris2 = iris[:100].transform(dom)
+    iris2 = iris[:100]
 
     def pred_error(data, *args, **kwargs):
         raise ValueError
@@ -1260,6 +1185,12 @@ def pred_error(data, *args, **kwargs):
     pred_error.name = "To err is human"
 
     if iris.domain.has_discrete_class:
+        idom = iris.domain
+        dom = Domain(
+            idom.attributes,
+            DiscreteVariable(idom.class_var.name, idom.class_var.values[1::-1])
+        )
+        iris2 = iris2.transform(dom)
         predictors_ = [
             Orange.classification.SVMLearner(probability=True)(iris2),
             Orange.classification.LogisticRegressionLearner()(iris),


Test patch:
diff --git a/Orange/widgets/evaluate/tests/test_owpredictions.py b/Orange/widgets/evaluate/tests/test_owpredictions.py
--- a/Orange/widgets/evaluate/tests/test_owpredictions.py
+++ b/Orange/widgets/evaluate/tests/test_owpredictions.py
@@ -6,15 +6,15 @@
 
 import numpy as np
 
-from AnyQt.QtCore import QItemSelectionModel, QItemSelection
-from AnyQt.QtGui import QStandardItemModel
+from AnyQt.QtCore import QItemSelectionModel, QItemSelection, Qt
 
 from Orange.base import Model
 from Orange.classification import LogisticRegressionLearner
 from Orange.data.io import TabReader
 from Orange.widgets.tests.base import WidgetTest
 from Orange.widgets.evaluate.owpredictions import (
-    OWPredictions, SortProxyModel, SharedSelectionModel, SharedSelectionStore)
+    OWPredictions, SharedSelectionModel, SharedSelectionStore, DataModel,
+    PredictionsModel)
 from Orange.widgets.evaluate.owcalibrationplot import OWCalibrationPlot
 from Orange.widgets.evaluate.owconfusionmatrix import OWConfusionMatrix
 from Orange.widgets.evaluate.owliftcurve import OWLiftCurve
@@ -202,11 +202,7 @@ def test_predictor_fails(self):
 
     def test_sort_matching(self):
         def get_items_order(model):
-            n = pred_model.rowCount()
-            return [
-                pred_model.mapToSource(model.index(i, 0)).row()
-                for i in range(n)
-            ]
+            return model.mapToSourceRows(np.arange(model.rowCount()))
 
         w = self.widget
 
@@ -216,46 +212,42 @@ def get_items_order(model):
         self.send_signal(self.widget.Inputs.data, titanic)
 
         pred_model = w.predictionsview.model()
-        data_model = w.predictionsview.model()
+        data_model = w.dataview.model()
         n = pred_model.rowCount()
 
         # no sort
         pred_order = get_items_order(pred_model)
         data_order = get_items_order(data_model)
-        self.assertListEqual(pred_order, list(range(n)))
-        self.assertListEqual(data_order, list(range(n)))
+        np.testing.assert_array_equal(pred_order, np.arange(n))
+        np.testing.assert_array_equal(data_order, np.arange(n))
 
         # sort by first column in prediction table
         pred_model.sort(0)
         w.predictionsview.horizontalHeader().sectionClicked.emit(0)
         pred_order = get_items_order(pred_model)
         data_order = get_items_order(data_model)
-        self.assertListEqual(pred_order, data_order)
+        np.testing.assert_array_equal(pred_order, data_order)
 
         # sort by second column in data table
         data_model.sort(1)
         w.dataview.horizontalHeader().sectionClicked.emit(0)
         pred_order = get_items_order(pred_model)
         data_order = get_items_order(data_model)
-        self.assertListEqual(pred_order, data_order)
+        np.testing.assert_array_equal(pred_order, data_order)
 
         # restore order
         w.reset_button.click()
         pred_order = get_items_order(pred_model)
         data_order = get_items_order(data_model)
-        self.assertListEqual(pred_order, list(range(n)))
-        self.assertListEqual(data_order, list(range(n)))
+        np.testing.assert_array_equal(pred_order, np.arange(n))
+        np.testing.assert_array_equal(data_order, np.arange(n))
 
     def test_sort_predictions(self):
         """
         Test whether sorting of probabilities by FilterSortProxy is correct.
         """
         def get_items_order(model):
-            n = pred_model.rowCount()
-            return [
-                pred_model.mapToSource(model.index(i, 0)).row()
-                for i in range(n)
-            ]
+            return model.mapToSourceRows(np.arange(model.rowCount()))
 
         log_reg_iris = LogisticRegressionLearner()(self.iris)
         self.send_signal(self.widget.Inputs.predictors, log_reg_iris)
@@ -528,21 +520,19 @@ def check_evres(expected):
 
 class SelectionModelTest(unittest.TestCase):
     def setUp(self):
-        self.sourceModel1 = QStandardItemModel(5, 2)
-        self.proxyModel1 = SortProxyModel()
-        self.proxyModel1.setSourceModel(self.sourceModel1)
-        self.store = SharedSelectionStore(self.proxyModel1)
-        self.model1 = SharedSelectionModel(self.store, self.proxyModel1, None)
+        iris = Table("iris")
 
-        self.sourceModel2 = QStandardItemModel(5, 3)
-        self.proxyModel2 = SortProxyModel()
-        self.proxyModel2.setSourceModel(self.sourceModel2)
-        self.model2 = SharedSelectionModel(self.store, self.proxyModel2, None)
+        self.datamodel1 = DataModel(iris[:5, :2])
+        self.store = SharedSelectionStore(self.datamodel1)
+        self.model1 = SharedSelectionModel(self.store, self.datamodel1, None)
+
+        self.datamodel2 = DataModel(iris[-5:, :3])
+        self.model2 = SharedSelectionModel(self.store, self.datamodel2, None)
 
     def itsel(self, rows):
         sel = QItemSelection()
         for row in rows:
-            index = self.store.proxy.index(row, 0)
+            index = self.store.model.index(row, 0)
             sel.select(index, index)
         return sel
 
@@ -560,33 +550,33 @@ def test_select_rows(self):
         store = self.store
         store.select_rows({1, 2}, QItemSelectionModel.Select)
         self.assertEqual(store.rows, {1, 2})
-        emit1.assert_called_with({1, 2}, set())
-        emit2.assert_called_with({1, 2}, set())
+        emit1.assert_called_with([1, 2], [])
+        emit2.assert_called_with([1, 2], [])
 
         store.select_rows({1, 2, 4}, QItemSelectionModel.Select)
         self.assertEqual(store.rows, {1, 2, 4})
-        emit1.assert_called_with({4}, set())
-        emit2.assert_called_with({4}, set())
+        emit1.assert_called_with([4], [])
+        emit2.assert_called_with([4], [])
 
         store.select_rows({3, 4}, QItemSelectionModel.Toggle)
         self.assertEqual(store.rows, {1, 2, 3})
-        emit1.assert_called_with({3}, {4})
-        emit2.assert_called_with({3}, {4})
+        emit1.assert_called_with([3], [4])
+        emit2.assert_called_with([3], [4])
 
         store.select_rows({0, 2}, QItemSelectionModel.Deselect)
         self.assertEqual(store.rows, {1, 3})
-        emit1.assert_called_with(set(), {2})
-        emit2.assert_called_with(set(), {2})
+        emit1.assert_called_with([], [2])
+        emit2.assert_called_with([], [2])
 
         store.select_rows({2, 3, 4}, QItemSelectionModel.ClearAndSelect)
         self.assertEqual(store.rows, {2, 3, 4})
-        emit1.assert_called_with({2, 4}, {1})
-        emit2.assert_called_with({2, 4}, {1})
+        emit1.assert_called_with([2, 4], [1])
+        emit2.assert_called_with([2, 4], [1])
 
         store.select_rows({2, 3, 4}, QItemSelectionModel.Clear)
         self.assertEqual(store.rows, set())
-        emit1.assert_called_with(set(), {2, 3, 4})
-        emit2.assert_called_with(set(), {2, 3, 4})
+        emit1.assert_called_with([], [2, 3, 4])
+        emit2.assert_called_with([], [2, 3, 4])
 
         store.select_rows({2, 3, 4}, QItemSelectionModel.ClearAndSelect)
         emit1.reset_mock()
@@ -600,13 +590,13 @@ def test_select_maps_from_proxy(self):
         store.select_rows = Mock()
 
         # Map QItemSelection
-        store.proxy.setSortIndices([1, 2, 3, 4, 0])
+        store.model.setSortIndices(np.array([4, 0, 1, 2, 3]))
         store.select(self.itsel([1, 2]), QItemSelectionModel.Select)
         store.select_rows.assert_called_with({0, 1}, QItemSelectionModel.Select)
 
         # Map QModelIndex
-        store.proxy.setSortIndices([1, 2, 3, 4, 0])
-        store.select(store.proxy.index(0, 0), QItemSelectionModel.Select)
+        store.model.setSortIndices(np.array([4, 0, 1, 2, 3]))
+        store.select(store.model.index(0, 0), QItemSelectionModel.Select)
         store.select_rows.assert_called_with({4}, QItemSelectionModel.Select)
 
         # Map empty selection
@@ -624,8 +614,8 @@ def test_clear(self):
 
         store.clear_selection()
         self.assertEqual(store.rows, set())
-        emit1.assert_called_with(set(), {1, 2, 4})
-        emit2.assert_called_with(set(), {1, 2, 4})
+        emit1.assert_called_with([], [1, 2, 4])
+        emit2.assert_called_with([], [1, 2, 4])
 
     def test_reset(self):
         store = self.store
@@ -644,20 +634,27 @@ def test_reset(self):
     def test_emit_changed_maps_to_proxy(self):
         store = self.store
         emit1 = self.model1.emit_selection_rows_changed = Mock()
+        self.model2.emit_selection_rows_changed = Mock()
+
+        def assert_called(exp_selected, exp_deselected):
+            # pylint: disable=unsubscriptable-object
+            selected, deselected = emit1.call_args[0]
+            self.assertEqual(list(selected), exp_selected)
+            self.assertEqual(list(deselected), exp_deselected)
+
 
-        store.proxy.setSortIndices([1, 2, 3, 4, 0])
+        store.model.setSortIndices([4, 0, 1, 2, 3])
         store.select_rows({3, 4}, QItemSelectionModel.Select)
-        emit1.assert_called_with({4, 0}, set())
+        assert_called([4, 0], [])
 
-        store.proxy.setSortIndices(None)
+        store.model.setSortIndices(None)
         store.select_rows({4}, QItemSelectionModel.Deselect)
-        emit1.assert_called_with(set(), {0})
+        assert_called([], [4])
 
-        store.proxy.setSortIndices([1, 0, 3, 4, 2])
-        store.proxy.setSortIndices(None)
-        self.proxyModel1.sort(-1)
+        store.model.setSortIndices([1, 0, 3, 4, 2])
+        store.model.setSortIndices(None)
         store.select_rows({2, 3}, QItemSelectionModel.Deselect)
-        emit1.assert_called_with(set(), {3})
+        assert_called([], [3])
 
 
 class SharedSelectionModelTest(SelectionModelTest):
@@ -671,9 +668,9 @@ def test_selection_from_rows(self):
         sel = self.model1.selection_from_rows({1, 2})
         self.assertEqual(len(sel), 2)
         ind1 = sel[0]
-        self.assertIs(ind1.model(), self.proxyModel1)
+        self.assertIs(ind1.model(), self.model1.model())
         self.assertEqual(ind1.left(), 0)
-        self.assertEqual(ind1.right(), self.proxyModel1.columnCount() - 1)
+        self.assertEqual(ind1.right(), self.model1.model().columnCount() - 1)
         self.assertEqual(ind1.top(), 1)
         self.assertEqual(ind1.bottom(), 1)
 
@@ -682,25 +679,26 @@ def test_emit_selection_rows_changed(self):
         m2 = Mock()
         self.model1.selectionChanged.connect(m1)
         self.model2.selectionChanged.connect(m2)
-        self.proxyModel1.setSortIndices([1, 0, 2, 4, 3])
+        self.model1.model().setSortIndices(np.array([1, 0, 2, 4, 3]))
         self.model1.select(self.itsel({1, 3}), QItemSelectionModel.Select)
         self.assertEqual(self.store.rows, {0, 4})
 
-        for model, m in zip((self.proxyModel1, self.proxyModel2), (m1, m2)):
+        for model, m in zip((self.model1, self.model2), (m1, m2)):
             sel = m.call_args[0][0]
             self.assertEqual(len(sel), 2)
             for ind, row in zip(sel, (1, 3)):
-                self.assertIs(ind.model(), model)
+                self.assertIs(ind.model(), model.model())
                 self.assertEqual(ind.left(), 0)
-                self.assertEqual(ind.right(), model.columnCount() - 1)
+                self.assertEqual(ind.right(), model.model().columnCount() - 1)
                 self.assertEqual(ind.top(), row)
                 self.assertEqual(ind.bottom(), row)
 
     def test_methods(self):
         def rowcol(sel):
             return {(index.row(), index.column()) for index in sel}
-        self.proxyModel1.setSortIndices([1, 0, 2, 4, 3])
-        self.proxyModel2.setSortIndices([1, 0, 2, 4, 3])
+
+        self.model1.model().setSortIndices(np.array([1, 0, 2, 4, 3]))
+        self.model2.model().setSortIndices(np.array([1, 0, 2, 4, 3]))
 
         self.assertFalse(self.model1.hasSelection())
         self.assertFalse(self.model1.isColumnSelected(1))
@@ -747,5 +745,139 @@ def rowcol(sel):
         self.assertEqual(self.model1.selectedIndexes(), [])
 
 
+class PredictionsModelTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.values = np.array([[0, 1, 1, 2, 0], [0, 0, 0, 1, 0]], dtype=float)
+        cls.probs = [np.array([[80, 10, 10],
+                               [30, 70, 0],
+                               [15, 80, 5],
+                               [0,  10, 90],
+                               [55, 40, 5]]) / 100,
+                     np.array([[80, 0, 20],
+                               [90, 5, 5],
+                               [70, 10, 20],
+                               [10, 60, 30],
+                               [50, 25, 25]]) / 100]
+        cls.no_probs = [np.zeros((5, 0)), np.zeros((5, 0))]
+
+    def test_model_classification(self):
+        model = PredictionsModel(self.values, self.probs)
+        self.assertEqual(model.rowCount(), 5)
+        self.assertEqual(model.columnCount(), 2)
+
+        val, prob = model.data(model.index(0, 1))
+        self.assertEqual(val, 0)
+        np.testing.assert_equal(prob, [0.8, 0, 0.2])
+
+        val, prob = model.data(model.index(3, 1))
+        self.assertEqual(val, 1)
+        np.testing.assert_equal(prob, [0.1, 0.6, 0.3])
+
+    def test_model_regression(self):
+        model = PredictionsModel(self.values, self.no_probs)
+        self.assertEqual(model.rowCount(), 5)
+        self.assertEqual(model.columnCount(), 2)
+
+        val, prob = model.data(model.index(0, 1))
+        self.assertEqual(val, 0)
+        np.testing.assert_equal(prob, [])
+
+        val, prob = model.data(model.index(3, 1))
+        self.assertEqual(val, 1)
+        np.testing.assert_equal(prob, [])
+
+    def test_model_header(self):
+        model = PredictionsModel(self.values, self.probs)
+        self.assertIsNone(model.headerData(0, Qt.Horizontal))
+        self.assertEqual(model.headerData(3, Qt.Vertical), "4")
+
+        model = PredictionsModel(self.values, self.probs, ["a", "b"])
+        self.assertEqual(model.headerData(0, Qt.Horizontal), "a")
+        self.assertEqual(model.headerData(1, Qt.Horizontal), "b")
+        self.assertEqual(model.headerData(3, Qt.Vertical), "4")
+
+        model = PredictionsModel(self.values, self.probs, ["a"])
+        self.assertEqual(model.headerData(0, Qt.Horizontal), "a")
+        self.assertIsNone(model.headerData(1, Qt.Horizontal))
+        self.assertEqual(model.headerData(3, Qt.Vertical), "4")
+
+    def test_model_empty(self):
+        model = PredictionsModel()
+        self.assertEqual(model.rowCount(), 0)
+        self.assertEqual(model.columnCount(), 0)
+        self.assertIsNone(model.headerData(1, Qt.Horizontal))
+
+    def test_sorting_classification(self):
+        model = PredictionsModel(self.values, self.probs)
+
+        val, prob = model.data(model.index(0, 1))
+        self.assertEqual(val, 0)
+        np.testing.assert_equal(prob, [0.8, 0, 0.2])
+
+        val, prob = model.data(model.index(3, 1))
+        self.assertEqual(val, 1)
+        np.testing.assert_equal(prob, [0.1, 0.6, 0.3])
+
+        model.setProbInd([2])
+        model.sort(0, Qt.DescendingOrder)
+        val, prob = model.data(model.index(0, 0))
+        self.assertEqual(val, 2)
+        np.testing.assert_equal(prob, [0, 0.1, 0.9])
+        val, prob = model.data(model.index(0, 1))
+        self.assertEqual(val, 1)
+        np.testing.assert_equal(prob, [0.1, 0.6, 0.3])
+
+        model.setProbInd([2])
+        model.sort(1, Qt.AscendingOrder)
+        val, prob = model.data(model.index(0, 1))
+        self.assertEqual(val, 0)
+        np.testing.assert_equal(prob, [0.9, 0.05, 0.05])
+        val, prob = model.data(model.index(0, 0))
+        self.assertEqual(val, 1)
+        np.testing.assert_equal(prob, [0.3, 0.7, 0])
+
+        model.setProbInd([1, 0])
+        model.sort(0, Qt.AscendingOrder)
+        np.testing.assert_equal(model.data(model.index(0, 0))[1], [0, .1, .9])
+        np.testing.assert_equal(model.data(model.index(1, 0))[1], [0.8, .1, .1])
+
+        model.setProbInd([1, 2])
+        model.sort(0, Qt.AscendingOrder)
+        np.testing.assert_equal(model.data(model.index(0, 0))[1], [0.8, .1, .1])
+        np.testing.assert_equal(model.data(model.index(1, 0))[1], [0, .1, .9])
+
+        model.setProbInd([])
+        model.sort(0, Qt.AscendingOrder)
+        self.assertEqual([model.data(model.index(i, 0))[0]
+                          for i in range(model.rowCount())], [0, 0, 1, 1, 2])
+
+        model.setProbInd([])
+        model.sort(0, Qt.DescendingOrder)
+        self.assertEqual([model.data(model.index(i, 0))[0]
+                          for i in range(model.rowCount())], [2, 1, 1, 0, 0])
+
+    def test_sorting_regression(self):
+        model = PredictionsModel(self.values, self.no_probs)
+
+        self.assertEqual(model.data(model.index(0, 1))[0], 0)
+        self.assertEqual(model.data(model.index(3, 1))[0], 1)
+
+        model.setProbInd([2])
+        model.sort(0, Qt.AscendingOrder)
+        self.assertEqual([model.data(model.index(i, 0))[0]
+                          for i in range(model.rowCount())], [0, 0, 1, 1, 2])
+
+        model.setProbInd([])
+        model.sort(0, Qt.DescendingOrder)
+        self.assertEqual([model.data(model.index(i, 0))[0]
+                          for i in range(model.rowCount())], [2, 1, 1, 0, 0])
+
+        model.setProbInd(None)
+        model.sort(0, Qt.AscendingOrder)
+        self.assertEqual([model.data(model.index(i, 0))[0]
+                          for i in range(model.rowCount())], [0, 0, 1, 1, 2])
+
+
 if __name__ == "__main__":
     unittest.main()