Basic matplotlib output for scatterplots

Author: markotoplak

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,35 @@ def write_image(cls, filename, scene):
         super().write_image(filename, scene)
 
 
+class MatplotlibFormat(FileFormat):
+    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 = ('.matplotlib.pdf',)  # file formats with same extension are not supported
+    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/utils/matplotlib_export.py

@@ -0,0 +1,108 @@
+from itertools import chain
+
+import numpy as np
+
+from pyqtgraph.graphicsItems.ScatterPlotItem import ScatterPlotItem
+
+
+def numpy_repr(a):
+    """ A numpy repr without summarization """
+    opts = np.get_printoptions()
+    try:
+        np.set_printoptions(threshold=10**10)
+        return repr(a)
+    finally:
+        np.set_printoptions(**opts)
+
+
+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")
+    code.append("sizes = {}".format(numpy_repr(sizes)))
+
+    def colortuple(color):
+        return color.redF(), color.greenF(), color.blueF(), color.alphaF()
+
+    edgecolors = np.array([colortuple(a.color()) for a in scatterplot_item.data["pen"]])
+    facecolors = np.array([colortuple(a.color()) for a in scatterplot_item.data["brush"]])
+
+    code.append("edgecolors = {}".format(numpy_repr(edgecolors)))
+    code.append("facecolors = {}".format(numpy_repr(facecolors)))
+
+    # 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 are unused
+            code.append("plt.scatter(x=x, y=y, s=sizes**2/4, marker={},".format(repr(m)))
+            code.append("            facecolors=facecolors, edgecolors=edgecolors)")
+        else:
+            code.append("indices = {}".format(numpy_repr(indices)))
+            code.append("plt.scatter(x=x[indices], y=y[indices], s=sizes[indices]**2/4, "
+                        "marker={},".format(repr(m)))
+            code.append("            facecolors=facecolors[indices], "
+                        "edgecolors=edgecolors[indices])")
+
+    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)
+
+

requirements-gui.txt

@@ -2,6 +2,7 @@
 AnyQt>=0.0.8
 
 pyqtgraph>=0.10.0
+matplotlib>=2.0.0
 
 # For add-ons' descriptions
 docutils