Scatterplot: Add tests for regression lines

Orange/widgets/visualize/tests/test_owscatterplot.py

@@ -1,19 +1,21 @@
 # Test methods with long descriptive names can omit docstrings
 # pylint: disable=missing-docstring,too-many-public-methods,protected-access
+# pylint: disable=too-many-lines
 from unittest.mock import MagicMock, patch, Mock
 import numpy as np
 
 from AnyQt.QtCore import QRectF, Qt
 from AnyQt.QtWidgets import QToolTip
+from AnyQt.QtGui import QColor
 
 from Orange.data import Table, Domain, ContinuousVariable, DiscreteVariable
 from Orange.widgets.tests.base import (
     WidgetTest, WidgetOutputsTestMixin, datasets, ProjectionWidgetTestMixin
 )
 from Orange.widgets.tests.utils import simulate
+from Orange.widgets.utils.colorpalette import DefaultRGBColors
 from Orange.widgets.visualize.owscatterplot import (
-    OWScatterPlot, ScatterPlotVizRank
-)
+    OWScatterPlot, ScatterPlotVizRank, OWScatterPlotGraph)
 from Orange.widgets.visualize.utils.widget import MAX_CATEGORIES
 from Orange.widgets.widget import AttributeList
 
@@ -735,12 +737,281 @@ def test_on_manual_change(self):
         selection = vizrank.rank_table.selectedIndexes()
         self.assertEqual(len(selection), 0)
 
-    def test_regression_line(self):
+    def test_regression_lines_appear(self):
         self.widget.graph.controls.show_reg_line.setChecked(True)
+        self.assertEqual(len(self.widget.graph.reg_line_items), 0)
         self.send_signal(self.widget.Inputs.data, self.data)
+        self.assertEqual(len(self.widget.graph.reg_line_items), 4)
+        simulate.combobox_activate_index(self.widget.controls.attr_color, 0)
+        self.assertEqual(len(self.widget.graph.reg_line_items), 1)
         data = self.data.copy()
         data[:, 0] = np.nan
         self.send_signal(self.widget.Inputs.data, data)
