From 670706158a8658cd1e3134a94adb7716b8741b8f Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Sun, 27 Sep 2015 19:59:52 -0400 Subject: [PATCH] ENH: Don't write files when rendering TermGraphs. Overhauls `zipline.modelling.visualize` to use in-memory buffers when shelling out to `dot` and friends. Also adds `svg`, `png`, and `jpeg` attributes to `TermGraph`, and adds a `_repr_png_` so that `TermGraph` renders as a PNG by default. --- zipline/modelling/graph.py | 23 +++++ zipline/modelling/visualize.py | 155 +++++++++++++++++---------------- 2 files changed, 102 insertions(+), 76 deletions(-) diff --git a/zipline/modelling/graph.py b/zipline/modelling/graph.py index 8abc8aa7..583a2181 100644 --- a/zipline/modelling/graph.py +++ b/zipline/modelling/graph.py @@ -7,6 +7,7 @@ from networkx import ( ) from six import itervalues, iteritems from zipline.utils.memoize import lazyval +from zipline.modelling.visualize import display_graph class CyclicDependency(Exception): @@ -41,6 +42,8 @@ class TermGraph(DiGraph): """ def __init__(self, terms): super(TermGraph, self).__init__(self) + + self._frozen = False parents = set() for term in itervalues(terms): self._add_to_graph(term, parents, extra_rows=0) @@ -50,6 +53,9 @@ class TermGraph(DiGraph): self._outputs = terms self._ordered = topological_sort(self) + # Mark that no more terms should be added to the graph. + self._frozen = True + @lazyval def offset(self): """ @@ -166,6 +172,8 @@ class TermGraph(DiGraph): """ Add `term` and all its inputs to the graph. """ + if self._frozen: + raise ValueError("Can't mutate `TermGraph` after construction.") # If we've seen this node already as a parent of the current traversal, # it means we have an unsatisifiable dependency. This should only be # possible if the term's inputs are mutated after construction. @@ -205,3 +213,18 @@ class TermGraph(DiGraph): """ attrs = self.node[term] attrs['extra_rows'] = max(N, attrs.get('extra_rows', 0)) + + @lazyval + def jpeg(self): + return display_graph(self, 'jpeg') + + @lazyval + def png(self): + return display_graph(self, 'png') + + @lazyval + def svg(self): + return display_graph(self, 'svg') + + def _repr_png_(self): + return self.png.data diff --git a/zipline/modelling/visualize.py b/zipline/modelling/visualize.py index d5c9e079..c59bf43f 100644 --- a/zipline/modelling/visualize.py +++ b/zipline/modelling/visualize.py @@ -1,19 +1,25 @@ """ Tools for visualizing dependencies between Terms. """ -from functools import partial +from __future__ import unicode_literals + from contextlib import contextmanager -from logbook import Logger, StderrHandler +import errno +from functools import partial +from io import BytesIO +from subprocess import Popen, PIPE + from networkx import topological_sort from six import iteritems -import subprocess + from zipline.data.dataset import BoundColumn from zipline.modelling import Filter, Factor, Classifier, Term from zipline.modelling.term import AssetExists -from zipline.modelling.graph import TermGraph -logger = Logger('Visualize') + +class NoIPython(Exception): + pass def delimit(delimiters, content): @@ -38,13 +44,13 @@ bracket = partial(delimit, '[]') def begin_graph(f, name, **attrs): writeln(f, "strict digraph %s {" % name) - writeln(f, "graph {attrs}".format(attrs=format_attrs(attrs))) + writeln(f, "graph {}".format(format_attrs(attrs))) def begin_cluster(f, name, **attrs): attrs.setdefault("label", quote(name)) writeln(f, "subgraph cluster_%s {" % name) - writeln(f, "graph {attrs}".format(attrs=format_attrs(attrs))) + writeln(f, "graph {}".format(format_attrs(attrs))) def end_graph(f): @@ -70,75 +76,96 @@ def roots(g): return set(n for n, d in iteritems(g.in_degree()) if d == 0) -def write_graph(g, filename, formats=('svg',), include_asset_exists=False): +def _render(g, out, format_, include_asset_exists=False): """ - Write the dependency graph of `terms` as a dot graph. + Draw `g` as a graph to `out`, in format `format`. - If `png` (default True), write a .png file using the system `dot` program. - If `pdf` (default False), write a .pdf file using the system `dot` program. + Parameters + ---------- + g : zipline.modelling.graph.TermGraph + Graph to render. + out : file-like object + format_ : str {'png', 'svg'} + Output format. + include_asset_exists : bool + Whether to filter out `AssetExists()` nodes. """ - dotfile = filename + '.dot' - graph_attrs = {'rankdir': 'TB', 'splines': 'ortho'} cluster_attrs = {'style': 'filled', 'color': 'lightgoldenrod1'} in_nodes = list(node for node in g if node.atomic) out_nodes = list(g.outputs.values()) - with open(dotfile, 'w') as f: - with graph(f, "G", **graph_attrs): - # Write outputs cluster. - with cluster(f, 'Output', labelloc='b', **cluster_attrs): - for term in out_nodes: - add_term_node(f, term) + f = BytesIO() + with graph(f, "G", **graph_attrs): - # Write inputs cluster. - with cluster(f, 'Input', **cluster_attrs): - for term in in_nodes: - if term is AssetExists() and not include_asset_exists: - continue - add_term_node(f, term) + # Write outputs cluster. + with cluster(f, 'Output', labelloc='b', **cluster_attrs): + for term in out_nodes: + add_term_node(f, term) - # Write intermediate results. - for term in topological_sort(g): - if term in in_nodes or term in out_nodes: + # Write inputs cluster. + with cluster(f, 'Input', **cluster_attrs): + for term in in_nodes: + if term is AssetExists() and not include_asset_exists: continue add_term_node(f, term) - # Write edges - for source, dest in g.edges(): - if source is AssetExists() and not include_asset_exists: - continue - add_edge(f, id(source), id(dest)) + # Write intermediate results. + for term in topological_sort(g): + if term in in_nodes or term in out_nodes: + continue + add_term_node(f, term) - outs = [] - for format_ in formats: - out = '.'.join([filename, format_]) - logger.info('Writing "%s"' % out) - subprocess.call(['dot', '-T', format_, dotfile, '-o', out]) - outs.append(out) - return outs + # Write edges + for source, dest in g.edges(): + if source is AssetExists() and not include_asset_exists: + continue + add_edge(f, id(source), id(dest)) + + cmd = ['dot', '-T', format_] + try: + proc = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE) + except OSError as e: + if e.errno == errno.ENOENT: + raise RuntimeError( + "Couldn't find `dot` graph layout program. " + "Make sure Graphviz is installed and `dot` is on your path." + ) + else: + raise + + f.seek(0) + proc_stdout, proc_stderr = proc.communicate(f.read()) + if proc_stderr: + raise RuntimeError( + "Error(s) while rendering graph: %s" % proc_stderr.decode('utf-8') + ) + + out.write(proc_stdout) -def show_graph(g, include_asset_exists=False): +def display_graph(g, format='svg', include_asset_exists=False): """ - Display a TermGraph interactively with IPython + Display a TermGraph interactively from within IPython. """ try: - from IPython.display import SVG + import IPython.display as display except ImportError: - raise Exception("IPython is not installed. Can't show term graph.") - result = write_graph( - g, - 'temp', - ('svg',), - include_asset_exists=include_asset_exists, - )[0] - return SVG(filename=result) + raise NoIPython("IPython is not installed. Can't display graph.") + + if format == 'svg': + display_cls = display.SVG + elif format in ("jpeg", "png"): + display_cls = partial(display.Image, format=format, embed=True) + + out = BytesIO() + _render(g, out, format, include_asset_exists=include_asset_exists) + return display_cls(data=out.getvalue()) def writeln(f, s): - f.write(s + '\n') + f.write((s + '\n').encode('utf-8')) def fmt(obj): @@ -197,27 +224,3 @@ def format_attrs(attrs): return '' entries = ['='.join((key, value)) for key, value in iteritems(attrs)] return '[' + ', '.join(entries) + ']' - - -if __name__ == '__main__': - from zipline.modelling.factor.technical import VWAP, MaxDrawdown, RSI - from zipline.data.equities import USEquityPricing - - with StderrHandler(): - vwap = VWAP(window_length=5) - dd_rank = MaxDrawdown([USEquityPricing.close], window_length=5).rank() - rsi_rank = (-RSI()).rank() - - score = ((dd_rank + rsi_rank) / 2) - - write_graph( - TermGraph( - { - 'vwap_pct': vwap.percentile_between(0, 20), - 'score': score, - 'score_gt': score > 10.0, - }, - ), - 'new', - ('png',), - )