API reference

Estimators

class difftree.DifferentiableTreeRegressor(depth=5, learning_rate=0.01, batch_size=256, n_epochs=200, gamma_init=1.0, gamma_max=5.0, delta_gamma=0.1, sigma_init=1.0, warmup_fraction=0.25, validation_fraction=0.1, n_calibration_steps=300, calibration_lr=0.02, standardize=True, leaf_type='constant', constant_outputs=None, jacobian_weight=1.0, hardening=False, gamma_hard=50.0, n_hard_epochs=50, hard_ridge=1e-06, continuous=False, continuity_weight=None, tied_splits=False, refine_splits=50, random_state=None, verbose=False)

Bases: RegressorMixin, BaseEstimator

Soft oblique decision tree of fixed depth with multivariate Gaussian leaves.

Predictive distribution at each input x is a mixture of multivariate Gaussians, one per leaf, weighted by soft path probabilities mu_l(x). For univariate y this reduces to the heteroscedastic Gaussian model; for multivariate y it is a structured conditional GMM (input-dependent MDN). All outputs share one partition (one set of splits).

Parameters:

• depth (int, default=5)

• learning_rate (float, default=0.01)

• batch_size (int, default=256)

• n_epochs (int, default=200)

• gamma_init (float) – Anneal split steepness 1 -> gamma_max during training.

• gamma_max (float) – Anneal split steepness 1 -> gamma_max during training.

• delta_gamma (float) – Anneal split steepness 1 -> gamma_max during training.

• sigma_init (float, default=1.0) – Initial leaf marginal standard deviation (in standardized y units).

• warmup_fraction (float, default=0.25)

• validation_fraction (float, default=0.1) – Fraction of training data held out for post-hoc gamma+alpha calibration.

• n_calibration_steps (calibration hyperparams.)

• calibration_lr (calibration hyperparams.)

• standardize (bool, default=True)

• leaf_type ({"constant", "affine"}, default="constant") – "constant": leaf l predicts the mean b_l. "affine": leaf l predicts A_l x + b_l. The Gaussian noise model is the same in both cases, so predict(return_std=True) works for either.

• constant_outputs (sequence of int or None, default=None) – With leaf_type="affine", output indices whose leaves stay constant (A_l row fixed at zero), e.g. intensive properties that enter bilinear blending terms. Ignored for constant leaves.

• jacobian_weight (float, default=1.0) – Weight of the Jacobian-matching (Sobolev) term, used only when fit receives dydx. The term is computed in standardized units.

• hardening (bool, default=False) – If True, fit ends with harden() on the same train/validation split.

• gamma_hard (float, default=50.0) – Final split steepness reached while hardening.

• n_hard_epochs (int, default=50) – Training epochs used to anneal gamma from gamma_max to gamma_hard before the partition is frozen.

• hard_ridge (float, default=1e-6) – Ridge strength of the least-squares leaf refit. Coefficients shrink toward the soft leaf values, which only matters for leaves with too few training points to fit on their own.

• continuous (bool or sequence of int, default=False) – Make the hardened model continuous across split boundaries for these outputs (True: every affine output). The hardening refit then requires neighbouring leaves to agree on every facet they share inside the training-data box, so a MILP optimum cannot exploit a jump between leaves. Only for affine outputs, and only appropriate where the true response is continuous.

• continuity_weight (float or None, default=None) – Untied trees only. None enforces continuity exactly. A float instead adds a penalty continuity_weight * ||mismatch||^2 on the coefficient mismatch (standardized units) across each shared facet, allowing small jumps.

• tied_splits (bool, default=False) – Every internal node at the same depth shares one hyperplane (an oblivious oblique tree). The partition is then an arrangement of depth hyperplanes, and a continuous model on it is the hinge model theta_0 . [1, x] + sum_d c_d max(0, w_d . x + b_d). With continuous, hardening fits that form directly and refines the hyperplanes, which keeps continuity cheap. Without tied splits, continuity forces leaves on either side of a deeper split to merge wherever that split meets a parent boundary.

