"""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``.
Four entry points:
- :func:`leaf_polytopes` returns one :class:`LeafPolytope` per leaf, with
empty leaves (one LP per leaf) and, optionally, unreached leaves pruned.
- :func:`to_spec` / :func:`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``.
- :func:`predict_spec` evaluates such a spec in NumPy, as a reference.
- :func:`to_pyomo` adds the tree to a Pyomo model as a disjunction over its
leaves (one ``Disjunct`` per leaf polytope and model), reformulated as a
MILP with ``gdp.bigm`` or ``gdp.hull``; :func:`leaf_from_pyomo` reads the
active leaf back after a solve. For tied, continuous trees it can instead
write the hinge form: exact for C0 models, a piecewise-linear MILP with a
guaranteed error bound (:func:`pwl_error_bound`) or a smooth NLP for C1
models (``smoothness=1``).
A C1 model is not piecewise affine, so the leaf-based exports
(:func:`leaf_polytopes`, :func:`to_spec`, :func:`to_json` and the ``"tree"``
formulation of :func:`to_pyomo`) refuse it rather than export its leaf models,
which describe it only outside its bands.
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:
- :func:`boundary_jumps` measures the largest jump of each output across every
shared facet (use ``harden(continuous=True)`` to remove them);
- :func:`repair_solution` moves a MILP solution a distance ``delta`` into the
interior of the leaf the solver chose (see :func:`leaf_from_binaries`), so
that ``predict(hard=True)`` at the returned point uses that leaf.
"""
from __future__ import annotations
import json
import warnings
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
from scipy.optimize import linprog
from sklearn.utils.validation import check_is_fitted
from . import _geometry, _model
SPEC_VERSION = 1
[docs]
@dataclass
class LeafPolytope:
"""One crisp leaf region and its model, in original units.
Attributes
----------
leaf : int
Leaf index in BFS leaf order, as returned by ``model.apply``.
path : list of (int, str)
Internal nodes on the root-to-leaf path and the branch taken
(``"left"`` or ``"right"``).
G, h : ndarray
Polytope ``{x : G x <= h}``: one row per path split, then the input
box (``x <= hi`` and ``-x <= -lo``).
coef : (d_y, p) ndarray
intercept : (d_y,) ndarray
Leaf mean ``coef @ x + intercept``.
cov : (d_y, d_y) ndarray
Leaf noise covariance (calibrated).
radius : float
Chebyshev radius of the polytope in standardized x units (radius of
the largest ball inside it, capped at 1). Empty polytopes have 0.
n_samples : int or None
Rows of the ``X`` passed to :func:`leaf_polytopes` that route here.
"""
leaf: int
path: list = field(default_factory=list)
G: np.ndarray = None
h: np.ndarray = None
coef: np.ndarray = None
intercept: np.ndarray = None
cov: np.ndarray = None
radius: float = 0.0
n_samples: int | None = None
[docs]
def contains(self, X, tol: float = 1e-9) -> np.ndarray:
"""Boolean mask of rows of ``X`` inside the (closed) polytope."""
X = np.atleast_2d(np.asarray(X, dtype=np.float64))
return np.all(X @ self.G.T <= self.h + tol, axis=1)
[docs]
def predict(self, X) -> np.ndarray:
"""Leaf mean at the rows of ``X``, shape (n, d_y)."""
X = np.atleast_2d(np.asarray(X, dtype=np.float64))
return X @ self.coef.T + self.intercept
# ---------------------------------------------------------------------------
# Parameters in original units
# ---------------------------------------------------------------------------
def _check_model(model):
check_is_fitted(model, "params_")
if not hasattr(model, "y_mean_"):
raise TypeError("export supports DifferentiableTreeRegressor only.")
if not getattr(model, "hardened_", False):
warnings.warn(
"Exporting a tree that was not hardened: the export reproduces "
"predict(hard=True), which can differ from the soft predict(). "
"Call model.harden(X, y) (or fit with hardening=True) first, and "
"check model.soft_hard_gap(X, y).",
UserWarning,
stacklevel=3,
)
def _require_tree_form(model, what):
if getattr(model, "smoothness_", 0) == 1:
raise ValueError(
f"{what} exports a piecewise-affine tree, but this model is C1 "
"(smoothness=1) and is quadratic near its splits. Use "
"to_pyomo(..., formulation='pwl', K=...) for a MILP approximation "
"with a guaranteed error bound, or formulation='smooth' for an NLP."
)
def _resolve_bounds(model, bounds):
if bounds is None:
lo, hi = model.x_min_, model.x_max_
else:
lo, hi = bounds
p = model.n_features_in_
lo = np.broadcast_to(np.asarray(lo, dtype=np.float64), (p,)).copy()
hi = np.broadcast_to(np.asarray(hi, dtype=np.float64), (p,)).copy()
if np.any(lo > hi):
raise ValueError("bounds: every lower bound must be <= its upper bound.")
return lo, hi
def _splits_original(model):
"""Split normals ``a`` (n_internal, p) and thresholds ``t`` (n_internal,):
node ``i`` goes left iff ``a_i . x <= t_i``."""
W = np.asarray(model.params_["W"], dtype=np.float64)
b = np.asarray(model.params_["b"], dtype=np.float64)
a = W / model.x_std_[None, :]
t = a @ model.x_mean_ - b
return a, t
def _leaves_original(model):
"""Leaf coefficients (L, d_y, p), intercepts (L, d_y), covariances
(L, d_y, d_y) in original units."""
import jax.numpy as jnp
P = model.params_
sy, my = model.y_std_, model.y_mean_
sx, mx = model.x_std_, model.x_mean_
coef = sy[None, :, None] * P["leaf_A"] / sx[None, None, :]
intercept = my[None, :] + sy[None, :] * P["leaf_mu"] - coef @ mx
L = np.asarray(
_model.build_chol(jnp.asarray(P["leaf_chol_raw"]), jnp.asarray(P["log_alpha"])),
dtype=np.float64,
)
cov = (L @ np.swapaxes(L, -1, -2)) * (sy[:, None] * sy[None, :])
return coef, intercept, cov
# ---------------------------------------------------------------------------
# Leaf polytopes
# ---------------------------------------------------------------------------
[docs]
def leaf_polytopes(
model,
bounds=None,
*,
prune_empty: bool = True,
X=None,
prune_unreached: bool = False,
tol: float = 1e-7,
) -> list[LeafPolytope]:
"""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``.
Returns
-------
list of LeafPolytope, in BFS leaf order.
"""
_check_model(model)
_require_tree_form(model, "leaf_polytopes")
return _leaf_polytopes(model, bounds, prune_empty, X, prune_unreached, tol)
def _leaf_polytopes(model, bounds, prune_empty, X, prune_unreached, tol):
if prune_unreached and X is None:
raise ValueError("prune_unreached=True requires X.")
lo, hi = _resolve_bounds(model, bounds)
depth = model.depth
p = model.n_features_in_
a, t = _splits_original(model)
coef, intercept, cov = _leaves_original(model)
counts = None
if X is not None:
counts = np.bincount(model.apply(X), minlength=_model.n_leaves(depth))
# LPs run in standardized units for conditioning; G, h are reported in
# original units. With xs = (x - mx) / sx, G_s xs <= h_s is equivalent
# to (G_s / sx) x <= h_s + (G_s / sx) . mx.
sx, mx = model.x_std_, model.x_mean_
W = np.asarray(model.params_["W"], dtype=np.float64)
b = np.asarray(model.params_["b"], dtype=np.float64)
eye = np.eye(p)
box_G = np.vstack([eye, -eye])
box_h = np.concatenate([hi, -lo])
box_Gs = box_G * sx[None, :]
box_hs = box_h - box_G @ mx
out = []
for leaf in range(_model.n_leaves(depth)):
nodes, sign, Gp, hp = _geometry.path_constraints(W, b, depth, leaf)
path = _geometry.leaf_path(depth, leaf)
_, radius = _geometry.chebyshev(np.vstack([Gp, box_Gs]),
np.concatenate([hp, box_hs]))
n_samples = None if counts is None else int(counts[leaf])
if prune_empty and radius <= tol:
continue
if prune_unreached and n_samples == 0:
continue
out.append(LeafPolytope(
leaf=leaf,
path=[(int(n), d) for n, d in path],
G=np.vstack([sign[:, None] * a[nodes], box_G]),
h=np.concatenate([sign * t[nodes], box_h]),
coef=coef[leaf],
intercept=intercept[leaf],
cov=cov[leaf],
radius=radius,
n_samples=n_samples,
))
if not out:
raise ValueError("Every leaf was pruned; check bounds and X.")
return out
# ---------------------------------------------------------------------------
# jax-ldt v1 JSON spec
# ---------------------------------------------------------------------------
[docs]
def to_spec(
model,
bounds=None,
*,
prune_empty: bool = True,
X=None,
prune_unreached: bool = False,
tol: float = 1e-7,
) -> dict:
"""Serialize the crisp tree to jax-ldt's solver-neutral v1 spec (a dict).
Arguments are as in :func:`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.
"""
_check_model(model)
_require_tree_form(model, "to_spec")
leaves = _leaf_polytopes(model, bounds, prune_empty, X, prune_unreached, tol)
lo, hi = _resolve_bounds(model, bounds)
kept = {lp.leaf: lp for lp in leaves}
depth = model.depth
n_int = _model.n_internal_nodes(depth)
p = model.n_features_in_
d_y = model.n_outputs_
a, t = _splits_original(model)
# has_kept[node]: does the subtree under ``node`` hold a kept leaf?
n_nodes = n_int + _model.n_leaves(depth)
has_kept = np.zeros(n_nodes, dtype=bool)
has_kept[n_int + np.array(sorted(kept))] = True
for node in range(n_int - 1, -1, -1):
has_kept[node] = has_kept[2 * node + 1] or has_kept[2 * node + 2]
def effective(node):
"""Skip internal nodes with a pruned side."""
while node < n_int:
left, right = 2 * node + 1, 2 * node + 2
if has_kept[left] and has_kept[right]:
break
node = left if has_kept[left] else right
return node
# Emit nodes in preorder; children ids are filled in after emission.
records = [] # (bfs node, left slot, right slot)
lifted = [] # BFS ids of kept internal nodes, in lifted-column order
def emit(node):
node = effective(node)
rid = len(records)
records.append([node, -1, -1])
if node < n_int:
lifted.append(node)
records[rid][1] = emit(2 * node + 1)
records[rid][2] = emit(2 * node + 2)
return rid
emit(0)
n_t = p + len(lifted)
col = {node: p + j for j, node in enumerate(lifted)}
A = np.hstack([np.eye(p), a[lifted].T]) if lifted else np.eye(p)
is_leaf, feature, threshold, left, right = [], [], [], [], []
leaf_params = np.zeros((len(records), 1 + n_t, d_y))
leaf_index, leaf_cov = [], []
for rid, (node, l_id, r_id) in enumerate(records):
if node < n_int:
is_leaf.append(False)
feature.append(col[node])
threshold.append(float(t[node]))
left.append(l_id)
right.append(r_id)
leaf_index.append(-1)
leaf_cov.append(None)
else:
lp = kept[node - n_int]
is_leaf.append(True)
feature.append(-1)
threshold.append(None)
left.append(-1)
right.append(-1)
leaf_params[rid, 0] = lp.intercept
leaf_params[rid, 1 : 1 + p] = lp.coef.T
leaf_index.append(lp.leaf)
leaf_cov.append(lp.cov.tolist())
from . import __version__
return {
"version": SPEC_VERSION,
"kind": "hyperplane_tree",
"n_features_in": int(p),
"n_features_transformed": int(n_t),
"n_targets": int(d_y),
"transform_matrix": A.tolist(),
"categorical_features": [],
"linear_features": list(range(p)),
"tree": {
"node_count": len(records),
"is_leaf": is_leaf,
"feature": feature,
"threshold": threshold,
"left": left,
"right": right,
"leaf_params": leaf_params.tolist(),
},
"difftree": {
"version": __version__,
"hardened": bool(getattr(model, "hardened_", False)),
"bounds": [lo.tolist(), hi.tolist()],
"leaf_index": leaf_index,
"leaf_cov": leaf_cov,
},
}
[docs]
def to_json(model, path: str | Path | None = None, **kwargs) -> str:
""":func:`to_spec` as a JSON string, also written to ``path`` if given."""
s = json.dumps(to_spec(model, **kwargs), indent=2, allow_nan=False)
if path is not None:
Path(path).write_text(s)
return s
[docs]
def predict_spec(spec: dict | str, X) -> np.ndarray:
"""Evaluate a v1 spec (dict or JSON string) at ``X`` with NumPy.
Returns (n,) for single-target specs, else (n, n_targets).
"""
if isinstance(spec, str):
spec = json.loads(spec)
X = np.atleast_2d(np.asarray(X, dtype=np.float64))
T = spec["tree"]
Xt = X @ np.asarray(spec["transform_matrix"], dtype=np.float64)
is_leaf = np.asarray(T["is_leaf"], dtype=bool)
feature = np.asarray(T["feature"], dtype=int)
thr = np.array([np.nan if v is None else v for v in T["threshold"]])
left = np.asarray(T["left"], dtype=int)
right = np.asarray(T["right"], dtype=int)
node = np.zeros(X.shape[0], dtype=int)
rows = np.arange(X.shape[0])
while not is_leaf[node].all():
active = ~is_leaf[node]
nd = node[active]
go_left = Xt[rows[active], feature[nd]] <= thr[nd]
node[active] = np.where(go_left, left[nd], right[nd])
params = np.asarray(T["leaf_params"], dtype=np.float64)[node] # (n, 1+n_t, T)
X_aug = np.hstack([np.ones((X.shape[0], 1)), Xt])
y = np.einsum("nf,nft->nt", X_aug, params)
return y[:, 0] if spec["n_targets"] == 1 else y
# ---------------------------------------------------------------------------
# Split boundaries: jumps and MILP solution repair
# ---------------------------------------------------------------------------
def _std_box(model, bounds):
lo, hi = _resolve_bounds(model, bounds)
return (lo - model.x_mean_) / model.x_std_, (hi - model.x_mean_) / model.x_std_
def _kept_leaves(model, lo_s, hi_s, tol):
W = np.asarray(model.params_["W"], dtype=np.float64)
b = np.asarray(model.params_["b"], dtype=np.float64)
box_G, box_h = _geometry.box_constraints(lo_s, hi_s)
kept = []
for leaf in range(_model.n_leaves(model.depth)):
_, _, Gp, hp = _geometry.path_constraints(W, b, model.depth, leaf)
_, r = _geometry.chebyshev(np.vstack([Gp, box_G]), np.concatenate([hp, box_h]))
if r > tol:
kept.append(leaf)
return kept
[docs]
def boundary_jumps(model, bounds=None, *, tol: float = 1e-7) -> list[dict]:
"""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.
"""
check_is_fitted(model, "params_")
return _boundary_jumps(model, *_std_box(model, bounds), tol)
def _boundary_jumps(model, lo_s, hi_s, tol=1e-7):
P = model.params_
W = np.asarray(P["W"], dtype=np.float64)
b = np.asarray(P["b"], dtype=np.float64)
facets = _geometry.shared_facets(W, b, model.depth, lo_s, hi_s,
leaves=_kept_leaves(model, lo_s, hi_s, tol), tol=tol)
out = []
for f in facets:
jump = np.zeros(model.n_outputs_)
worst = np.tile(f.center, (model.n_outputs_, 1))
for k in range(model.n_outputs_):
c0 = P["leaf_mu"][f.left, k] - P["leaf_mu"][f.right, k]
c = P["leaf_A"][f.left, k] - P["leaf_A"][f.right, k]
(lo_v, lo_x), (hi_v, hi_x) = _geometry.affine_range(
c0, c, f.G, f.h, W[f.node][None, :], np.array([-b[f.node]])
)
v, xk = (lo_v, lo_x) if abs(lo_v) >= abs(hi_v) else (hi_v, hi_x)
jump[k] = abs(v) * model.y_std_[k]
if xk is not None:
worst[k] = xk
out.append({
"node": f.node,
"left": f.left,
"right": f.right,
"point": f.center * model.x_std_ + model.x_mean_,
"max_jump": jump,
"worst_point": worst * model.x_std_ + model.x_mean_,
})
out.sort(key=lambda d: -float(d["max_jump"].max()))
return out
[docs]
def leaf_from_binaries(spec: dict | str, z) -> int:
"""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.
"""
if isinstance(spec, str):
spec = json.loads(spec)
leaf_nodes = np.flatnonzero(spec["tree"]["is_leaf"])
z = np.asarray(z, dtype=np.float64).ravel()
if z.shape[0] != leaf_nodes.shape[0]:
raise ValueError(f"z has {z.shape[0]} entries; spec has {leaf_nodes.shape[0]} leaves.")
return int(spec["difftree"]["leaf_index"][leaf_nodes[int(np.argmax(z))]])
[docs]
def repair_solution(model, x, leaf: int, *, delta: float = 1e-6, bounds=None) -> np.ndarray:
"""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 :func:`leaf_from_binaries`).
delta : float, default 1e-6
bounds : (lo, hi), optional
Input box in original units; defaults to the training-data box.
Returns
-------
(p,) ndarray in original units.
"""
check_is_fitted(model, "params_")
x = np.asarray(x, dtype=np.float64).ravel()
p = model.n_features_in_
W = np.asarray(model.params_["W"], dtype=np.float64)
b = np.asarray(model.params_["b"], dtype=np.float64)
lo_s, hi_s = _std_box(model, bounds)
_, _, Gp, hp = _geometry.path_constraints(W, b, model.depth, int(leaf))
box_G, box_h = _geometry.box_constraints(lo_s, hi_s)
G = np.vstack([Gp, box_G])
h = np.concatenate([hp, box_h]) - delta * np.linalg.norm(G, axis=1)
xs0 = (x - model.x_mean_) / model.x_std_
# variables [xs, s]: min s s.t. G xs <= h, |xs - xs0|_inf <= s
eye = np.eye(p)
A_ub = np.block([
[G, np.zeros((G.shape[0], 1))],
[eye, -np.ones((p, 1))],
[-eye, -np.ones((p, 1))],
])
b_ub = np.concatenate([h, xs0, -xs0])
c = np.zeros(p + 1)
c[-1] = 1.0
res = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=[(None, None)] * (p + 1), method="highs")
if res.status != 0:
raise ValueError(
f"Leaf {leaf} has no point {delta} inside all of its faces within the box."
)
return res.x[:p] * model.x_std_ + model.x_mean_
# ---------------------------------------------------------------------------
# Pyomo
# ---------------------------------------------------------------------------
def _var_list(v, n, what):
"""A Pyomo indexed Var (in index order) or a sequence of scalar vars."""
items = list(v.values()) if hasattr(v, "values") else list(v)
if len(items) != n:
raise ValueError(f"{what} has {len(items)} variables; the model needs {n}.")
return items
def _tighten(var, lo, hi):
if var.lb is None or var.lb < lo:
var.setlb(float(lo))
if var.ub is None or var.ub > hi:
var.setub(float(hi))
def _output_range(lp):
"""Exact range (d_y, 2) of each output of ``lp``'s model over its polytope."""
d_y = lp.coef.shape[0]
out = np.empty((d_y, 2))
free = [(None, None)] * lp.G.shape[1]
for k in range(d_y):
c = lp.coef[k]
if not np.any(c):
out[k] = lp.intercept[k]
continue
for j, s in enumerate((1.0, -1.0)):
res = linprog(s * c, A_ub=lp.G, b_ub=lp.h, bounds=free, method="highs")
if res.status != 0:
raise ValueError(f"Output range LP failed on leaf {lp.leaf}: {res.message}")
out[k, j] = s * res.fun + lp.intercept[k]
return out
[docs]
def to_pyomo(
model,
m=None,
x=None,
y=None,
bounds=None,
*,
formulation: str = "auto",
K: int | None = None,
method: str | None = "bigm",
name: str = "difftree",
prune_empty: bool = True,
X=None,
prune_unreached: bool = False,
tol: float = 1e-7,
):
"""Embed a crisp tree in a Pyomo model.
``formulation`` chooses the form:
- ``"tree"``: a disjunction over the leaves, described below. Exact for
any piecewise-affine (C0 or discontinuous) model.
- ``"pwl"``: the hinge form of a tied, continuous tree, ``y = theta_0 .
[1, x] + sum_d c_d max(0, h_d(x))``, with one big-M binary per hinge
term that changes sign inside the box (about ``depth`` binaries
instead of ``2**depth`` leaves). Exact for a C0 model. For a C1 model
each smoothed hinge is replaced by its ``K``-segment piecewise-linear
interpolant across the band (``K + 1`` hinge terms per depth and
output), which overestimates the C1 model where ``c_d > 0`` and
underestimates it where ``c_d < 0``, by at most
:func:`pwl_error_bound` everywhere.
- ``"smooth"``: the hinge form as an NLP with no binaries, each smoothed
hinge written with ``Expr_if``. C1 for a C1 model, so suitable for
IPOPT and similar solvers (for a C0 model it is only C0).
- ``"auto"`` (default): ``"tree"`` for piecewise-affine models. A C1
model has no exact MILP form, so ``"auto"`` raises and asks for
``"pwl"`` or ``"smooth"``.
``"pwl"`` and ``"smooth"`` need every output in hinge form
(``tied_splits=True`` and ``continuous`` covering all outputs).
The ``"tree"`` formulation:
Leaf ``l`` becomes a ``Disjunct`` holding its path half-spaces (``G_l x
<= h_l`` without the box rows) and its affine model ``y = coef_l x +
intercept_l``; a ``Disjunction`` makes exactly one leaf active. The input
box becomes bounds on ``x``, and each output gets the exact bounds of its
leaf models over their polytopes (one LP per leaf, output and direction),
so a big-M or hull reformulation needs no user-supplied constants. Bounds
already on ``x`` or ``y`` are only ever tightened.
Requires ``pyomo`` (``pip install difftree[pyomo]``).
Parameters
----------
model : fitted DifferentiableTreeRegressor (ideally hardened).
m : pyomo Block, optional
Model to add the tree to; a new ``ConcreteModel`` by default.
x : pyomo indexed Var or sequence of p vars, optional
Inputs, in feature order. Created on the tree block if omitted.
y : pyomo indexed Var or sequence of d_y vars, optional
Outputs, in output order. Created on the tree block if omitted.
bounds : (lo, hi), optional
Input box in original units; defaults to the training-data box.
formulation : {"auto", "tree", "pwl", "smooth"}, default "auto"
See above.
K : int, optional
``"pwl"`` with a C1 model: segments per band (required). Ignored for
a C0 model, whose hinge form is exact.
method : {"bigm", "hull", None}, default "bigm"
``"tree"`` only: Pyomo GDP transformation applied to the tree block
(``gdp.bigm`` or ``gdp.hull``). ``None`` leaves the disjunction
untransformed, e.g. to transform the whole model yourself or use a
GDP solver. ``"pwl"`` always uses big-M and ``"smooth"`` has no
binaries, so both ignore it.
name : str, default "difftree"
Name of the tree block on ``m``.
prune_empty, X, prune_unreached, tol
Passed to :func:`leaf_polytopes`.
Returns
-------
pyomo Block ``b`` (also ``getattr(m, name)``; ``b.model()`` is ``m``)
with ``b.x`` / ``b.y`` (references to the input and output vars),
``b.leaf`` (Disjuncts indexed by BFS leaf index, as returned by
``model.apply``) and ``b.choose`` (the Disjunction). After a solve,
:func:`leaf_from_pyomo` gives the active leaf, to pass to
:func:`repair_solution`. ``b.formulation`` records the formulation used.
A ``"pwl"`` block also has ``b.z`` (hinge term values), ``b.on`` (their
binaries), ``b.K`` and ``b.error_bound`` (the dict returned by
:func:`pwl_error_bound`).
"""
try:
import pyomo.environ as pyo
from pyomo.gdp import Disjunct, Disjunction
except ImportError as e: # pragma: no cover
raise ImportError("to_pyomo requires pyomo: pip install difftree[pyomo]") from e
if method not in ("bigm", "hull", None):
raise ValueError(f"method must be 'bigm', 'hull' or None, got {method!r}.")
if formulation not in ("auto", "tree", "pwl", "smooth"):
raise ValueError(
f"formulation must be 'auto', 'tree', 'pwl' or 'smooth', got {formulation!r}."
)
_check_model(model)
if formulation in ("pwl", "smooth"):
return _hinge_to_pyomo(model, m, x, y, bounds, formulation, K, name)
_require_tree_form(model, "to_pyomo" if formulation == "auto" else "to_pyomo(formulation='tree')")
polys = _leaf_polytopes(model, bounds, prune_empty, X, prune_unreached, tol)
lo, hi = _resolve_bounds(model, bounds)
p, d_y = model.n_features_in_, model.n_outputs_
if m is None:
m = pyo.ConcreteModel()
b = pyo.Block()
m.add_component(name, b)
if x is None:
b.x_vars = pyo.Var(range(p))
x = b.x_vars
if y is None:
b.y_vars = pyo.Var(range(d_y))
y = b.y_vars
xs = _var_list(x, p, "x")
ys = _var_list(y, d_y, "y")
b.x = pyo.Reference(xs)
b.y = pyo.Reference(ys)
for j in range(p):
_tighten(xs[j], lo[j], hi[j])
ranges = np.stack([_output_range(lp) for lp in polys]) # (L, d_y, 2)
for k in range(d_y):
_tighten(ys[k], ranges[:, k, 0].min(), ranges[:, k, 1].max())
by_leaf = {lp.leaf: lp for lp in polys}
def _leaf_rule(d, leaf):
lp = by_leaf[leaf]
n_path = lp.G.shape[0] - 2 * p # drop the box rows: they are x bounds
d.region = pyo.Constraint(range(n_path), rule=lambda _, r: (
sum(float(lp.G[r, j]) * xs[j] for j in range(p)) <= float(lp.h[r])))
d.output = pyo.Constraint(range(d_y), rule=lambda _, k: (
ys[k] == sum(float(lp.coef[k, j]) * xs[j] for j in range(p) if lp.coef[k, j] != 0)
+ float(lp.intercept[k])))
b.leaf = Disjunct(list(by_leaf), rule=_leaf_rule)
b.choose = Disjunction(expr=[b.leaf[leaf] for leaf in by_leaf])
if method is not None:
pyo.TransformationFactory(f"gdp.{method}").apply_to(m, targets=[b])
b.formulation = "tree"
return b
def _hinge_form(model, what):
"""The model's hinge form, checking that it covers every output."""
h = getattr(model, "hinge_", None)
if h is None or len(h["outputs"]) != model.n_outputs_:
raise ValueError(
f"{what} needs every output in hinge form: fit with tied_splits=True "
"and continuous=True, and harden."
)
return h
[docs]
def pwl_error_bound(model, K: int | None = None) -> dict:
"""Guaranteed error of the ``K``-segment piecewise-linear hinge model.
For a C1 model (``smoothness=1``), replacing each smoothed hinge by its
interpolant through ``K + 1`` equally spaced points across the band
changes output ``k`` by an amount in ``[-under[k], over[k]]`` at every
``x``, where ``over = sum_d max(c_d, 0) eps_d / (4 K^2)`` and ``under =
sum_d max(-c_d, 0) eps_d / (4 K^2)`` in original units. The bound
follows from the interpolation error ``phi'' L^2 / 8`` of a quadratic on
a segment of length ``L = 2 eps / K``, and is attained at segment
midpoints. For a C0 model both are zero and ``K`` is ignored.
Returns a dict with ``over``, ``under`` and ``bound = max(over, under)``,
each of shape (d_y,).
"""
check_is_fitted(model, "params_")
h = _hinge_form(model, "pwl_error_bound")
zero = np.zeros(model.n_outputs_)
if h["eps"] is None:
return {"over": zero, "under": zero.copy(), "bound": zero.copy()}
if K is None or int(K) != K or K < 1:
raise ValueError(f"K must be a positive integer for a C1 model, got {K!r}.")
p = model.n_features_in_
c = h["coef"][p + 1 :] # (D, m)
scale = model.y_std_[h["outputs"]] / (4 * int(K) ** 2)
over, under = zero.copy(), zero.copy()
over[h["outputs"]] = (np.maximum(c, 0) * h["eps"]).sum(axis=0) * scale
under[h["outputs"]] = (np.maximum(-c, 0) * h["eps"]).sum(axis=0) * scale
return {"over": over, "under": under, "bound": np.maximum(over, under)}
def _hinge_to_pyomo(model, m, x, y, bounds, formulation, K, name):
"""``"pwl"`` and ``"smooth"`` formulations of :func:`to_pyomo`."""
import pyomo.environ as pyo
from . import _hinge
h = _hinge_form(model, f"to_pyomo(formulation={formulation!r})")
smooth = h["eps"] is not None
if formulation == "pwl" and smooth:
bound = pwl_error_bound(model, K) # validates K
elif formulation == "pwl":
bound = pwl_error_bound(model)
lo, hi = _resolve_bounds(model, bounds)
lo_s, hi_s = (lo - model.x_mean_) / model.x_std_, (hi - model.x_mean_) / model.x_std_
p, d_y = model.n_features_in_, model.n_outputs_
W, bh, coef = h["W"], h["b"], h["coef"]
D, q = W.shape[0], p + 1
if m is None:
m = pyo.ConcreteModel()
b = pyo.Block()
m.add_component(name, b)
if x is None:
b.x_vars = pyo.Var(range(p))
x = b.x_vars
if y is None:
b.y_vars = pyo.Var(range(d_y))
y = b.y_vars
xv = _var_list(x, p, "x")
yv = _var_list(y, d_y, "y")
b.x = pyo.Reference(xv)
b.y = pyo.Reference(yv)
for j in range(p):
_tighten(xv[j], lo[j], hi[j])
xs = [(xv[j] - float(model.x_mean_[j])) / float(model.x_std_[j]) for j in range(p)]
hexpr = [sum(float(W[d, j]) * xs[j] for j in range(p)) + float(bh[d]) for d in range(D)]
# Range of h_d over the box (interval arithmetic in standardized units).
h_lo = bh + np.minimum(W * lo_s, W * hi_s).sum(axis=1)
h_hi = bh + np.maximum(W * lo_s, W * hi_s).sum(axis=1)
# Hinge terms max(0, h_d - t) * a, keyed per output o (C1 bands differ
# per output) or shared by all outputs (C0).
terms = [] # (o or None, d, t, a)
if formulation == "pwl":
if smooth:
t_all, a = _hinge.pwl_terms(h["eps"], int(K)) # (D, m, K + 1), (K + 1,)
for o in range(len(h["outputs"])):
for d in range(D):
terms += [(o, d, float(t_all[d, o, k]), float(a[k])) for k in range(int(K) + 1)]
else:
terms = [(None, d, 0.0, 1.0) for d in range(D)]
b.T = pyo.Set(initialize=range(len(terms)))
b.z = pyo.Var(b.T, within=pyo.NonNegativeReals)
b.on = pyo.Var(b.T, within=pyo.Binary)
b.hinge = pyo.ConstraintList()
for i, (_, d, t, _) in enumerate(terms):
L, U = h_lo[d] - t, h_hi[d] - t
g = hexpr[d] - t
b.z[i].setub(float(max(U, 0.0)))
if U <= 0:
b.z[i].fix(0.0)
b.on[i].fix(0)
elif L >= 0:
b.hinge.add(b.z[i] == g)
b.on[i].fix(1)
else:
b.hinge.add(b.z[i] >= g)
b.hinge.add(b.z[i] <= g - float(L) * (1 - b.on[i]))
b.hinge.add(b.z[i] <= float(U) * b.on[i])
def hinge_sum(col):
"""sum_d c_d * (hinge d) for output column ``col`` of the hinge form."""
if formulation == "pwl":
return sum(float(coef[q + d, col]) * a * b.z[i]
for i, (oo, d, _, a) in enumerate(terms) if oo is None or oo == col)
out = 0.0
for d in range(D):
hd = hexpr[d]
if smooth:
e = float(h["eps"][d, col])
ph = pyo.Expr_if(IF=hd <= -e, THEN=0.0,
ELSE=pyo.Expr_if(IF=hd >= e, THEN=hd, ELSE=(hd + e) ** 2 / (4 * e)))
else:
ph = pyo.Expr_if(IF=hd >= 0, THEN=hd, ELSE=0.0)
out = out + float(coef[q + d, col]) * ph
return out
b.output = pyo.ConstraintList()
for col, o in enumerate(h["outputs"]):
lin = float(coef[0, col]) + sum(float(coef[1 + j, col]) * xs[j] for j in range(p))
b.output.add(yv[o] == float(model.y_mean_[o])
+ float(model.y_std_[o]) * (lin + hinge_sum(col)))
b.formulation = formulation
b.K = int(K) if formulation == "pwl" and smooth else None
if formulation == "pwl":
b.error_bound = bound
b._difftree_model = model
return b
[docs]
def leaf_from_pyomo(block) -> int:
"""BFS leaf index of the active leaf of a solved :func:`to_pyomo` block.
For the ``"pwl"`` and ``"smooth"`` formulations there are no leaf
binaries; the leaf is the crisp routing of the solved ``x``.
"""
import pyomo.environ as pyo
if getattr(block, "formulation", "tree") != "tree":
xv = np.array([[pyo.value(v) for v in block.x.values()]])
return int(block._difftree_model.apply(xv)[0])
vals = {leaf: pyo.value(d.binary_indicator_var) for leaf, d in block.leaf.items()}
return int(max(vals, key=vals.get))