Estimators¶

Single Shape Generalized Tree estimators following the scikit-learn API.

Both estimators accept tao_n_runs and tao_lambda; TAO runs automatically at the end of fit() when tao_n_runs > 0. See TAO refinement for behaviour and post-hoc TAO_refine().

After fitting, feature_importances_ gives normalized importances over logical features, aligned with processed_features_. Forests expose mean_feature_importances_ and std_feature_importance_ instead (see Ensembles). Prefer built-in importances when TAO is off (tao_n_runs=0); see Feature importance for permutation importance with categoricals.

SGTClassifier¶

class sgtlearn.SGTClassifier(*, criterion='gini', num_partitions=2, max_depth=None, max_leaf_nodes=None, min_samples_leaf=1, min_impurity_decrease=0.0, inner_max_depth=3, inner_max_leaf_nodes=8, inner_min_samples_leaf=1, inner_min_impurity_decrease=0.0, coordinate_descent_max_iters=20, coordinate_descent_patience=5, coordinate_descent_smart_init=True, random_state=42, max_features=None, class_weight=None, tao_n_runs=10, tao_lambda=0.0)[source]¶

Bases: ClassifierMixin, BaseShapeCART

Shape Generalized Tree classifier.

A decision tree where each internal node applies a learnable, axis-aligned shape function to a single feature rather than a single threshold. The shape function is itself an inner tree (univariate, depth-limited) that partitions the feature’s value range into num_partitions bins; the outer tree then routes samples through those bins to grow the overall classifier. Training is performed by the native ShapeCART C++ trainer exposed through ClassificationShapeGeneralizedTree.

The estimator follows the scikit-learn ClassifierMixin contract and is compatible with sklearn pipelines, cross-validators, and metaestimators.

Parameters:

criterion ({"gini", "entropy"}, default="gini") – Impurity criterion used at outer-tree splits, forwarded to the native trainer.
num_partitions (int, default=2) – Number of branches each outer split fans out into (i.e. the arity of the shape function). 2 reproduces standard binary tree; larger values yield the SGT_K multi-way variant.
max_depth (int, optional) – Maximum depth of the outer tree. None (default) means grow until another stopping criterion fires.
max_leaf_nodes (int, optional) – Maximum number of leaves in the outer tree. None means unlimited.
min_samples_leaf (int, default=1) – Minimum number of training samples required at an outer leaf.
min_impurity_decrease (float, default=0.0) – Minimum impurity decrease required to accept an outer split.
inner_max_depth (int, default=3) – Maximum depth of the inner tree that defines the shape function on each feature. 1 corresponds to a standard CART threshold split; larger values produce richer shapes.
inner_max_leaf_nodes (int, default=8) – Maximum number of bins (leaves) the inner tree may form on a feature.
inner_min_samples_leaf (int, default=1) – Minimum samples per inner-tree leaf.
inner_min_impurity_decrease (float, default=0.0) – Minimum impurity decrease required to accept an inner split.
coordinate_descent_max_iters (int, default=20) – Maximum coordinate-descent iterations used to refine the bin-to-branch assignment after the inner tree is built.
coordinate_descent_patience (int, default=5) – Number of non-improving iterations tolerated before coordinate descent terminates early.
coordinate_descent_smart_init (bool, default=True) – If True, seed coordinate descent with a k-means clustering of the bin statistics; if False, use round-robin assignment.
random_state (int, optional, default=42) – Seed forwarded to the native trainer. None is treated as 42.
max_features (int, float, {"sqrt", "log2"} or None, default=None) –
Number of features sampled (without replacement) when searching for a split:
- None: use all n_features columns.
- int k >= 1: use min(k, n_features) columns.
- float c in (0, 1]: use max(1, int(c * n_features)) columns.
- "sqrt": use max(1, int(sqrt(n_features))) columns.
- "log2": use max(1, int(log2(n_features))) columns.
String values are case-insensitive in the native binding.
class_weight (dict, optional) – Per-class weights multiplied into sample_weight before training. Keys are class labels as in y; values are non-negative floats.
tao_n_runs (int)
tao_lambda (float)

classes_¶

Unique class labels observed during fit(), in the order used by predict_proba() columns.

