Visualize decision trees interactively in Jupyter, JupyterLab, and Google Colab. Zoom, pan, collapse nodes, and trace sample paths - all inside your notebook.
Works with scikit-learn, XGBoost, LightGBM, and ONNX.
Installation
Quick Start
Visualize Decision Tree classifier on iris data
from sklearn.tree import DecisionTreeClassifier from sklearn.datasets import load_iris from supertree import SuperTree # Load the iris dataset iris = load_iris() # Train model model = DecisionTreeClassifier(max_depth=3) model.fit(iris.data, iris.target) # Initialize supertree super_tree = SuperTree(model, iris.data, iris.target, iris.feature_names, iris.target_names) # show tree in your notebook super_tree.show_tree()
It works with trees ensembles too - Random Forest Regressor Example
from sklearn.ensemble import RandomForestRegressor from sklearn.datasets import load_diabetes from supertree import SuperTree # <- import supertree :) # Load the diabetes dataset diabetes = load_diabetes() X = diabetes.data y = diabetes.target # Train model model = RandomForestRegressor(n_estimators=100, max_depth=3, random_state=42) model.fit(X, y) # Initialize supertree super_tree = SuperTree(model,X, y) # show tree with index 2 in your notebook super_tree.show_tree(2)
There are more code snippets in the examples directory.
Supported Libraries
- scikit-learn (
sklearn) - LightGBM
- XGBoost
- ONNX:
Supported Algorithms
The package is compatible with a wide range of classifiers and regressors from these libraries, specifically:
Scikit-learn
DecisionTreeClassifierExtraTreeClassifierExtraTreesClassifierRandomForestClassifierGradientBoostingClassifierHistGradientBoostingClassifierDecisionTreeRegressorExtraTreeRegressorExtraTreesRegressorRandomForestRegressorGradientBoostingRegressorHistGradientBoostingRegressor
LightGBM
LGBMClassifierLGBMRegressorBooster
XGBoost
XGBClassifierXGBRFClassifierXGBRegressorXGBRFRegressorBooster
If we do not support the model you want to use, please let us know.
Articles
- Visualize decision tree from scikit-learn package
- 4 ways to vizualize decision tree from LightGBM
- How to visualize decision tree from Xgboost
Support
If you encounter any issues, find a bug, or have a feature request, we would love to hear from you! Please don't hesitate to reach out to us at supertree/issues. We are committed to improving this package and appreciate any feedback or suggestions you may have.
License
supertree is open source under the Apache License 2.0.