[ENH] Matplotlib output for Scatter plot

Orange/tests/test_evaluation_scoring.py

@@ -1,10 +1,7 @@
 # Test methods with long descriptive names can omit docstrings
 # pylint: disable=missing-docstring
 
-import sys
-import contextlib
 import unittest
-from unittest.mock import Mock
 import numpy as np
 
 from Orange.data import DiscreteVariable, ContinuousVariable, Domain
@@ -337,23 +334,9 @@ def test_compute_CD(self):
         cd = scoring.compute_CD(avranks, 30, test="bonferroni-dunn")
         np.testing.assert_almost_equal(cd, 0.798)
 
-        @contextlib.contextmanager
-        def mock_module(name):
-            if not name in sys.modules:
-                try:
-                    sys.modules[name] = Mock()
-                    yield
-                finally:
-                    del sys.modules[name]
-            else:
-                yield
-
         # Do what you will, just don't crash
-        with mock_module("matplotlib"), \
-                mock_module("matplotlib.pyplot"), \
-                mock_module("matplotlib.backends.backend_agg"):
-            scoring.graph_ranks(avranks, "abcd", cd)
-            scoring.graph_ranks(avranks, "abcd", cd, cdmethod=0)
+        scoring.graph_ranks(avranks, "abcd", cd)
+        scoring.graph_ranks(avranks, "abcd", cd, cdmethod=0)
 
 
 class TestLogLoss(unittest.TestCase):

Orange/widgets/io.py

@@ -9,6 +9,7 @@
 )
 
 from Orange.data.io import FileFormat
+from Orange.widgets.utils.matplotlib_export import scene_code
 
 # Importing WebviewWidget can fail if neither QWebKit (old, deprecated) nor
 # QWebEngine (bleeding-edge, hard to install) are available
@@ -181,6 +182,36 @@ def write_image(cls, filename, scene):
         super().write_image(filename, scene)
 
 
+class MatplotlibFormat:
+    # not registered as a FileFormat as it only works with scatter plot
+    EXTENSIONS = ('.py',)
+    DESCRIPTION = 'Python Code (with Matplotlib)'
+    PRIORITY = 300
+
+    @classmethod
+    def write_image(cls, filename, scene):
+        code = scene_code(scene) + "\n\nplt.show()"
+        with open(filename, "wt") as f:
+            f.write(code)
+
+    @classmethod
+    def write(cls, filename, scene):
+        if type(scene) == dict:
+            scene = scene['scene']
+        cls.write_image(filename, scene)
+
+
+class MatplotlibPDFFormat(MatplotlibFormat):
+    EXTENSIONS = ('.pdf',)
+    DESCRIPTION = 'Portable Document Format (from Matplotlib)'
+    PRIORITY = 200
+
+    @classmethod
+    def write_image(cls, filename, scene):
+        code = scene_code(scene) + "\n\nplt.savefig({})".format(repr(filename))
+        exec(code, {})  # will generate a pdf
+
+
 if hasattr(QtGui, "QPdfWriter"):
     class PdfFormat(ImgFormat):
         EXTENSIONS = ('.pdf', )

Orange/widgets/tests/test_io.py

@@ -1,11 +1,18 @@
 import os
 import tempfile
 import unittest
+from unittest.mock import patch
 
 from AnyQt.QtWidgets import QGraphicsScene, QGraphicsRectItem
-from Orange.widgets.tests.base import GuiTest
+
+import Orange
+from Orange.tests import named_file
+from Orange.widgets.tests.base import GuiTest, WidgetTest
 
 from Orange.widgets import io as imgio
+from Orange.widgets.io import MatplotlibFormat, MatplotlibPDFFormat
+from Orange.widgets.visualize.owscatterplot import OWScatterPlot
+
 
 @unittest.skipUnless(hasattr(imgio, "PdfFormat"), "QPdfWriter not available")
 class TestIO(GuiTest):
@@ -18,3 +25,33 @@ def test_pdf(self):
             imgio.PdfFormat.write_image(fname, sc)
         finally:
             os.unlink(fname)
+
+
+class TestMatplotlib(WidgetTest):
+
+    def test_python(self):
+        iris = Orange.data.Table("iris")
+        self.widget = self.create_widget(OWScatterPlot)
+        self.send_signal(OWScatterPlot.Inputs.data, iris[::10])
+        with named_file("", suffix=".py") as fname:
+            with patch("Orange.widgets.utils.filedialogs.open_filename_dialog_save",
+                       lambda *x: (fname, MatplotlibFormat, None)):
+                self.widget.save_graph()
+                with open(fname, "rt") as f:
+                    code = f.read()
+                self.assertIn("plt.show()", code)
+                self.assertIn("plt.scatter", code)
+                # test if the runs
+                exec(code.replace("plt.show()", ""), {})
+
+    def test_pdf(self):
+        iris = Orange.data.Table("iris")
+        self.widget = self.create_widget(OWScatterPlot)
+        self.send_signal(OWScatterPlot.Inputs.data, iris[::10])
+        with named_file("", suffix=".pdf") as fname:
+            with patch("Orange.widgets.utils.filedialogs.open_filename_dialog_save",
+                       lambda *x: (fname, MatplotlibPDFFormat, None)):
+                self.widget.save_graph()
+                with open(fname, "rb") as f:
+                    code = f.read()
+                self.assertTrue(code.startswith(b"%PDF"))