+        self.assertEqual(len(self.widget.graph.reg_line_items), 0)
+
+    def test_regression_line_coeffs(self):
+        widget = self.widget
+        graph = widget.graph
+        xy = np.array([[0, 0], [1, 0], [1, 2], [2, 2],
+                       [0, 1], [1, 3], [2, 5]], dtype=np.float)
+        colors = np.array([0, 0, 0, 0, 1, 1, 1], dtype=np.float)
+        widget.get_coordinates_data = lambda: xy.T
+        widget.get_color_data = lambda: colors
+        widget.is_continuous_color = lambda: False
+        graph.palette = DefaultRGBColors
+        graph.controls.show_reg_line.setChecked(True)
+
+        graph.update_regression_line()
+
+        line1 = graph.reg_line_items[1]
+        self.assertEqual(line1.pos().x(), 0)
+        self.assertEqual(line1.pos().y(), 0)
+        self.assertEqual(line1.angle, 45)
+        self.assertEqual(line1.pen.color().getRgb()[:3], graph.palette[0])
+
+        line2 = graph.reg_line_items[2]
+        self.assertEqual(line2.pos().x(), 0)
+        self.assertEqual(line2.pos().y(), 1)
+        self.assertAlmostEqual(line2.angle, np.degrees(np.arctan2(2, 1)))
+        self.assertEqual(line2.pen.color().getRgb()[:3], graph.palette[1])
+
+        graph.orthonormal_regression = True
+        graph.update_regression_line()
+
+        line1 = graph.reg_line_items[1]
+        self.assertEqual(line1.pos().x(), 0)
+        self.assertAlmostEqual(line1.pos().y(), -0.6180339887498949)
+        self.assertEqual(line1.angle, 58.28252558853899)
+        self.assertEqual(line1.pen.color().getRgb()[:3], graph.palette[0])
+
+        line2 = graph.reg_line_items[2]
+        self.assertEqual(line2.pos().x(), 0)
+        self.assertEqual(line2.pos().y(), 1)
+        self.assertAlmostEqual(line2.angle, np.degrees(np.arctan2(2, 1)))
+        self.assertEqual(line2.pen.color().getRgb()[:3], graph.palette[1])
+
+    def test_orthonormal_line(self):
+        color = QColor(1, 2, 3)
+        width = 42
+        # Normal line
+        line = OWScatterPlotGraph._orthonormal_line(
+            np.array([0, 1, 1, 2]), np.array([0, 0, 2, 2]), color, width)
+        self.assertEqual(line.pos().x(), 0)
+        self.assertAlmostEqual(line.pos().y(), -0.6180339887498949)
+        self.assertEqual(line.angle, 58.28252558853899)
+        self.assertEqual(line.pen.color(), color)
+        self.assertEqual(line.pen.width(), width)
+
+        # Normal line, negative slope
+        line = OWScatterPlotGraph._orthonormal_line(
+            np.array([1, 2, 3]), np.array([3, 2, 1]), color, width)
+        self.assertEqual(line.pos().x(), 1)
+        self.assertEqual(line.pos().y(), 3)
+        self.assertEqual(line.angle % 360, 315)
+
+        # Horizontal line
+        line = OWScatterPlotGraph._orthonormal_line(
+            np.array([10, 11, 12]), np.array([42, 42, 42]), color, width)
+        self.assertEqual(line.pos().x(), 10)
+        self.assertEqual(line.pos().y(), 42)
+        self.assertEqual(line.angle, 0)
+
+        # Vertical line
+        line = OWScatterPlotGraph._orthonormal_line(
+            np.array([42, 42, 42]), np.array([10, 11, 12]), color, width)
+        self.assertEqual(line.pos().x(), 42)
+        self.assertEqual(line.pos().y(), 10)
+        self.assertEqual(line.angle, 90)
+
+        # No line because all points coincide
+        line = OWScatterPlotGraph._orthonormal_line(
+            np.array([1, 1, 1]), np.array([42, 42, 42]), color, width)
+        self.assertIsNone(line)
+
+        # No line because the group is symmetric
+        line = OWScatterPlotGraph._orthonormal_line(
+            np.array([1, 1, 2, 2]), np.array([42, 5, 5, 42]), color, width)
+        self.assertIsNone(line)
+
+    def test_regression_line(self):
+        color = QColor(1, 2, 3)
+        width = 42
+        # Normal line
+        line = OWScatterPlotGraph._regression_line(
+            np.array([0, 1, 1, 2]), np.array([0, 0, 2, 2]), color, width)
+        self.assertEqual(line.pos().x(), 0)
+        self.assertAlmostEqual(line.pos().y(), 0)
+        self.assertEqual(line.angle, 45)
+        self.assertEqual(line.pen.color(), color)
+        self.assertEqual(line.pen.width(), width)
+
+        # Normal line, negative slope
+        line = OWScatterPlotGraph._regression_line(
+            np.array([1, 2, 3]), np.array([3, 2, 1]), color, width)
+        self.assertEqual(line.pos().x(), 1)
+        self.assertEqual(line.pos().y(), 3)
+        self.assertEqual(line.angle % 360, 315)
+
+        # Horizontal line
+        line = OWScatterPlotGraph._regression_line(
+            np.array([10, 11, 12]), np.array([42, 42, 42]), color, width)
+        self.assertEqual(line.pos().x(), 10)
+        self.assertEqual(line.pos().y(), 42)
+        self.assertEqual(line.angle, 0)
+
+        # Vertical line
+        line = OWScatterPlotGraph._regression_line(
+            np.array([42, 42, 42]), np.array([10, 11, 12]), color, width)
+        self.assertIsNone(line)
+
+        # No line because all points coincide
+        line = OWScatterPlotGraph._regression_line(
+            np.array([1, 1, 1]), np.array([42, 42, 42]), color, width)
+        self.assertIsNone(line)
+
+    def test_add_line_calls_proper_regressor(self):
+        graph = self.widget.graph
+        graph._orthonormal_line = Mock(return_value=None)
+        graph._regression_line = Mock(return_value=None)
+        x, y, c, w = Mock(), Mock(), Mock(), Mock()
+
+        graph.orthonormal_regression = True
+        graph._add_line(x, y, c, w)
+        graph._orthonormal_line.assert_called_once_with(x, y, c, w)
+        graph._regression_line.assert_not_called()
+        graph._orthonormal_line.reset_mock()
+
+        graph.orthonormal_regression = False
+        graph._add_line(x, y, c, w)
+        graph._regression_line.assert_called_with(x, y, c, w)
+        graph._orthonormal_line.assert_not_called()
+
+    def test_no_regression_line(self):
+        graph = self.widget.graph
+        graph._orthonormal_line = lambda *_: None
+        graph.orthonormal_regression = True
+
+        graph.plot_widget.addItem = Mock()
+
+        x, y, c, w = Mock(), Mock(), Mock(), Mock()
+        graph._add_line(x, y, c, w)
+        graph.plot_widget.addItem.assert_not_called()
+        self.assertEqual(graph.reg_line_items, [])
+
+    def test_update_regression_line_calls_add_line(self):
+        widget = self.widget
+        graph = widget.graph
+        x, y = np.array([[0, 0], [1, 0], [1, 2], [2, 2],
+                         [0, 1], [1, 3], [2, 5]], dtype=np.float).T
+        colors = np.array([0, 0, 0, 0, 1, 1, 1], dtype=np.float)
+        widget.get_coordinates_data = lambda: (x, y)
+        widget.get_color_data = lambda: colors
+        widget.is_continuous_color = lambda: False
+        graph.palette = DefaultRGBColors
+        graph.controls.show_reg_line.setChecked(True)
+
+        graph._add_line = Mock()
+
+        graph.update_regression_line()
+        (args1, kwargs1), (args2, kwargs2), (args3, kwargs3) = \
+            graph._add_line.call_args_list
+        np.testing.assert_equal(args1[0], x)
+        np.testing.assert_equal(args1[1], y)
+        self.assertEqual(args1[2], QColor("#505050"))
+        self.assertEqual(kwargs1["width"], 1)
+
+        np.testing.assert_equal(args2[0], x[:4])
+        np.testing.assert_equal(args2[1], y[:4])
+        self.assertEqual(args2[2], graph.palette[0])
+        self.assertEqual(kwargs2["width"], 3)
+
+        np.testing.assert_equal(args3[0], x[4:])
+        np.testing.assert_equal(args3[1], y[4:])
+        self.assertEqual(args3[2], graph.palette[1])
+        self.assertEqual(kwargs3["width"], 3)
+        graph._add_line.reset_mock()
+
+        # Continuous color - just a single line
+        widget.is_continuous_color = lambda: True
+        graph.update_regression_line()
+        graph._add_line.assert_called_once()
+        args1, kwargs1 = graph._add_line.call_args_list[0]
+        np.testing.assert_equal(args1[0], x)
+        np.testing.assert_equal(args1[1], y)
+        self.assertEqual(args1[2], QColor("#505050"))
+        self.assertEqual(kwargs1["width"], 1)
+        graph._add_line.reset_mock()
+        widget.is_continuous_color = lambda: False
+
+        # No palette - just a single line
+        graph.palette = None
+        graph.update_regression_line()
+        graph._add_line.assert_called_once()
+        graph._add_line.reset_mock()
+        graph.palette = DefaultRGBColors
+
+        # Regression line is disabled
+        graph.show_reg_line = False
+        graph.update_regression_line()
+        graph._add_line.assert_not_called()
+        graph.show_reg_line = True
+
+        # No colors - just one line
+        widget.get_color_data = lambda: None
+        graph.update_regression_line()
+        graph._add_line.assert_called_once()
+        graph._add_line.reset_mock()
+
+        # No data
+        widget.get_coordinates_data = lambda: (None, None)
+        graph.update_regression_line()
+        graph._add_line.assert_not_called()
+        graph.show_reg_line = True
+        widget.get_coordinates_data = lambda: (x, y)
+
+        # One color group contains just one point - skip that line
+        widget.get_color_data = lambda: np.array([0] + [1] * (len(x) - 1))
+
+        graph.update_regression_line()
+        (args1, kwargs1), (args2, kwargs2) = graph._add_line.call_args_list
+        np.testing.assert_equal(args1[0], x)
+        np.testing.assert_equal(args1[1], y)
+        self.assertEqual(args1[2], QColor("#505050"))
+        self.assertEqual(kwargs1["width"], 1)
+
+        np.testing.assert_equal(args2[0], x[1:])
+        np.testing.assert_equal(args2[1], y[1:])
+        self.assertEqual(args2[2], graph.palette[1])
+        self.assertEqual(kwargs2["width"], 3)
+
+    def test_update_regression_line_is_called(self):
+        widget = self.widget
+        graph = widget.graph
+        urline = graph.update_regression_line = Mock()
+
+        self.send_signal(widget.Inputs.data, self.data)
+        urline.assert_called_once()
+        urline.reset_mock()
+
+        self.send_signal(widget.Inputs.data, None)
+        urline.assert_called_once()
+        urline.reset_mock()
+
+        self.send_signal(widget.Inputs.data, self.data)
+        urline.assert_called_once()
+        urline.reset_mock()
+
+        simulate.combobox_activate_index(self.widget.controls.attr_color, 0)
+        urline.assert_called_once()
+        urline.reset_mock()
+
+        simulate.combobox_activate_index(self.widget.controls.attr_color, 2)
+        urline.assert_called_once()
+        urline.reset_mock()
+
+        simulate.combobox_activate_index(self.widget.controls.attr_x, 3)
+        urline.assert_called_once()
+        urline.reset_mock()
 
 
 if __name__ == "__main__":