Package: pybaobabdt
pybaobabdt PackageThe pybaobabdt package provides a python implementation for the visualization of decision trees.
The technique is based on the scientific paperBaobabView: Interactive construction and analysis of decision treesdeveloped by the TU/e.
A typical decision tree is visualized using a standard node link diagram:The problem, however, is that information is not easily extracted from this. Which classes are
easy to separate for example, which classes are similar, where does the main flow of items go etc.
Therefore, we developed techniques to answer these questions with a scalable visualization:Note, this is the same decision tree as the standard node-link diagram above. Each class is represented by a
color, the width of the link represents the number of items flowing from one node to the other.InstallationCurrently it is supported on Python3.6 onwards. The package can be installed through pip:$pipinstallpybaobabdtRequirementsThis implementation requiresGraphviz. Graphviz can be installed using:$sudoapt-getinstallgraphvizgraphviz-devFurthermore it depends on the following python packages (sklearn, numpy, pygraphviz, matplotlib, scipy, pandas), which can be installed through pip:$python3-mpipinstall-rrequirements.txtUsageThe following example illustrates the ease of use of this package. First build (or load) a decision tree classifier
with sklearn:importpybaobabdtimportpandasaspdfromscipy.ioimportarfffromsklearn.treeimportDecisionTreeClassifierdata=arff.loadarff('winequality-red.arff')df=pd.DataFrame(data[0])y=list(df['class'])features=list(df.columns)features.remove('class')X=df.loc[:,features]clf=DecisionTreeClassifier().fit(X,y)Next, use pybaobab to visualize it:ax=pybaobabdt.drawTree(clf,size=10,dpi=72,features=features)You can then save it to a file with for example:ax.get_figure().savefig('tree.png',format='png',dpi=300,transparent=True)Also, trees from a RandomForest classifier can be visualized and saved to a high-resolution image for inspection:importpybaobabdtimportpandasaspdfromscipy.ioimportarffimportmatplotlib.pyplotaspltfromsklearn.ensembleimportRandomForestClassifierdata=arff.loadarff('vehicle.arff')df=pd.DataFrame(data[0])y=list(df['class'])features=list(df.columns)features.remove('class')X=df.loc[:,features]clf=RandomForestClassifier(n_estimators=20,n_jobs=-1,random_state=0)clf.fit(X,y)Save to image:size=(15,15)plt.rcParams['figure.figsize']=sizefig=plt.figure(figsize=size,dpi=300)foridx,treeinenumerate(clf.estimators_):ax1=fig.add_subplot(5,4,idx+1)pybaobabdt.drawTree(tree,model=clf,size=15,dpi=300,features=features,ax=ax1)fig.savefig('random-forest.png',format='png',dpi=1200,transparent=True)OptionsThere are several different options that can be used in the drawTree function.colormap='plasma' (all matplotlibcolormapsare supported)You can also define your own colormap, which could be useful to highlight a specific class for example:#colors = [[1,0,0], [0,1,0], [0,0,1], [1,1,0]]colors=["gray","gray","purple","gray"]colorMap=ListedColormap(colors)ax=pybaobabdt.drawTree(clf,size=10,dpi=72,features=features,colormap=colorMap)maxdepth=3 (set the maximum depth of the tree to render, this can be useful for large trees, to inspect only the top splits.)ratio=0.5 (sets the aspect ratio of the tree, default = 1)Note that examples can be found in the 'notebooks' folder containing jupyter notebook examples.LicenseGNU General Public License v3.0ReferenceIf you need to reference this work please use the following bibtex entry:@INPROCEEDINGS{Elzen2011,author={van den Elzen, Stef and van Wijk, Jarke J.},booktitle={2011 IEEE Conference on Visual Analytics Science and Technology (VAST)},title={BaobabView: Interactive construction and analysis of decision trees},year={2011},pages={151-160},doi={10.1109/VAST.2011.6102453}}S. van den Elzen and J. J. van Wijk, "BaobabView: Interactive construction and analysis of decision trees," 2011 IEEE Conference on Visual Analytics Science and Technology (VAST), 2011, pp. 151-160, doi: 10.1109/VAST.2011.6102453.