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(a.tolist()) + ")" 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 if isinstance(a, np.ndarray): a = a.tolist() else: a = list(a) return "array(" + repr(a) + ", dtype='int')" def compress_if_all_same(l): s = set(l) if len(s) == 1: v = s.pop() return v.item() if isinstance(v, np.generic) else v else: return 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("import numpy as np") 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)