• refine_splits (int, default=50) – Tied trees with continuous only: maximum Powell iterations used to move the hyperplanes so that the continuous fit’s least-squares error is smallest (0 keeps the trained hyperplanes).

• random_state (int or None)

• verbose (bool)

fit(X, y, dydx=None)

Fit the tree.

Parameters:

• X ((n, p) array)

• y ((n,) or (n, d_y) array)

• dydx ((n, d_y, p) array, optional) – Jacobian of y with respect to x at each training point, in original units. For 1D y a (n, p) array is also accepted. Requires leaf_type="affine"; adds a Jacobian-matching term weighted by jacobian_weight.

harden(X, y, dydx=None, *, continuous=None, bounds=None)

Turn the fitted soft tree into a crisp piecewise model.

Steps:

1. Keep training splits and leaves while gamma grows geometrically from gamma_max to gamma_hard over n_hard_epochs epochs, so the splits settle where crisp routing wants them.

2. Freeze the partition: route each training point crisply (apply()).

3. Refit every leaf by least squares on the points it holds (affine outputs regress on [1, x], constant outputs take the mean; dydx rows join the least-squares system when given). With continuous, the refit is one joint constrained least-squares problem in which neighbouring leaves agree on their shared facets; with tied_splits it is instead a hinging-hyperplanes fit, with the hyperplanes refined (refine_splits), which is continuous everywhere. Leaf noise covariances are refit from the residuals.

4. Fit the variance scale alpha in closed form on the held-out validation slice under crisp routing, and set every per-depth gamma to gamma_hard.

X, y are split with validation_fraction exactly as in fit(). Passing the training data used in fit is the usual choice. The result of comparing soft and hard predictions on the validation slice before and after hardening is stored in hardening_report_ (see soft_hard_gap()), together with max_boundary_jump, the largest jump of each output across any shared facet in the box (see difftree.export.boundary_jumps()).

continuous overrides the constructor setting for this call. bounds = (lo, hi) (original units) sets the box in which leaves must join up continuously; it defaults to the training-data box and should cover the region a MILP will search.

Returns self.

soft_hard_gap(X, y)

Compare soft and hard (crisp) predictions on (X, y).

Returns a dict with, per output, rmse_soft, rmse_hard, rmse_gap = rmse_hard - rmse_soft, r2_soft, r2_hard; the mean negative log-likelihoods nll_soft and nll_hard (original y units); mean_max_leaf_prob, the average probability of each point’s most likely leaf under soft routing (1.0 means routing is already crisp); and routing_agreement, the fraction of points whose most likely soft leaf is the crisp leaf. A large gap means the soft fit relies on blending leaves and is a misleading guide to the crisp (MILP-embeddable) model.

apply(X)

Leaf index (0 .. 2**depth - 1, BFS leaf order) each row routes to under crisp routing: node i goes right iff w_i . x + b_i > 0.

Return type:

ndarray

predict(X, return_std=False, return_components=False, hard=False)

Predict marginal mean (and optionally per-output marginal std).

For multivariate y, std is the per-output marginal standard deviation (sqrt of diagonal of the predictive covariance). Use predict_full_cov for the full covariance matrix.

With hard=True each row is routed crisply to one leaf (see apply()) and the prediction is that leaf’s Gaussian: the std is the leaf noise only and the routing component is zero.

Parameters:

• return_std (bool)

• return_components (bool)

• hard (bool)

predict_full_cov(X, hard=False)

Predict mean and full predictive covariance per sample.

Returns (mean, cov) of shapes (n, d_y) and (n, d_y, d_y), in original (un-standardized) y units. hard=True uses crisp routing, as in predict().

Parameters:

hard (bool)

sample_y(X, n_samples=1, random_state=None)

Marginal Monte Carlo draws from p(y|x).

Returns shape (n_samples, n, d_y) for multivariate or (n_samples, n) for univariate.

