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.
pip install supertree
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()
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.
- scikit-learn (
sklearn) - LightGBM
- XGBoost
- ONNX:
The package is compatible with a wide range of classifiers and regressors from these libraries, specifically:
DecisionTreeClassifierExtraTreeClassifierExtraTreesClassifierRandomForestClassifierGradientBoostingClassifierHistGradientBoostingClassifierDecisionTreeRegressorExtraTreeRegressorExtraTreesRegressorRandomForestRegressorGradientBoostingRegressorHistGradientBoostingRegressor
LGBMClassifierLGBMRegressorBooster
XGBClassifierXGBRFClassifierXGBRegressorXGBRFRegressorBooster
If we do not support the model you want to use, please let us know.
- Visualize decision tree from scikit-learn package
- 4 ways to vizualize decision tree from LightGBM
- How to visualize decision tree from Xgboost
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.
supertree is open source under the Apache License 2.0.