Orange/widgets/tests/test_matplotlib_export.py

@@ -0,0 +1,43 @@
+import pyqtgraph as pg
+
+from AnyQt.QtCore import QRectF
+
+import Orange
+
+from Orange.widgets.tests.base import WidgetTest
+from Orange.widgets.utils.matplotlib_export import scatterplot_code
+from Orange.widgets.visualize.owscatterplot import OWScatterPlot
+
+
+def add_intro(a):
+    r = "import matplotlib.pyplot as plt\n" + \
+        "from numpy import array\n" + \
+        "plt.clf()"
+    return r + a
+
+
+class TestScatterPlot(WidgetTest):
+
+    def test_owscatterplot_ignore_empty(self):
+        iris = Orange.data.Table("iris")
+        self.widget = self.create_widget(OWScatterPlot)
+        self.send_signal(OWScatterPlot.Inputs.data, iris[::10])
+        code = scatterplot_code(self.widget.graph.scatterplot_item)
+        self.assertIn("plt.scatter", code)
+        # tbe selected graph has to generate nothing
+        code = scatterplot_code(self.widget.graph.scatterplot_item_sel)
+        self.assertEqual(code, "")
+        # for a selection the selected graph has to be non-empty
+        self.widget.graph.select_by_rectangle(QRectF(4, 3, 3, 1))
+        code = scatterplot_code(self.widget.graph.scatterplot_item_sel)
+        self.assertIn("plt.scatter", code)
+        exec(add_intro(code), {})
+
+    def test_scatterplot_simple(self):
+        plotWidget = pg.PlotWidget(background="w")
+        scatterplot = pg.ScatterPlotItem()
+        scatterplot.setData(x=[1, 2, 3], y=[3, 2, 1])
+        plotWidget.addItem(scatterplot)
+        code = scatterplot_code(scatterplot)
+        self.assertIn("plt.scatter", code)
+        exec(add_intro(code), {})

Orange/widgets/utils/matplotlib_export.py