Type:: ndarray of shape (n_classes,)

n_classes_¶

Number of classes.

Type:: int

n_features_in_¶

Number of features in X passed to fit().

Type:: int

feature_importances_¶

Normalized impurity-based importances from the fitted tree. Index i corresponds to processed_features_ entry i (same order as the logical features passed to the native trainer).

Type:: ndarray of shape (n_logical_features,)

processed_features_¶

Resolved logical features used at fit(). features[i] and logical_names[i] align with feature_importances_[i]. logical_names are stringified feature_dict keys (auto-filled columns use "0", "1", …).

Type:: ProcessedFeatures

Notes

Internally, X is cast to C-contiguous float32 and y to uint64 before being passed to the native trainer. Sparse input is not supported. NaN in X is handled by the native trainer (non-finite values are sorted to the feature tail and routed to the last bin at inference). Infinity in X is rejected.

References

Upadhya, N. and Cohen, E. “Empowering Decision Trees via Shape Function Branching.” NeurIPS, 2025.

See also

SGTRegressor: Regression counterpart.
sgtlearn.ensemble.RandomSGForestClassifier: Bootstrap ensemble over SGTClassifier.

Examples

>>> from sklearn.datasets import make_classification
>>> from sgtlearn import SGTClassifier
>>> X, y = make_classification(n_samples=500, random_state=0)
>>> clf = SGTClassifier(max_depth=4, random_state=42).fit(X, y)
>>> clf.predict(X[:5]).shape
(5,)

fit(X, y, sample_weight=None, *, feature_dict=None, processed_features=None, check_input=True)[source]¶

Fit the tree on X and class labels y.

Parameters:

X (array-like of shape (n_samples, n_features)) – Training features.
y (array-like of shape (n_samples,)) – Target class labels.
sample_weight (array-like of shape (n_samples,), optional) – Per-sample weights.
feature_dict (mapping, optional) – {logical_name: [columns]} grouping of columns into logical features. Keys are int or str logical names. Values are column indices (int) or column names (str); names require a pandas DataFrame X or an explicit column_names. A group of more than one column is treated as a single categorical feature (routed through the one-hot inner discretizer); singletons are continuous. Columns not mentioned default to continuous singletons. See configure_feature_dict() for the full resolution rules.
processed_features (ProcessedFeatures, optional) – Pre-resolved features from configure_feature_dict() (for ensembles that should resolve features once).
check_input (bool, default=True) – If False, X and y are not validated (for callers that already ran check_X_y()).

Return type:

SGTClassifier

predict(X)[source]¶

Predict class labels (via the fitted label encoder).

Parameters:: X (ndarray)
Return type:: ndarray

predict_proba(X)[source]¶

Return shape (n_samples, n_classes) probabilities aligned with classes_ order.

Parameters:: X (ndarray)
Return type:: ndarray

tree_export()[source]¶

Return a flat dict snapshot of the fitted tree.

See sgtlearn._export.plot_tree for the canonical consumer. Keys: num_partitions, num_nodes, root_index, num_classes, criterion, nodes (list of per-node dicts).

Return type:: dict

set_fit_request(*, check_input='$UNCHANGED$', feature_dict='$UNCHANGED$', processed_features='$UNCHANGED$', sample_weight='$UNCHANGED$')¶

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

