Improve decision tree plotting in Jupyter environment

[nvdv]

Recently I used built-in visualizer for DecisionTreeClassifier in Jupyter and can say that
its interface could be better (example is taken from docs):

>>> from IPython.display import Image  
>>> dot_data = StringIO()  
>>> tree.export_graphviz(clf, out_file=dot_data,  
                         feature_names=iris.feature_names,  
                         class_names=iris.target_names,  
                         filled=True, rounded=True,  
                         special_characters=True)  
>>> graph = pydot.graph_from_dot_data(dot_data.getvalue())  
>>> Image(graph.create_png())  

In addition this does not work on Python 3, since at the time of writing pydot cannot be installed for Python 3.
The ideal solution will be something like

from sklearn import tree
tc = tree.DecisionTreeClassifier()
...
tree.plot(tc) # or even tc.plot()

but in this case tree module should depend on pydot and IPython.display.image modules.
I can fix this issue, but what is the best way to do this?

[nvdv]

One of the possible solutions might be something like

class OptionalImportError(Exception):
    ...


def plot(params):
    try:
        import pydot
        from IPython.display import Image  
    except ImportError:
        raise OptionalImportError
    # do actual plotting
    ...

[nvdv]

anyone?

[amueller]

I think it would be nice to make it easier to plot trees. Currently it's a bit of a pain. But just returning the tree structure from export_graphviz would mean you can show the tree in IPython by just calling export_graphviz. You could just add a tree_to_graphiviz function or something next to it, that returns a graphvis object instead of writing to a file.

[amueller]

btw, this works:

from sklearn.tree import export_graphviz
import graphviz

export_graphviz(tree, out_file="mytree.dot")
with open("mytree.dot") as f:
    dot_graph = f.read()
graphviz.Source(dot_graph)

And the returned graphviz value is automatically rendered in Jupyter notebook. It would be nice not to need to create a tempfile, though. So I'd like to do

from sklearn.tree import convert_to_graphviz
convert_to_graphviz(tree)

That would have an optional dependency on graphviz, though.
The other option would be to not create the object and just return the string:

from sklearn.tree import convert_to_graphviz
import graphviz

graphviz.Source(export_graphviz(tree))

[nvdv]

Thanks!

[raghavrv]

+(np.inf ** np.inf)

[raghavrv]

This is my current setup for diagnosing trees

import matplotlib.pyplot as plt
import pygraphviz as pgv
import networkx as nx
import pygraphviz
import matplotlib.image as img
import matplotlib.pyplot as plt

from sklearn.tree import export_graphviz
from StringIO import StringIO
from io import BytesIO

def get_graph(dtc, n_classes, feat_names=None, size=[7, 7]):
    dot_file = StringIO()
    image_file = BytesIO()

    # Get the dot graph of our decision tree
    export_graphviz(dtc, out_file=dot_file, feature_names=feat_names, rounded=True, filled=True,
                    special_characters=True, class_names=map(str, range(n_classes)), max_depth=10)
    dot_file.seek(0)

    # Convert this dot graph into an image
    g = pygraphviz.AGraph(dot_file.read())
    g.layout('dot')
    # g.draw doesn't work when the image object doesn't have a name (with a proper extension)
    image_file.name = "image.png"
    image_file.seek(0)
    g.draw(path=image_file)
    image_file.seek(0)

    # Plot it
    plt.figure().set_size_inches(*size)
    plt.axis('off')
    plt.imshow(img.imread(fname=image_file))
    plt.show()

[ogrisel]

This was closed by github when merging #7342. However I think we could still do better as suggested in #6261 (comment).