Parameters:

• n_samples (int)

• random_state (int | None)

predict_quantile(X, q, n_samples=1000, random_state=None)

Approximate marginal quantiles by Monte Carlo. Univariate only.

Parameters:

• n_samples (int)

• random_state (int | None)

predict_conditional(X, observed)

Conditional moments p(y_unobs | x, y_obs).

Parameters:

• X ((n, p) array)

• observed (dict {output_idx: value}) – Each value is a scalar (broadcast across rows) or an (n,) array.

Returns:

• cond_mean ((n, d_unobs) array)

• cond_cov ((n, d_unobs, d_unobs) array)

• log_evidence ((n,) array) – log p(y_obs | x), useful for ranking how well x is consistent with the observation.

sample_y_conditional(X, observed, n_samples=1, random_state=None)

Monte Carlo draws from p(y_unobs | x, y_obs).

Returns shape (n_samples, n, d_unobs) in original y units.

Parameters:

• observed (dict)

• n_samples (int)

• random_state (int | None)

inverse_design(targets, *, n_starts=16, n_steps=400, learning_rate=0.05, bounds=None, init_x=None, random_state=0)

Find x such that p(y_target | x) is maximized.

Parameters:

• targets (dict {output_idx: value}) – Output values to match. Values are scalars in original y units. Subset of outputs is supported (constrained inverse design).

• n_starts (number of random restarts.)

• n_steps (Adam steps per restart.)

• learning_rate (Adam learning rate.)

• bounds (tuple of (lo, hi) arrays in original x units, optional.) – Restricts the search to a feature box.

• init_x ((n_starts, p) array of starts in original x units, optional.)

• random_state (int)

Returns:

• x_best ((p,) ndarray in original x units.)

• log_lik_best (float) – log p(y_target | x_best) under the trained likelihood.

• x_all ((n_starts, p) ndarray of all restart results.)

• log_liks ((n_starts,) ndarray of all final log-likelihoods.)

score(X, y, sample_weight=None)

R^2 (univariate) or uniform-average R^2 across outputs (multivariate).

class difftree.DifferentiableTreeClassifier(depth=5, learning_rate=0.01, batch_size=256, n_epochs=200, gamma_init=1.0, gamma_max=5.0, delta_gamma=0.1, validation_fraction=0.1, n_calibration_steps=300, calibration_lr=0.05, standardize=True, random_state=None, verbose=False)

Bases: ClassifierMixin, BaseEstimator

Soft oblique decision tree of fixed depth with categorical leaves.

Predictive distribution at each x is a mixture of categoricals (one per leaf), weighted by the soft path probabilities. The total predictive entropy decomposes into an aleatoric (within-leaf) part and an epistemic (mutual-information / between-leaf-routing) part.

Calibration is performed post-hoc on a validation slice by fitting the per-depth gamma (split steepness) by minimizing validation NLL.

See DifferentiableTreeRegressor for shared hyperparameter docs.

Parameters:

• depth (int)

• learning_rate (float)

• batch_size (int)

• n_epochs (int)

• gamma_init (float)

• gamma_max (float)

• delta_gamma (float)

• validation_fraction (float)

• n_calibration_steps (int)

• calibration_lr (float)

• standardize (bool)

• random_state (int | None)

• verbose (bool)

apply(X)

Leaf index (0 .. 2**depth - 1, BFS leaf order) each row routes to under crisp routing: node i goes right iff w_i . x + b_i > 0.

Return type:

ndarray

predict_proba(X, return_components=False, hard=False)

Predict class probabilities.

Parameters:

• return_components (bool) – If True, also return aleatoric and epistemic predictive entropies (per sample, both shape (n,)).

• hard (bool) – If True, route each row crisply to one leaf (see apply()) and return that leaf’s class distribution. The epistemic (routing) entropy is then zero.

sample_y(X, n_samples=1, random_state=None)