check_input (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for check_input parameter in fit.
feature_dict (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for feature_dict parameter in fit.
processed_features (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for processed_features parameter in fit.
sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.
self (SGTClassifier)

Returns:

self – The updated object.

Return type:

object

set_score_request(*, sample_weight='$UNCHANGED$')¶

Configure whether metadata should be requested to be passed to the score method.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in score.
self (SGTClassifier)

Returns:

self – The updated object.

Return type:

object

SGTRegressor¶

class sgtlearn.SGTRegressor(*, criterion='squared_error', num_partitions=2, max_depth=None, max_leaf_nodes=None, min_samples_leaf=1, min_impurity_decrease=0.0, inner_max_depth=3, inner_max_leaf_nodes=8, inner_min_samples_leaf=1, inner_min_impurity_decrease=0.0, coordinate_descent_max_iters=20, coordinate_descent_patience=5, coordinate_descent_smart_init=True, random_state=42, max_features=None, tao_n_runs=10, tao_lambda=0.0)[source]¶

Bases: RegressorMixin, BaseShapeCART

Shape Generalized Tree regressor.

Regression analogue of SGTClassifier. Each internal node applies a learnable shape function (an inner univariate tree) to a single feature, and the outer tree routes samples through the resulting bins. Training is performed by the native ShapeCART C++ trainer exposed through RegressionShapeGeneralizedTree.

Compatible with the scikit-learn RegressorMixin contract.

Parameters:

criterion ({"squared_error", "mse", "absolute_error", "mae"}, default="squared_error") – Loss used at outer-tree splits. "mse" and "mae" are accepted as aliases.
num_partitions (int, default=2) – Number of branches each outer split fans out into (i.e. the arity of the shape function). 2 reproduces standard binary tree; larger values yield the SGT_K multi-way variant.
max_depth (int, optional) – Maximum depth of the outer tree. None means grow until another stopping criterion fires.
max_leaf_nodes (int, optional) – Maximum number of leaves in the outer tree. None means unlimited.
min_samples_leaf (int, default=1) – Minimum number of samples required at an outer leaf.
min_impurity_decrease (float, default=0.0) – Minimum impurity decrease required to accept an outer split.
inner_max_depth (int, default=3) – Maximum depth of the inner tree defining the shape function on each feature. 1 reduces to a standard threshold split.
inner_max_leaf_nodes (int, default=8) – Maximum number of bins (leaves) the inner tree may form on a feature.
inner_min_samples_leaf (int, default=1) – Minimum samples per inner-tree leaf.
inner_min_impurity_decrease (float, default=0.0) – Minimum impurity decrease required to accept an inner split.
coordinate_descent_max_iters (int, default=20) – Maximum coordinate-descent iterations used to refine the bin-to-branch assignment after the inner tree is built.
coordinate_descent_patience (int, default=5) – Number of non-improving iterations tolerated before coordinate descent terminates early.
coordinate_descent_smart_init (bool, default=True) – Accepted for API symmetry with SGTClassifier but ignored by the regression trainer: regression always seeds inner bin-to-partition assignments round-robin (no k-means initialisation).
random_state (int, optional, default=42) – Seed forwarded to the native trainer. None is treated as 42.
max_features (int, float, {"sqrt", "log2"} or None, default=None) – Per-split feature subsampling. Same semantics as SGTClassifier.
tao_n_runs (int)
tao_lambda (float)

n_features_in_¶

Number of features seen during fit().

Type:: int

feature_importances_¶

Normalized impurity-based importances from the fitted tree. Index i corresponds to processed_features_ entry i (same order as the logical features passed to the native trainer).

Type:: ndarray of shape (n_logical_features,)

processed_features_¶

Type:: ProcessedFeatures

Notes

Internally, X and y are cast to C-contiguous float32 before being passed to the native trainer. Sparse input is not supported. NaN in X is handled by the native trainer (non-finite values are sorted to the feature tail and routed to the last bin at inference). Infinity in X is rejected.

For squared_error/mse, the trainer runs coordinate descent after the round-robin seed and keeps the refined assignment only if branch MSE improves clearly; otherwise it restores the seed. absolute_error/mae skips coordinate descent entirely.

References

Upadhya, N. and Cohen, E. “Empowering Decision Trees via Shape Function Branching.” NeurIPS, 2025.

See also

SGTClassifier: Classification counterpart.
sgtlearn.ensemble.RandomSGForestRegressor: Bootstrap ensemble over SGTRegressor.

Examples

>>> from sklearn.datasets import make_regression
>>> from sgtlearn import SGTRegressor
>>> X, y = make_regression(n_samples=500, random_state=0)
>>> reg = SGTRegressor(max_depth=4, random_state=42).fit(X, y)
>>> reg.predict(X[:5]).shape
(5,)

fit(X, y, sample_weight=None, *, feature_dict=None, processed_features=None, check_input=True)[source]¶

Fit the tree on X and continuous targets y.

Parameters:

X (array-like of shape (n_samples, n_features)) – Training features.
y (array-like of shape (n_samples,)) – Target values.
sample_weight (array-like of shape (n_samples,), optional) – Per-sample weights.
feature_dict (mapping, optional) – {logical_name: [columns]} grouping of columns into logical features. Keys are int or str logical names. Values are column indices (int) or column names (str); names require a pandas DataFrame X or an explicit column_names. A group of more than one column is treated as a single categorical feature (routed through the one-hot inner discretizer); singletons are continuous. Columns not mentioned default to continuous singletons. See configure_feature_dict() for the full resolution rules.
processed_features (ProcessedFeatures, optional) – Pre-resolved features from configure_feature_dict().
check_input (bool, default=True) – If False, X and y are not validated (for callers that already ran check_X_y()).

Return type:

SGTRegressor

tree_export()[source]¶

Return a flat dict snapshot of the fitted tree.

See sgtlearn._export.plot_tree for the canonical consumer. Keys: num_partitions, num_nodes, root_index, criterion, nodes (list of per-node dicts). Regression has no num_classes.

Return type:: dict

set_fit_request(*, check_input='$UNCHANGED$', feature_dict='$UNCHANGED$', processed_features='$UNCHANGED$', sample_weight='$UNCHANGED$')¶

Configure whether metadata should be requested to be passed to the fit method.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

check_input (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for check_input parameter in fit.
feature_dict (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for feature_dict parameter in fit.
processed_features (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for processed_features parameter in fit.
sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.
self (SGTRegressor)

Returns:

self – The updated object.

Return type:

object

set_score_request(*, sample_weight='$UNCHANGED$')¶

Configure whether metadata should be requested to be passed to the score method.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in score.
self (SGTRegressor)

Returns:

self – The updated object.

Return type:

object

BaseShapeCART¶

class sgtlearn.BaseShapeCART[source]¶

Bases: BaseEstimator

Shared sklearn BaseEstimator hook point for shape-generalized trees.

Subclasses own the native backend handle (_est) and validation rules.

property feature_importances_: ndarray¶

Normalized per-feature importances from the fitted tree.

Length matches the number of logical features passed to fit (one-to-one with processed_features_). Available only after training.

property processed_features_: ProcessedFeatures¶

Resolved logical features used at fit, aligned with importances.

features[i] and logical_names[i] correspond to feature_importances_[i]. logical_names are the feature_dict keys (stringified); omitted columns are filled as "0", "1", …

Feature configuration¶

Group columns into logical features (e.g. one-hot categorical blocks) for training. Pass the mapping directly as fit(..., feature_dict=...), or pre-resolve it once with configure_feature_dict() and pass the result as fit(..., processed_features=...).

sgtlearn.configure_feature_dict(n_features, feature_dict=None, *, column_names=None)[source]¶

Resolve logical features for tree training.

Parameters:

n_features (int) – Number of columns in X.
feature_dict (Mapping[int | str, Sequence[int | str]] | None) – Mapping {logical_key: [column indices or names]}. Keys may be int or str (logical feature names). Values are column indices (int) or column names (str) when column_names or a pandas DataFrame was used for training. A group with more than one column is categorical; singletons are continuous. Unmentioned columns are filled in as continuous singletons. When omitted, each column is its own continuous feature.
column_names (Sequence[str] | None) – Names of columns in X, used to resolve string column references in feature_dict.

Returns:

Feature list consumed by the native fit(..., features=...) binding.

Return type:

ProcessedFeatures

class sgtlearn._features.ProcessedFeatures(features, logical_names=())[source]¶

Bases: object

Resolved logical features ready for the native trainer.

Parameters:

features (list[dict[str, Any]])
logical_names (tuple[str, ...])

features¶

List of {"type": "continuous"|"categorical", "indices": [...]} dicts in trainer order. Index i aligns with estimator.feature_importances_[i] after fit().

Type:: list[dict[str, Any]]

logical_names¶

Parallel names for features. When feature_dict is supplied, these are the stringified keys; omitted columns are filled as "0", "1", …. Default (no feature_dict) uses "0" … "n_features-1" even if X is a pandas DataFrame (DataFrame column names are stored on feature_names_in_ instead).

Type:: tuple[str, …]