@@ -0,0 +1,215 @@
+from itertools import chain
+
+import numpy as np
+
+from matplotlib.colors import to_hex
+from pyqtgraph.graphicsItems.ScatterPlotItem import ScatterPlotItem
+from AnyQt.QtCore import Qt
+from AnyQt.QtGui import QPen, QBrush
+
+
+def numpy_repr(a):
+    """ A numpy repr without summarization """
+    opts = np.get_printoptions()
+    # avoid numpy repr as it changes between versions
+    # TODO handle numpy repr differences
+    if isinstance(a, np.ndarray):
+        return "array(" + repr(list(a)) + ")"
+    try:
+        np.set_printoptions(threshold=10**10)
+        return repr(a)
+    finally:
+        np.set_printoptions(**opts)
+
+
+def numpy_repr_int(a):
+    # avoid numpy repr as it changes between versions
+    # TODO handle numpy repr differences
+    return "array(" + repr(list(a)) + ", dtype='int')"
+
+
+def compress_if_all_same(l):
+    s = set(l)
+    return s.pop() if len(s) == 1 else l
+
+
+def is_sequence_not_string(a):
+    if isinstance(a, str):
+        return False
+    try:
+        iter(a)
+        return True
+    except TypeError:
+        pass
+    return False
+
+
+def code_with_indices(data, data_name, indices, indices_name):
+    if is_sequence_not_string(data) and indices is not None:
+        return data_name + "[" + indices_name + "]"
+    else:
+        return data_name
+
+
+def index_per_different(l):
+    different = []
+    different_ind = {}
+    index = []
+    for e in l:
+        if e not in different_ind:
+            different_ind[e] = len(different)
+            different.append(e)
+        index.append(different_ind[e])
+    return different, index
+
+
+def scatterplot_code(scatterplot_item):
+    x = scatterplot_item.data['x']
+    y = scatterplot_item.data['y']
+    sizes = scatterplot_item.data["size"]
+
+    code = []
+
+    code.append("# data")
+    code.append("x = {}".format(numpy_repr(x)))
+    code.append("y = {}".format(numpy_repr(y)))
+
+    code.append("# style")
+    sizes = compress_if_all_same(sizes)
+    if sizes == -1:
+        sizes = None
+    code.append("sizes = {}".format(numpy_repr(sizes)))
+
+    def colortuple(pen):
+        if isinstance(pen, (QPen, QBrush)):
+            color = pen.color()
+            return color.redF(), color.greenF(), color.blueF(), color.alphaF()
+        return pen
+
+    def width(pen):
+        if isinstance(pen, QPen):
+            return pen.widthF()
+        return pen
+
+    linewidths = np.array([width(a) for a in scatterplot_item.data["pen"]])
+
+    def shown(a):
+        if isinstance(a, (QPen, QBrush)):
+            s = a.style()
+            if s == Qt.NoPen or s == Qt.NoBrush or a.color().alpha() == 0:
+                return False
+        return True
+
+    shown_edge = [shown(a) for a in scatterplot_item.data["pen"]]
+    shown_brush = [shown(a) for a in scatterplot_item.data["brush"]]
+
+    # return early if the scatterplot is all transparent
+    if not any(shown_edge) and not any(shown_brush):
+        return ""
+
+    def do_colors(code, data_column, show, name):
+        colors = [colortuple(a) for a in data_column]
+        if all(a is None for a in colors):
+            colors, index = None, None
+        else:
+            # replace None values with blue colors
+            colors = np.array([((0, 0, 1, 1) if a is None else a)
+                               for a in colors])
+            # set alpha for hidden (Qt.NoPen, Qt.NoBrush) elements to zero
+            colors[:, 3][np.array(show) == 0] = 0
+            # shorter color names for printout
+            colors = [to_hex(c, keep_alpha=True) for c in colors]
+            colors, index = index_per_different(colors)
+
+        code.append("{} = {}".format(name, repr(colors)))
+        if index is not None:
+            code.append("{}_index = {}".format(name, numpy_repr_int(index)))
+
+        decompresssed_code = name
+        if index is not None:
+            decompresssed_code = "array({})[{}_index]".format(name, name)
+            colors = np.array(colors)[index]
+
+        return colors, decompresssed_code
+
+    edgecolors, edgecolors_code = do_colors(code, scatterplot_item.data["pen"], shown_edge, "edgecolors")
+    facecolors, facecolors_code = do_colors(code, scatterplot_item.data["brush"], shown_brush, "facecolors")
+
+    linewidths = compress_if_all_same(linewidths)
+    code.append("linewidths = {}".format(numpy_repr(linewidths)))
+
+    # possible_markers for scatterplot are in .graph.CurveSymbols
+    def matplotlib_marker(m):
+        if m == "t":
+            return "^"
+        elif m == "t2":
+            return ">"
+        elif m == "t3":
+            return "<"
+        elif m == "star":
+            return "*"
+        elif m == "+":
+            return "P"
+        elif m == "x":
+            return "X"
+        return m
+
+    # TODO labels are missing
+
+    # each marker requires one call to matplotlib's scatter!
+    markers = np.array([matplotlib_marker(m) for m in scatterplot_item.data["symbol"]])
+    for m in set(markers):
+        indices = np.where(markers == m)[0]
+        if np.all(indices == np.arange(x.shape[0])):
+            indices = None
+        if indices is not None:
+            code.append("indices = {}".format(numpy_repr_int(indices)))
+
+        def indexed(data, data_name, indices=indices):
+            return code_with_indices(data, data_name, indices, "indices")
+
+        code.append("plt.scatter(x={}, y={}, s={}, marker={},\n"
+                    "            facecolors={}, edgecolors={},\n"
+                    "            linewidths={})"
+                    .format(indexed(x, "x"),
+                            indexed(y, "y"),
+                            (indexed(sizes, "sizes") + "**2/4") if sizes is not None else "sizes",
+                            repr(m),
+                            indexed(facecolors, facecolors_code),
+                            indexed(edgecolors, edgecolors_code),
+                            indexed(linewidths, "linewidths")
+                           ))
+
+    return "\n".join(code)
+
+
+def scene_code(scene):
+
+    code = []
+
+    code.append("import matplotlib.pyplot as plt")
+    code.append("from numpy import array")
+
+    code.append("")
+    code.append("plt.clf()")
+
+    code.append("")
+
+    for item in scene.items:
+        if isinstance(item, ScatterPlotItem):
+            code.append(scatterplot_code(item))
+
+    # TODO currently does not work for graphs without axes and for multiple axes!
+    for position, set_ticks, set_label in [("bottom", "plt.xticks", "plt.xlabel"),
+                                           ("left", "plt.yticks", "plt.ylabel")]:
+        axis = scene.getAxis(position)
+        code.append("{}({})".format(set_label, repr(str(axis.labelText))))
+
+        # textual tick labels
+        if axis._tickLevels is not None:
+            major_minor = list(chain(*axis._tickLevels))
+            locs = [a[0] for a in major_minor]
+            labels = [a[1] for a in major_minor]
+            code.append("{}({}, {})".format(set_ticks, locs, repr(labels)))
+
+    return "\n".join(code)