Monte Carlo class-label samples from the predictive distribution.

Returns labels in the original class space, shape (n_samples, n).

Parameters:

• n_samples (int)

• random_state (int | None)

Export

Export a crisp (hardened) regressor for MILP embedding, in original units.

Under crisp routing (predict(hard=True)), internal node i sends x left iff a_i . x <= t_i, with the split undone from standardized units:

a_i = W_i / x_std, t_i = a_i . x_mean - b_i.

Each leaf is then a polytope {x : G_l x <= h_l}: the half-spaces along its root-to-leaf path intersected with an input box. Its leaf model is affine, y = coef_l x + intercept_l (coef_l has zero rows for outputs held constant per leaf), with Gaussian noise covariance cov_l.

Three entry points:

• leaf_polytopes() returns one LeafPolytope per leaf, with empty leaves (one LP per leaf) and, optionally, unreached leaves pruned.

• to_spec() / to_json() emit the solver-neutral v1 JSON spec of jax-ldt (jax_ldt.export.spec), so jax_ldt.export.from_json and its discopt adapter can consume a difftree model. Each oblique split becomes one lifted column of transform_matrix: the matrix is [I_p | a_1 ... a_k], the first p lifted features are the raw inputs (the leaf models use only these), and node i splits lifted feature p + j at t_i.

• predict_spec() evaluates such a spec in NumPy, as a reference.

Two helpers deal with split boundaries, where a piecewise-affine tree can jump. A MILP treats leaf polytopes as closed sets, so its optimum can sit on a boundary and use either neighbouring leaf’s model there:

• boundary_jumps() measures the largest jump of each output across every shared facet (use harden(continuous=True) to remove them);

• repair_solution() moves a MILP solution a distance delta into the interior of the leaf the solver chose (see leaf_from_binaries()), so that predict(hard=True) at the returned point uses that leaf.

class difftree.export.LeafPolytope(leaf, path=<factory>, G=None, h=None, coef=None, intercept=None, cov=None, radius=0.0, n_samples=None)

Bases: object

One crisp leaf region and its model, in original units.

Parameters:

• leaf (int)

• path (list)

• G (ndarray)

• h (ndarray)

• coef (ndarray)

• intercept (ndarray)

• cov (ndarray)

• radius (float)

• n_samples (int | None)

leaf

Leaf index in BFS leaf order, as returned by model.apply.

Type:

int

path

Internal nodes on the root-to-leaf path and the branch taken ("left" or "right").

Type:

list of (int, str)

G, h

Polytope {x : G x <= h}: one row per path split, then the input box (x <= hi and -x <= -lo).

Type:

ndarray

coef

Type:

(d_y, p) ndarray

intercept

Leaf mean coef @ x + intercept.

Type:

(d_y,) ndarray

cov

Leaf noise covariance (calibrated).

Type:

(d_y, d_y) ndarray

radius

Chebyshev radius of the polytope in standardized x units (radius of the largest ball inside it, capped at 1). Empty polytopes have 0.

Type:

float

n_samples

Rows of the X passed to leaf_polytopes() that route here.

Type:

int or None

contains(X, tol=1e-09)

Boolean mask of rows of X inside the (closed) polytope.

Parameters:

tol (float)

Return type:

ndarray

predict(X)

Leaf mean at the rows of X, shape (n, d_y).

Return type:

ndarray

difftree.export.leaf_polytopes(model, bounds=None, *, prune_empty=True, X=None, prune_unreached=False, tol=1e-07)

Leaf polytopes {x : G_l x <= h_l} and leaf models in original units.

Parameters:

• model (fitted DifferentiableTreeRegressor (ideally hardened).)

• bounds ((lo, hi), optional) – Input box in original units (arrays of length p, or scalars). Defaults to the training-data box (model.x_min_, model.x_max_).

• prune_empty (bool, default True) – Drop leaves whose polytope has no interior inside the box (one LP per leaf).

