sgtlearn

sgtlearn is a Python package for learning Shape Generalized Trees (SGTs) — a class of decision trees where each internal node applies a learnable, axis-aligned shape function to a feature, producing non-linear yet interpretable splits.

It implements the algorithms from the NeurIPS 2025 paper Empowering Decision Trees via Shape Function Branching, with a scikit-learn-compatible estimator API.

Highlights

• 🌳 Shape Generalized Trees (SGTs): each node applies a learnable, axis-aligned shape function to a feature for non-linear, interpretable splits.

• 👁 Interpretability: every node’s shape function can be visualized directly with plot_tree().

• ⚡ ShapeCART algorithm: an efficient native (C++/pybind11) induction method for learning SGTs from data.

• 🔀 Extensions: multi-way branching (\(\mathrm{SGT}_K\)) and bootstrap ensembling via RandomSGForestClassifier / RandomSGForestRegressor.

Note

This codebase is an efficient, working implementation of the paper’s algorithms. See the Release Roadmap for which features are implemented and which are planned. The canonical research code lives at optimal-uoft/Empowering-DTs-via-Shape-Functions.

Getting Started

• Installation
  • From PyPI
  • Requirements
  • Developer setup (build from source)
  • Build workflow (scikit-build + CMake)
• Quick Start
  • Inner vs. outer depth
  • Regression
  • Ensembles
  • TAO refinement
  • Categorical features

Tutorials

• Shape Functions — what an SGT node actually does
  • The knobs
  • What to notice
• Categorical features with feature_dict
  • The shape of a feature_dict
  • Building the dict for real data
• Feature importance
  • 1. Built-in impurity importances
  • 2. Permutation importance with categoricals
• SGT_K — multi-way branching
  • When does K > 2 help?
  • What to notice
• Tree Structure vs. Train/Test Accuracy
  • Setup
  • What to notice — over/underfitting
  • Takeaways
• Inspecting a Fitted SGT
  • Reading a node
  • The export is read-only
• Ensembling SGTs — Random SG Forests
  • What to notice
  • The knobs
  • How does it compare to a classic random forest?
• Regression with SGTs
  • Split criterion: squared_error vs absolute_error
  • Forests regress too
• TAO Refinement — a walkthrough
  • What is TAO?
  • Regularizing with lambda_
  • Forests
  • How to turn TAO off
  • Fit-time vs. post-hoc

API Reference

• API Reference
  • Estimators
  • Ensembles
  • TAO refinement
  • Plotting & Export
  • Datasets

Project

• Release Roadmap

Indices

• Index

• Module Index

• Search Page