• X ((n, p) array, optional) – Points to count per leaf (n_samples), e.g. the training inputs.

• prune_unreached (bool, default False) – Also drop leaves that no row of X reaches. Requires X.

• tol (float) – A leaf is empty when its Chebyshev radius (standardized x units) is at most tol.

Return type:

list of LeafPolytope, in BFS leaf order.

difftree.export.to_spec(model, bounds=None, *, prune_empty=True, X=None, prune_unreached=False, tol=1e-07)

Serialize the crisp tree to jax-ldt’s solver-neutral v1 spec (a dict).

Arguments are as in leaf_polytopes(). Pruning collapses every internal node with a pruned side into its other child, so the spec only contains splits that separate kept leaves. Collapsing nodes whose pruned side is empty is exact inside the box; collapsing nodes whose side is only unreached (prune_unreached) reassigns that region to a neighbouring leaf.

Layout: transform_matrix is [I_p | a_1 ... a_k] (one lifted column per kept internal node); internal nodes split lifted feature p + j at t and go left iff x'_f <= t; leaf_params[node] is [bias; weights] over the lifted features, with weights only on the first p (raw) columns. Extra difftree information (bounds, BFS leaf ids, leaf covariances) sits under the top-level "difftree" key, which jax-ldt ignores.

Parameters:

• prune_empty (bool)

• prune_unreached (bool)

• tol (float)

Return type:

dict

difftree.export.to_json(model, path=None, **kwargs)

to_spec() as a JSON string, also written to path if given.

Parameters:

path (str | Path | None)

Return type:

str

difftree.export.predict_spec(spec, X)

Evaluate a v1 spec (dict or JSON string) at X with NumPy.

Returns (n,) for single-target specs, else (n, n_targets).

Parameters:

spec (dict | str)

Return type:

ndarray

difftree.export.boundary_jumps(model, bounds=None, *, tol=1e-07)

Largest jump of each output across every shared leaf facet in the box.

For every pair of leaves that share a (p-1)-dimensional facet inside the box, the difference of their affine models is maximized (in absolute value) over that facet with two LPs per output.

Parameters:

• model (fitted DifferentiableTreeRegressor.)

• bounds ((lo, hi), optional) – Input box in original units; defaults to the training-data box.

• tol (float) – Minimum facet Chebyshev radius (standardized x units) to count.

Returns:

• list of dict, sorted by largest jump first, each with node (internal

• node whose hyperplane holds the facet), left / right (BFS leaf

• indices), point (the facet’s Chebyshev center, original units),

• max_jump ((d_y,) array, original y units) and worst_point

• ((d_y, p) array (where on the facet each output’s jump is largest). A)

• continuous output has max_jump at round-off level on every facet.

Return type:

list[dict]

difftree.export.leaf_from_binaries(spec, z)

BFS leaf index of the active leaf binary z of a MILP embedding.

z holds one value per leaf of spec, in the order the spec’s leaves appear (np.where(spec["tree"]["is_leaf"])), which is the order jax-ldt’s discopt adapter uses for its leaf binaries.

Parameters:

spec (dict | str)

Return type:

int

difftree.export.repair_solution(model, x, leaf, *, delta=1e-06, bounds=None)

Move x into the interior of leaf’s polytope.

Solves one LP: the closest point (max-norm, standardized x units) to x that lies at least delta inside every face of the leaf’s polytope within the box. Crisp routing of the result is leaf, so predict(hard=True) there uses the leaf model the MILP optimized, and the surrogate output changes by at most delta times its slope. Choose delta above the solver’s feasibility tolerance (in standardized units).

Parameters:

• model (fitted DifferentiableTreeRegressor.)

• x ((p,) array, original units (e.g. the MILP's optimal x).)

• leaf (int) – BFS leaf index the solver selected (see leaf_from_binaries()).

• delta (float, default 1e-6)

• bounds ((lo, hi), optional) – Input box in original units; defaults to the training-data box.

Return type:

(p,) ndarray in original units.