From 87daf57b8910617c990f2a1fbaeb8c7f1dbbe9fa Mon Sep 17 00:00:00 2001 From: J Xu Date: Mon, 17 Aug 2026 09:39:17 -0700 Subject: [PATCH 1/3] feat: unified boosting trainer with FormulaBoost and WeibullAFT Collapse the distributional and standard boosting loops into one multi-channel trainer (fit_boosting) driven by an Objective protocol, so a model is a configuration rather than a separate training fork. - FormulaObjective / FormulaBoost: boost the parameters of an arbitrary differentiable formula f(theta(z), x) with damped generalized Gauss-Newton preconditioning. The off-diagonal GGN term is what a black-box GBDT and an XGBoost custom objective's diagonal-only Hessian cannot express. - WeibullAFTObjective / WeibullAFT: right-censored survival regression that boosts BOTH Weibull scale(z) and shape(z) via an expected-Fisher natural gradient. XGBoost's survival:aft holds the distribution scale as a single global hyperparameter and cannot vary the shape with covariates. - DistributionObjective: NaturalBoost / DistributionalGBDT rebuilt on the shared trainer and the old per-model loop deleted, with device-resident Normal/Poisson steps. GPU path keeps raw scores device-resident when the objective is device-capable and builds one tree per channel with fit_tree_gpu_native; the formula and survival objectives build trees on GPU with host-side preconditioning. Co-authored-by: Cursor --- src/openboost/__init__.py | 4 + src/openboost/_backends/_cuda.py | 128 +++++ src/openboost/_models/__init__.py | 4 + src/openboost/_models/_distributional.py | 258 ++------- src/openboost/_models/_formula.py | 189 +++++++ src/openboost/_models/_survival.py | 190 +++++++ src/openboost/_objectives.py | 637 +++++++++++++++++++++++ src/openboost/_trainer.py | 303 +++++++++++ 8 files changed, 1509 insertions(+), 204 deletions(-) create mode 100644 src/openboost/_models/_formula.py create mode 100644 src/openboost/_models/_survival.py create mode 100644 src/openboost/_objectives.py create mode 100644 src/openboost/_trainer.py diff --git a/src/openboost/__init__.py b/src/openboost/__init__.py index 950a945..472bba7 100644 --- a/src/openboost/__init__.py +++ b/src/openboost/__init__.py @@ -183,6 +183,7 @@ DART, # Phase 15/16: Distributional GBDT (NaturalBoost) DistributionalGBDT, + FormulaBoost, GradientBoosting, # Phase 15: Linear Leaf GBDT LinearLeafGBDT, @@ -204,6 +205,7 @@ OpenBoostLinearLeafRegressor, # Phase 13: sklearn-compatible wrappers OpenBoostRegressor, + WeibullAFT, ) from ._models import ( # Backward compatibility aliases (deprecated, accessed via __getattr__) @@ -346,6 +348,8 @@ def __getattr__(name: str): "NaturalBoostStudentT", "NaturalBoostTweedie", "NaturalBoostNegBin", + "FormulaBoost", + "WeibullAFT", "LegacyTree", # Backward compatibility (deprecated) "NGBoost", diff --git a/src/openboost/_backends/_cuda.py b/src/openboost/_backends/_cuda.py index 4d862fb..5ca1d42 100644 --- a/src/openboost/_backends/_cuda.py +++ b/src/openboost/_backends/_cuda.py @@ -3584,3 +3584,131 @@ def build_tree_symmetric_gpu_native( # Already on GPU - no transfer needed! return level_features_gpu, level_thresholds_gpu, leaf_values + + +# ============================================================================= +# Multi-parameter objective kernels (unified trainer, K>1) +# ============================================================================= + +def _blocks_threads(n: int, threads: int = 256) -> tuple[int, int]: + return (n + threads - 1) // threads, threads + + +@cuda.jit +def _normal_ordinary_kernel(raw_loc, raw_scale, y, g_loc, h_loc, g_scale, h_scale, n): + """Ordinary NLL grad/hess for Normal: loc identity, scale exp-link.""" + i = cuda.grid(1) + if i >= n: + return + mu = raw_loc[i] + s = raw_scale[i] + if s > 20.0: + s = 20.0 + elif s < -20.0: + s = -20.0 + sigma = math.exp(s) + var = sigma * sigma + if var < 1e-12: + var = 1e-12 + resid = y[i] - mu + g_loc[i] = -resid / var + h_loc[i] = 1.0 / var + g_scale[i] = 1.0 - (resid * resid) / var + h_scale[i] = 2.0 + + +@cuda.jit +def _normal_natural_kernel(raw_loc, raw_scale, y, g_loc, h_loc, g_scale, h_scale, n): + """Natural gradient for Normal (diagonal Fisher): F=diag(1/σ², 2).""" + i = cuda.grid(1) + if i >= n: + return + mu = raw_loc[i] + s = raw_scale[i] + if s > 20.0: + s = 20.0 + elif s < -20.0: + s = -20.0 + sigma = math.exp(s) + var = sigma * sigma + if var < 1e-12: + var = 1e-12 + resid = y[i] - mu + g_loc[i] = -resid + h_loc[i] = 1.0 + g_scale[i] = 0.5 * (1.0 - (resid * resid) / var) + h_scale[i] = 1.0 + + +@cuda.jit +def _poisson_ordinary_kernel(raw_rate, y, g, h, n): + """Ordinary NLL grad/hess for Poisson with exp link: λ=exp(raw).""" + i = cuda.grid(1) + if i >= n: + return + r = raw_rate[i] + if r > 20.0: + r = 20.0 + elif r < -20.0: + r = -20.0 + lam = math.exp(r) + g[i] = lam - y[i] + h[i] = lam if lam > 1e-6 else 1e-6 + + +@cuda.jit +def _poisson_natural_kernel(raw_rate, y, g, h, n): + """Natural gradient for Poisson: F=λ, nat=(λ-y)/λ, unit hessian.""" + i = cuda.grid(1) + if i >= n: + return + r = raw_rate[i] + if r > 20.0: + r = 20.0 + elif r < -20.0: + r = -20.0 + lam = math.exp(r) + if lam < 1e-12: + lam = 1e-12 + g[i] = 1.0 - y[i] / lam + h[i] = 1.0 + + +@cuda.jit +def _scale_gh_kernel(grad, hess, weight, n): + i = cuda.grid(1) + if i < n: + w = weight[i] + grad[i] *= w + hess[i] *= w + + +def normal_step_gpu(raw_loc, raw_scale, y, *, natural: bool): + """In-place Normal (grad, hess) on device. Returns four device arrays.""" + n = int(y.shape[0]) + g_loc = cuda.device_array(n, dtype=np.float32) + h_loc = cuda.device_array(n, dtype=np.float32) + g_scale = cuda.device_array(n, dtype=np.float32) + h_scale = cuda.device_array(n, dtype=np.float32) + blocks, threads = _blocks_threads(n) + kern = _normal_natural_kernel if natural else _normal_ordinary_kernel + kern[blocks, threads](raw_loc, raw_scale, y, g_loc, h_loc, g_scale, h_scale, n) + return g_loc, h_loc, g_scale, h_scale + + +def poisson_step_gpu(raw_rate, y, *, natural: bool): + """In-place Poisson (grad, hess) on device.""" + n = int(y.shape[0]) + g = cuda.device_array(n, dtype=np.float32) + h = cuda.device_array(n, dtype=np.float32) + blocks, threads = _blocks_threads(n) + kern = _poisson_natural_kernel if natural else _poisson_ordinary_kernel + kern[blocks, threads](raw_rate, y, g, h, n) + return g, h + + +def scale_gh_gpu(grad, hess, weight) -> None: + """Multiply grad and hess by per-sample weights, in place.""" + n = int(grad.shape[0]) + blocks, threads = _blocks_threads(n) + _scale_gh_kernel[blocks, threads](grad, hess, weight, n) diff --git a/src/openboost/_models/__init__.py b/src/openboost/_models/__init__.py index 810829a..03e3b49 100644 --- a/src/openboost/_models/__init__.py +++ b/src/openboost/_models/__init__.py @@ -33,6 +33,7 @@ NGBoostStudentT, NGBoostTweedie, ) +from ._formula import FormulaBoost from ._gam import OpenBoostGAM # Phase 15: Linear Leaf GBDT @@ -45,6 +46,7 @@ OpenBoostLinearLeafRegressor, OpenBoostRegressor, ) +from ._survival import WeibullAFT __all__ = [ # Standard GBDT @@ -82,4 +84,6 @@ # Phase 15: Linear Leaf GBDT "LinearLeafGBDT", "LinearLeafTree", + "FormulaBoost", + "WeibullAFT", ] diff --git a/src/openboost/_models/_distributional.py b/src/openboost/_models/_distributional.py index a4151e0..120a8de 100644 --- a/src/openboost/_models/_distributional.py +++ b/src/openboost/_models/_distributional.py @@ -40,25 +40,20 @@ import numpy as np -from .._array import BinnedArray, array -from .._callbacks import ( - Callback, - CallbackManager, - EarlyStopping, - TrainingState, - warn_if_early_stopping_without_eval_set, -) +from .._array import BinnedArray +from .._callbacks import Callback from .._core._growth import TreeStructure -from .._core._tree import fit_tree from .._distributions import ( Distribution, DistributionOutput, Normal, get_distribution, ) +from .._objectives import DistributionObjective from .._persistence import PersistenceMixin +from .._trainer import TrainerConfig, fit_boosting, predict_raw from .._utils import crps_empirical, crps_gaussian, interval_score, pinball_loss -from .._validation import validate_eval_set, validate_sample_weight +from .._validation import validate_sample_weight if TYPE_CHECKING: from numpy.typing import NDArray @@ -184,7 +179,8 @@ class DistributionalGBDT(PersistenceMixin): X_binned_: BinnedArray | None = field(default=None, init=False, repr=False) _base_scores: dict[str, float] = field(default_factory=dict, init=False, repr=False) n_features_in_: int = field(default=0, init=False, repr=False) - + _use_natural_gradient: bool = field(default=False, init=False, repr=False) + def fit( self, X: NDArray, @@ -254,11 +250,8 @@ def fit( f"Available: {', '.join(EVAL_METRICS)}." ) - # Split optional (X, y, exposure) eval entries before validation eval_pairs, eval_exposures = _split_eval_exposures(eval_set) - # Resolve which raw parameter carries the log-exposure offset - # (raises ValueError for families without a log-link mean) exposure_param: str | None = None exposure_sign = 0.0 needs_exposure = exposure is not None or any( @@ -272,172 +265,56 @@ def fit( exposure = _validate_exposure(exposure, n_samples, context="fit") train_log_offset = (exposure_sign * np.log(exposure)).astype(np.float32) - # Bin features - if isinstance(X, BinnedArray): - self.X_binned_ = X - else: - self.X_binned_ = array(X, n_bins=self.n_bins) - - self.n_features_in_ = self.X_binned_.n_features - - # Initialize tree storage - self.trees_ = {} - for param_name in self.distribution_.param_names: - self.trees_[param_name] = [] - - # Initialize raw predictions (in link space) using data statistics - init_params = self.distribution_.init_params(y) - raw_preds = {} - - for param_name in self.distribution_.param_names: - raw_init = init_params[param_name] - self._base_scores[param_name] = float(raw_init) - raw_preds[param_name] = np.full(n_samples, raw_init, dtype=np.float32) - - # Setup callbacks (early_stopping_rounds is sugar for EarlyStopping) - cb_list = list(callbacks) if callbacks else [] - if early_stopping_rounds is not None: - cb_list.append( - EarlyStopping(patience=early_stopping_rounds, restore_best=True) - ) - cb_manager = CallbackManager(cb_list) - state = TrainingState(model=self, n_rounds=self.n_trees) - cb_manager.on_train_begin(state) - - eval_pairs = validate_eval_set(eval_pairs, self.X_binned_.n_features) - warn_if_early_stopping_without_eval_set(cb_list, eval_pairs) + objective = DistributionObjective( + self.distribution_, + natural=self._use_natural_gradient, + exposure_param=exposure_param, + exposure_sign=exposure_sign, + ) - # Per-eval-set state: binned features, targets, incrementally - # maintained raw scores, and optional log-exposure offset - eval_data = [] + eval_sets = None if eval_pairs: + eval_sets = [] for (X_e, y_e), exp_e in zip(eval_pairs, eval_exposures, strict=True): - X_e_binned = ( - X_e if isinstance(X_e, BinnedArray) - else self.X_binned_.transform(X_e) - ) - raw_e = { - p: np.full( - X_e_binned.n_samples, self._base_scores[p], dtype=np.float32 - ) - for p in self.distribution_.param_names - } - log_off_e = None + extra_e: dict = {} if exp_e is not None: - exp_e = _validate_exposure( - exp_e, X_e_binned.n_samples, context="eval" - ) - log_off_e = (exposure_sign * np.log(exp_e)).astype(np.float32) - eval_data.append((X_e_binned, y_e, raw_e, log_off_e)) - - self.evals_result_ = { - f'eval_{i}': {eval_metric: []} for i in range(len(eval_data)) - } - - # Training loop - for _round_idx in range(self.n_trees): - # Constrained params from raw scores (+ exposure offset) - params = self._constrained_params( - raw_preds, exposure_param, train_log_offset + y_e_arr = np.asarray(y_e).ravel() + exp_e = _validate_exposure(exp_e, len(y_e_arr), context="eval") + extra_e["log_offset"] = ( + exposure_sign * np.log(exp_e) + ).astype(np.float32) + eval_sets.append({"X": X_e, "y": y_e, "extra": extra_e}) + + def _score_eval(y_e, raw_e, extra_e): + params_e = objective.constrain(raw_e, extra_e) + return self._eval_metric_value( + y_e, params_e, eval_metric, quantiles, interval_alpha ) - # Get gradients for each parameter - grads_dict = self._compute_gradients(y, params) - - if sample_weight is not None: - # Weighted likelihood: the objective is sum_i w_i * NLL_i. - # - # Ordinary-gradient path (DistributionalGBDT): grad_i/hess_i - # are per-sample derivatives of NLL_i, so both scale linearly - # in w_i; the Newton leaf value -Σ w_i g_i / (Σ w_i h_i + λ) - # is then the correct weighted step. - # - # Natural-gradient path (NaturalBoost): under the weighted - # likelihood the per-sample Fisher information also scales by - # w_i, so the per-sample natural gradient - # (w_i F_i)^{-1} (w_i g_i) = F_i^{-1} g_i is weight-INVARIANT. - # The correct weighted natural-gradient aggregate is obtained - # by scaling the per-sample natural gradient (returned by - # _compute_gradients with unit hessians) by w_i, and scaling - # its unit hessian by w_i so leaf aggregation becomes the - # weighted mean -Σ w_i g̃_i / (Σ w_i + λ). Post-scaling both - # grad and hess by w_i therefore covers both paths. - grads_dict = { - p: ( - (g * sample_weight).astype(np.float32), - (h * sample_weight).astype(np.float32), - ) - for p, (g, h) in grads_dict.items() - } - - # Train one tree per parameter - for param_name in self.distribution_.param_names: - grad, hess = grads_dict[param_name] - - tree = fit_tree( - self.X_binned_, - grad, - hess, - max_depth=self.max_depth, - min_child_weight=self.min_child_weight, - reg_lambda=self.reg_lambda, - reg_alpha=self.reg_alpha, - subsample=self.subsample, - colsample_bytree=self.colsample_bytree, - ) - - self.trees_[param_name].append(tree) - - # Update raw predictions (add tree prediction, as in standard GBDT) - # Tree is trained on gradients, so it outputs the negative gradient direction - tree_pred = tree(self.X_binned_) - if hasattr(tree_pred, 'copy_to_host'): - tree_pred = tree_pred.copy_to_host() - - raw_preds[param_name] += self.learning_rate * tree_pred - - # Keep eval-set raw scores in sync (incremental, avoids - # re-running all trees each round) - for X_e_binned, _y_e, raw_e, _log_off_e in eval_data: - pred_e = tree(X_e_binned) - if hasattr(pred_e, 'copy_to_host'): - pred_e = pred_e.copy_to_host() - raw_e[param_name] += self.learning_rate * pred_e - - # Evaluate ALL eval sets and record per-round history - last_metric = None - for i, (_X_e_binned, y_e, raw_e, log_off_e) in enumerate(eval_data): - params_e = self._constrained_params( - raw_e, exposure_param, log_off_e - ) - last_metric = self._eval_metric_value( - y_e, params_e, eval_metric, quantiles, interval_alpha - ) - self.evals_result_[f'eval_{i}'][eval_metric].append(last_metric) - - # Callbacks - state.round_idx = _round_idx - if cb_manager.callbacks: - # Train loss: (weighted) mean NLL of the post-round params, - # matching the timing of the eval-set metrics. Recomputed from - # raw_preds because identity-link params alias raw_preds and - # were mutated in place by the tree updates above. - params_report = self._constrained_params( - raw_preds, exposure_param, train_log_offset - ) - state.train_loss = float( - np.average( - self.distribution_.nll(y, params_report), - weights=sample_weight, - ) - ) - if last_metric is not None: - # Early stopping monitors the LAST eval set's metric - state.val_loss = last_metric - if not cb_manager.on_round_end(state): - break - - cb_manager.on_train_end(state) + fit_boosting( + self, + objective, + X, + y, + config=TrainerConfig( + n_trees=self.n_trees, + max_depth=self.max_depth, + learning_rate=self.learning_rate, + min_child_weight=self.min_child_weight, + reg_lambda=self.reg_lambda, + reg_alpha=self.reg_alpha, + subsample=self.subsample, + colsample_bytree=self.colsample_bytree, + n_bins=self.n_bins, + ), + sample_weight=sample_weight, + extra={"log_offset": train_log_offset}, + callbacks=callbacks, + early_stopping_rounds=early_stopping_rounds, + eval_sets=eval_sets, + eval_fn=_score_eval if eval_sets else None, + eval_metric_name=eval_metric, + ) return self def _resolve_exposure_offset(self) -> tuple[str, float]: @@ -550,35 +427,7 @@ def _predict_raw(self, X: NDArray | BinnedArray) -> dict[str, NDArray]: Returns: Dictionary mapping param_name -> raw predictions """ - if not self.trees_: - raise RuntimeError("Model not fitted. Call fit() first.") - - # Bin the data if needed, using training bin edges for consistency - if isinstance(X, BinnedArray): - X_binned = X - elif self.X_binned_ is not None: - # Use transform to apply training bin edges to new data - X_binned = self.X_binned_.transform(X) - else: - X_binned = array(X, n_bins=self.n_bins) - - n_samples = X_binned.n_samples - raw_preds = {} - - for param_name in self.distribution_.param_names: - # Start with base score - pred = np.full(n_samples, self._base_scores[param_name], dtype=np.float32) - - # Accumulate tree predictions - for tree in self.trees_[param_name]: - tree_pred = tree(X_binned) - if hasattr(tree_pred, 'copy_to_host'): - tree_pred = tree_pred.copy_to_host() - pred += self.learning_rate * tree_pred - - raw_preds[param_name] = pred - - return raw_preds + return predict_raw(self, X) def predict_params( self, @@ -792,7 +641,8 @@ class NaturalBoost(DistributionalGBDT): # Override defaults for NaturalBoost max_depth: int = 4 # Shallower trees often work better learning_rate: float = 0.1 - + _use_natural_gradient: bool = field(default=True, init=False, repr=False) + def _compute_gradients( self, y: NDArray, diff --git a/src/openboost/_models/_formula.py b/src/openboost/_models/_formula.py new file mode 100644 index 0000000..af31c26 --- /dev/null +++ b/src/openboost/_models/_formula.py @@ -0,0 +1,189 @@ +"""FormulaBoost: boost every parameter of a user formula. + +Varying-coefficient / semi-parametric boosting. Features ``Z`` determine +parameter surfaces ``theta(Z)`` via trees; a user formula +``y ≈ f(theta, x)`` consumes those parameters and a structural input ``x``. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +from .._array import BinnedArray +from .._callbacks import Callback +from .._core._growth import TreeStructure +from .._objectives import FormulaObjective +from .._persistence import PersistenceMixin +from .._trainer import TrainerConfig, fit_boosting, predict_raw + + +def _as_1d(a, n: int, name: str) -> NDArray: + arr = np.asarray(a, dtype=np.float64).ravel() + if arr.shape[0] != n: + raise ValueError( + f"{name} has length {arr.shape[0]}, expected {n} (matching y)." + ) + return arr + + +@dataclass +class FormulaBoost(PersistenceMixin): + """Boost the parameters of an arbitrary differentiable formula. + + Args: + formula: ``formula(theta, x) -> yhat``. ``theta`` is a tuple of K + arrays in constrained parameter space; ``x`` is the structural + input (e.g. spend, dose). + n_params: Number of formula parameters K. + links: Per-parameter link (``identity``, ``log``, ``softplus``, + ``sigmoid``), length K. + loss: Training loss. Currently ``mse`` only. + precond: GGN preconditioner: ``full`` (default), ``diag``, or + ``plain`` (raw gradient — usually a bad idea). + damp: Levenberg–Marquardt damping added to the GGN matrix. + param_names: Optional names for the K parameters. Defaults to + ``theta_0``, ``theta_1``, ... + n_trees, max_depth, learning_rate, ...: standard tree knobs. + + Example: + ```python + def curve(theta, x): + a, b = theta + return a * x ** (1.0 / (1.0 + np.exp(-b * x))) + + m = FormulaBoost( + formula=curve, n_params=2, links=("log", "identity"), + param_names=("a", "b"), + ) + m.fit(Z, y, model_input=x) + params = m.predict_params(Z) # per-sample (a, b) + yhat = m.predict(Z, model_input=x_new) + ``` + """ + + formula: Callable + n_params: int + links: tuple[str, ...] + loss: str = "mse" + precond: str = "full" + damp: float = 1.0 + param_names: tuple[str, ...] | None = None + n_trees: int = 100 + max_depth: int = 3 + learning_rate: float = 0.1 + min_child_weight: float = 1.0 + reg_lambda: float = 1.0 + reg_alpha: float = 0.0 + subsample: float = 1.0 + colsample_bytree: float = 1.0 + n_bins: int = 254 + + trees_: dict[str, list[TreeStructure]] = field( + default_factory=dict, init=False, repr=False + ) + evals_result_: dict[str, dict[str, list[float]]] = field( + default_factory=dict, init=False, repr=False + ) + X_binned_: BinnedArray | None = field(default=None, init=False, repr=False) + _base_scores: dict[str, float] = field(default_factory=dict, init=False, repr=False) + n_features_in_: int = field(default=0, init=False, repr=False) + _objective: FormulaObjective | None = field(default=None, init=False, repr=False) + + def _make_objective(self) -> FormulaObjective: + return FormulaObjective( + self.formula, + self.n_params, + self.links, + loss=self.loss, + precond=self.precond, + damp=self.damp, + param_names=self.param_names, + ) + + def fit( + self, + X: NDArray, + y: NDArray, + model_input: NDArray, + sample_weight: NDArray | None = None, + callbacks: list[Callback] | None = None, + eval_set: list[tuple] | None = None, + early_stopping_rounds: int | None = None, + ) -> FormulaBoost: + """Fit parameter surfaces ``theta(Z)`` from ``(X, y, model_input)``. + + ``eval_set`` entries are ``(X_val, y_val, model_input_val)``. + """ + y = np.asarray(y, dtype=np.float64).ravel() + x = _as_1d(model_input, len(y), "model_input") + self._objective = self._make_objective() + + eval_sets: list[dict[str, Any]] | None = None + if eval_set is not None: + if ( + isinstance(eval_set, tuple) + and len(eval_set) == 3 + and not isinstance(eval_set[0], tuple) + ): + eval_set = [eval_set] + eval_sets = [] + for item in eval_set: + if not (isinstance(item, tuple) and len(item) == 3): + raise ValueError( + "FormulaBoost eval_set entries must be " + "(X_val, y_val, model_input_val)." + ) + X_e, y_e, x_e = item + y_e = np.asarray(y_e, dtype=np.float64).ravel() + eval_sets.append( + { + "X": X_e, + "y": y_e, + "extra": {"model_input": _as_1d(x_e, len(y_e), "model_input")}, + } + ) + + fit_boosting( + self, + self._objective, + X, + y, + config=TrainerConfig( + n_trees=self.n_trees, + max_depth=self.max_depth, + learning_rate=self.learning_rate, + min_child_weight=self.min_child_weight, + reg_lambda=self.reg_lambda, + reg_alpha=self.reg_alpha, + subsample=self.subsample, + colsample_bytree=self.colsample_bytree, + n_bins=self.n_bins, + ), + sample_weight=sample_weight, + extra={"model_input": x}, + callbacks=callbacks, + early_stopping_rounds=early_stopping_rounds, + eval_sets=eval_sets, + eval_metric_name="mse", + ) + return self + + def predict_params(self, X: NDArray | BinnedArray) -> dict[str, NDArray]: + """Predict constrained formula parameters for each row of ``X``.""" + if self._objective is None: + self._objective = self._make_objective() + raw = predict_raw(self, X) + return self._objective.constrain(raw) + + def predict(self, X: NDArray | BinnedArray, model_input: NDArray) -> NDArray: + """Evaluate the formula at ``predict_params(X)`` and ``model_input``.""" + params = self.predict_params(X) + names = self._objective.channel_names + x = _as_1d(model_input, next(iter(params.values())).shape[0], "model_input") + theta = tuple(params[name] for name in names) + return np.asarray(self.formula(theta, x), dtype=np.float64).ravel() diff --git a/src/openboost/_models/_survival.py b/src/openboost/_models/_survival.py new file mode 100644 index 0000000..d337ff2 --- /dev/null +++ b/src/openboost/_models/_survival.py @@ -0,0 +1,190 @@ +"""WeibullAFT: boosted Weibull survival regression with censoring. + +Both the scale ``lambda(z)`` and the shape ``k(z)`` are boosting ensembles over +the features ``z`` and trained on a right-censored negative log-likelihood. +XGBoost's ``survival:aft`` learns only the location and holds the distribution +scale as a single global hyperparameter, so it cannot vary the Weibull shape +with covariates; NaturalBoost's built-in distributions have no censored +likelihood at all. WeibullAFT expresses both. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +from .._array import BinnedArray +from .._callbacks import Callback +from .._core._growth import TreeStructure +from .._objectives import WeibullAFTObjective +from .._persistence import PersistenceMixin +from .._trainer import TrainerConfig, fit_boosting, predict_raw + + +def _as_1d(a, n: int, name: str) -> NDArray: + arr = np.asarray(a, dtype=np.float64).ravel() + if arr.shape[0] != n: + raise ValueError( + f"{name} has length {arr.shape[0]}, expected {n} (matching y)." + ) + return arr + + +@dataclass +class WeibullAFT(PersistenceMixin): + """Boosted Weibull accelerated-failure-time model with right censoring. + + Args: + damp: Levenberg-Marquardt damping on the per-sample 2x2 information. + n_trees, max_depth, learning_rate, ...: standard tree knobs. + + Example: + ```python + m = WeibullAFT(n_trees=300, max_depth=3, learning_rate=0.1) + m.fit(Z, time, event=observed) # event: 1 seen, 0 censored + params = m.predict_params(Z) # per-sample {scale, shape} + t_hat = m.predict(Z) # predicted median time + s = m.predict_survival(Z, t=5.0) # S(5.0 | z) per sample + ``` + """ + + damp: float = 1.0 + n_trees: int = 100 + max_depth: int = 3 + learning_rate: float = 0.1 + min_child_weight: float = 1.0 + reg_lambda: float = 1.0 + reg_alpha: float = 0.0 + subsample: float = 1.0 + colsample_bytree: float = 1.0 + n_bins: int = 254 + + trees_: dict[str, list[TreeStructure]] = field( + default_factory=dict, init=False, repr=False + ) + evals_result_: dict[str, dict[str, list[float]]] = field( + default_factory=dict, init=False, repr=False + ) + X_binned_: BinnedArray | None = field(default=None, init=False, repr=False) + _base_scores: dict[str, float] = field(default_factory=dict, init=False, repr=False) + n_features_in_: int = field(default=0, init=False, repr=False) + _objective: WeibullAFTObjective | None = field( + default=None, init=False, repr=False + ) + + def _make_objective(self) -> WeibullAFTObjective: + return WeibullAFTObjective(damp=self.damp) + + def fit( + self, + X: NDArray, + y: NDArray, + event: NDArray | None = None, + sample_weight: NDArray | None = None, + callbacks: list[Callback] | None = None, + eval_set: list[tuple] | None = None, + early_stopping_rounds: int | None = None, + ) -> WeibullAFT: + """Fit ``lambda(z)``, ``k(z)`` from ``(X, time, event)``. + + ``y`` is the observed time (>0). ``event`` is 1 for an observed event + and 0 for right-censored (default: all observed). ``eval_set`` entries + are ``(X_val, y_val, event_val)``. + """ + y = np.asarray(y, dtype=np.float64).ravel() + if np.any(y <= 0): + raise ValueError("WeibullAFT requires strictly positive times y.") + ev = None if event is None else _as_1d(event, len(y), "event") + self._objective = self._make_objective() + + eval_sets: list[dict[str, Any]] | None = None + if eval_set is not None: + if ( + isinstance(eval_set, tuple) + and len(eval_set) in (2, 3) + and not isinstance(eval_set[0], tuple) + ): + eval_set = [eval_set] + eval_sets = [] + for item in eval_set: + if not (isinstance(item, tuple) and len(item) in (2, 3)): + raise ValueError( + "WeibullAFT eval_set entries must be " + "(X_val, y_val[, event_val])." + ) + X_e, y_e = item[0], np.asarray(item[1], dtype=np.float64).ravel() + ev_e = (_as_1d(item[2], len(y_e), "event") + if len(item) == 3 else None) + eval_sets.append( + {"X": X_e, "y": y_e, "extra": {"event": ev_e}} + ) + + fit_boosting( + self, + self._objective, + X, + y, + config=TrainerConfig( + n_trees=self.n_trees, + max_depth=self.max_depth, + learning_rate=self.learning_rate, + min_child_weight=self.min_child_weight, + reg_lambda=self.reg_lambda, + reg_alpha=self.reg_alpha, + subsample=self.subsample, + colsample_bytree=self.colsample_bytree, + n_bins=self.n_bins, + ), + sample_weight=sample_weight, + extra={"event": ev}, + callbacks=callbacks, + early_stopping_rounds=early_stopping_rounds, + eval_sets=eval_sets, + eval_metric_name="nll", + ) + return self + + def predict_params(self, X: NDArray | BinnedArray) -> dict[str, NDArray]: + """Per-sample constrained Weibull parameters ``{scale, shape}``.""" + if self._objective is None: + self._objective = self._make_objective() + return self._objective.constrain(predict_raw(self, X)) + + def predict_quantile(self, X: NDArray | BinnedArray, q: float = 0.5) -> NDArray: + """Predict the time ``t`` at which ``P(T <= t) = q`` for each row.""" + if not 0.0 < q < 1.0: + raise ValueError("q must be in (0, 1).") + params = self.predict_params(X) + lam, k = params["scale"], params["shape"] + return lam * (-np.log(1.0 - q)) ** (1.0 / k) + + def predict_median(self, X: NDArray | BinnedArray) -> NDArray: + return self.predict_quantile(X, 0.5) + + def predict(self, X: NDArray | BinnedArray) -> NDArray: + """Point prediction = predicted median survival time.""" + return self.predict_median(X) + + def predict_survival(self, X: NDArray | BinnedArray, t: float | NDArray) -> NDArray: + """Survival probability ``S(t | z) = exp(-(t / lambda)^k)`` per row.""" + params = self.predict_params(X) + lam, k = params["scale"], params["shape"] + t = np.asarray(t, dtype=np.float64) + return np.exp(-((t / lam) ** k)) + + def nll( + self, + X: NDArray | BinnedArray, + y: NDArray, + event: NDArray | None = None, + ) -> float: + """Mean censored negative log-likelihood on ``(X, y, event)``.""" + if self._objective is None: + self._objective = self._make_objective() + y = np.asarray(y, dtype=np.float64).ravel() + ev = None if event is None else _as_1d(event, len(y), "event") + raw = predict_raw(self, X) + return self._objective.loss_value(raw, y, None, {"event": ev}) diff --git a/src/openboost/_objectives.py b/src/openboost/_objectives.py new file mode 100644 index 0000000..1ee716b --- /dev/null +++ b/src/openboost/_objectives.py @@ -0,0 +1,637 @@ +"""Objectives for the unified boosting trainer. + +An objective turns current raw scores F (one channel per boosted parameter) +into per-sample step directions (grad, hess) that ``fit_tree`` consumes. + +Two first-class implementations: + +- ``DistributionObjective`` — NLL / natural-gradient path used by + DistributionalGBDT and NaturalBoost. +- ``FormulaObjective`` — user formula ``f(theta, x)`` with damped generalized + Gauss-Newton preconditioning (the FormulaBoost path). +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Protocol + +import numpy as np +from numpy.typing import NDArray + +from ._distributions import Distribution + +RawScores = dict[str, NDArray] +GradHess = dict[str, tuple[NDArray, NDArray]] + + +class Objective(Protocol): + """Internal protocol consumed by ``fit_boosting``.""" + + channel_names: list[str] + + def init_raw( + self, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> dict[str, float]: + """Return per-channel base scores in raw (unconstrained) space.""" + ... + + def step( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> GradHess: + """Per-channel (gradient, hessian) w.r.t. raw scores.""" + ... + + def loss_value( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> float: + """Scalar loss (lower is better) for train reporting / default eval.""" + ... + + def constrain( + self, + raw: RawScores, + extra: dict[str, Any] | None = None, + ) -> dict[str, NDArray]: + """Map raw scores to constrained parameters.""" + ... + + +def _is_device(arr) -> bool: + return hasattr(arr, "__cuda_array_interface__") + + +def _apply_sample_weight(grads: GradHess, sample_weight: NDArray | None) -> GradHess: + if sample_weight is None: + return grads + w = sample_weight.astype(np.float32, copy=False) + return { + name: ((g * w).astype(np.float32), (h * w).astype(np.float32)) + for name, (g, h) in grads.items() + } + + +# ============================================================================= +# Distribution objective (NaturalBoost / DistributionalGBDT) +# ============================================================================= + + +class DistributionObjective: + """Boost each distribution parameter from an NLL (optionally natural).""" + + def __init__( + self, + distribution: Distribution, + *, + natural: bool = False, + exposure_param: str | None = None, + exposure_sign: float = 0.0, + ): + self.distribution = distribution + self.natural = natural + self.exposure_param = exposure_param + self.exposure_sign = exposure_sign + self.channel_names = list(distribution.param_names) + self._device_kernels_ok = True + + @property + def device_capable(self) -> bool: + """True when grad/hess can be computed on device-resident raw scores.""" + return type(self.distribution).__name__ in ("Normal", "Poisson") + + @property + def unit_hessian(self) -> bool: + return bool(self.natural) + + def init_raw( + self, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> dict[str, float]: + return {k: float(v) for k, v in self.distribution.init_params(y).items()} + + def constrain( + self, + raw: RawScores, + extra: dict[str, Any] | None = None, + ) -> dict[str, NDArray]: + log_offset = None if extra is None else extra.get("log_offset") + params = {} + for name in self.channel_names: + score = raw[name] + if log_offset is not None and name == self.exposure_param: + score = score + log_offset + params[name] = self.distribution.link(name, score) + return params + + def step( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> GradHess: + first = next(iter(raw.values())) + if ( + self.device_capable + and self._device_kernels_ok + and _is_device(first) + and (extra is None or extra.get("log_offset") is None) + ): + return self._step_device(raw, y, sample_weight) + params = self.constrain(raw, extra) + if self.natural: + grads = self.distribution.natural_gradient(y, params) + else: + grads = self.distribution.nll_gradient(y, params) + return _apply_sample_weight(grads, sample_weight) + + def _step_device( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None, + ) -> GradHess: + try: + from ._backends._cuda import normal_step_gpu, poisson_step_gpu, scale_gh_gpu + + name = type(self.distribution).__name__ + if name == "Normal": + g_loc, h_loc, g_scale, h_scale = normal_step_gpu( + raw["loc"], raw["scale"], y, natural=self.natural + ) + grads: GradHess = {"loc": (g_loc, h_loc), "scale": (g_scale, h_scale)} + elif name == "Poisson": + g, h = poisson_step_gpu(raw["rate"], y, natural=self.natural) + grads = {"rate": (g, h)} + else: # pragma: no cover + raise RuntimeError(f"No device step for {name}") + if sample_weight is not None: + for g, h in grads.values(): + scale_gh_gpu(g, h, sample_weight) + return grads + except Exception: + self._device_kernels_ok = False + # Kernel compile can fail on some numpy/numba-cuda combos; + # fall back to the host path and let the trainer upload grads. + raw_host = { + k: (v.copy_to_host() if hasattr(v, "copy_to_host") else np.asarray(v)) + for k, v in raw.items() + } + y_host = y.copy_to_host() if hasattr(y, "copy_to_host") else np.asarray(y) + sw_host = None + if sample_weight is not None: + sw_host = ( + sample_weight.copy_to_host() + if hasattr(sample_weight, "copy_to_host") + else np.asarray(sample_weight) + ) + params = self.constrain(raw_host) + if self.natural: + grads = self.distribution.natural_gradient(y_host, params) + else: + grads = self.distribution.nll_gradient(y_host, params) + return _apply_sample_weight(grads, sw_host) + + def loss_value( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> float: + params = self.constrain(raw, extra) + nll = self.distribution.nll(y, params) + return float(np.average(nll, weights=sample_weight)) + + +# ============================================================================= +# Formula objective (FormulaBoost) +# ============================================================================= + +# link(raw) -> constrained; link_prime(raw) -> d(constrained)/d(raw) +_LINKS: dict[str, tuple[Callable[[NDArray], NDArray], Callable[[NDArray], NDArray]]] = { + "identity": (lambda r: r, lambda r: np.ones_like(r, dtype=np.float64)), + "log": (np.exp, np.exp), + "softplus": ( + lambda r: np.logaddexp(0.0, r), + lambda r: 1.0 / (1.0 + np.exp(-r)), + ), + "sigmoid": ( + lambda r: 1.0 / (1.0 + np.exp(-r)), + lambda r: (s := 1.0 / (1.0 + np.exp(-r))) * (1.0 - s), + ), +} + + +def _link_pair(name: str): + if name not in _LINKS: + raise ValueError( + f"Unknown link '{name}'. Available: {', '.join(sorted(_LINKS))}." + ) + return _LINKS[name] + + +def _formula_and_jac( + formula: Callable, + theta: list[NDArray], + x: NDArray, + eps: float = 1e-5, +) -> tuple[NDArray, NDArray]: + """Evaluate ``formula(theta, x)`` and a forward-difference Jacobian. + + Returns ``(f, J)`` with ``f`` shape ``(n,)`` and ``J`` shape ``(n, K)`` + where ``J[:, k] = df / d theta_k``. + """ + f0 = np.asarray(formula(tuple(theta), x), dtype=np.float64).ravel() + k = len(theta) + jac = np.empty((f0.shape[0], k), dtype=np.float64) + for j in range(k): + bumped = list(theta) + bumped[j] = theta[j] + eps + f1 = np.asarray(formula(tuple(bumped), x), dtype=np.float64).ravel() + jac[:, j] = (f1 - f0) / eps + return f0, jac + + +def _ggn_step( + residual: NDArray, + jac_raw: NDArray, + precond: str, + damp: float, +) -> NDArray: + """Per-sample GGN step directions, shape ``(n, K)``. + + MSE / 0.5*(f-y)^2 => g = residual * J, G = J^T J (per sample). + """ + g = residual[:, None] * jac_raw # (n, K) + if precond == "plain": + return g + if precond == "diag": + return g / (jac_raw * jac_raw + damp) + if precond != "full": + raise ValueError( + f"Unknown precond '{precond}'. Use 'plain', 'diag', or 'full'." + ) + + n, k = jac_raw.shape + if k == 1: + return g / (jac_raw * jac_raw + damp) + if k == 2: + g00 = jac_raw[:, 0] * jac_raw[:, 0] + damp + g11 = jac_raw[:, 1] * jac_raw[:, 1] + damp + g01 = jac_raw[:, 0] * jac_raw[:, 1] + det = g00 * g11 - g01 * g01 + det = np.where(np.abs(det) < 1e-12, 1e-12, det) + d0 = (g11 * g[:, 0] - g01 * g[:, 1]) / det + d1 = (g00 * g[:, 1] - g01 * g[:, 0]) / det + return np.stack([d0, d1], axis=1) + + # General K: batched (J J^T + damp I)^{-1} g via small dense solves + eye = np.eye(k, dtype=np.float64) + out = np.empty_like(g) + for i in range(n): + ji = jac_raw[i] + gmat = np.outer(ji, ji) + damp * eye + try: + out[i] = np.linalg.solve(gmat, g[i]) + except np.linalg.LinAlgError: + out[i] = np.linalg.solve(gmat + 1e-6 * eye, g[i]) + return out + + +def _fit_global_raw( + formula: Callable, + links: list[str], + x: NDArray, + y: NDArray, + sample_weight: NDArray | None, +) -> np.ndarray: + """Fit a single global raw-parameter vector by L-BFGS-B.""" + from scipy.optimize import minimize + + k = len(links) + x = np.asarray(x, dtype=np.float64).ravel() + y = np.asarray(y, dtype=np.float64).ravel() + + v0 = np.zeros(k, dtype=np.float64) + if links[0] == "log" and np.all(y > 0): + v0[0] = float(np.log(np.average(y, weights=sample_weight))) + + def packed_loss(v: np.ndarray) -> float: + theta = [] + for j, name in enumerate(links): + fn, _ = _link_pair(name) + raw_j = np.full_like(x, v[j]) + theta.append(fn(raw_j)) + pred = np.asarray(formula(tuple(theta), x), dtype=np.float64).ravel() + if not np.all(np.isfinite(pred)): + return 1e12 + err = 0.5 * (pred - y) ** 2 + return float(np.average(err, weights=sample_weight)) + + result = minimize(packed_loss, v0, method="L-BFGS-B") + return result.x if result.success else v0 + + +class FormulaObjective: + """Boost the parameters of a user formula ``f(theta, x)``. + + ``formula(theta, x)`` receives ``theta`` as a tuple of K arrays (constrained + space) and ``x`` as the structural input (e.g. spend). Features ``Z`` that + the trees split on never enter the formula — they only determine + ``theta(Z)``. + """ + + def __init__( + self, + formula: Callable, + n_params: int, + links: tuple[str, ...] | list[str], + *, + loss: str = "mse", + precond: str = "full", + damp: float = 1.0, + param_names: tuple[str, ...] | list[str] | None = None, + fd_eps: float = 1e-5, + ): + if loss != "mse": + raise ValueError("FormulaObjective currently supports loss='mse' only.") + if len(links) != n_params: + raise ValueError( + f"links has length {len(links)}, expected n_params={n_params}." + ) + for name in links: + _link_pair(name) + if precond not in ("plain", "diag", "full"): + raise ValueError( + f"Unknown precond '{precond}'. Use 'plain', 'diag', or 'full'." + ) + + self.formula = formula + self.n_params = n_params + self.links = list(links) + self.loss = loss + self.precond = precond + self.damp = float(damp) + self.fd_eps = float(fd_eps) + if param_names is None: + self.channel_names = [f"theta_{j}" for j in range(n_params)] + else: + if len(param_names) != n_params: + raise ValueError("param_names length must equal n_params.") + self.channel_names = list(param_names) + + @property + def device_capable(self) -> bool: + # Formula + GGN stay on host; the trainer still builds trees on GPU. + return False + + @property + def unit_hessian(self) -> bool: + return True + + def _require_x(self, extra: dict[str, Any] | None) -> NDArray: + if extra is None or extra.get("model_input") is None: + raise ValueError( + "FormulaObjective requires extra['model_input'] " + "(the structural input x that enters the formula)." + ) + return np.asarray(extra["model_input"], dtype=np.float64).ravel() + + def init_raw( + self, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> dict[str, float]: + x = self._require_x(extra) + v = _fit_global_raw(self.formula, self.links, x, y, sample_weight) + return {name: float(v[j]) for j, name in enumerate(self.channel_names)} + + def constrain( + self, + raw: RawScores, + extra: dict[str, Any] | None = None, + ) -> dict[str, NDArray]: + params = {} + for name, link_name in zip(self.channel_names, self.links, strict=True): + fn, _ = _link_pair(link_name) + params[name] = fn(np.asarray(raw[name], dtype=np.float64)) + return params + + def _predict_f(self, raw: RawScores, extra: dict[str, Any] | None) -> NDArray: + x = self._require_x(extra) + params = self.constrain(raw, extra) + theta = [params[name] for name in self.channel_names] + return np.asarray(self.formula(tuple(theta), x), dtype=np.float64).ravel() + + def step( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> GradHess: + x = self._require_x(extra) + y = np.asarray(y, dtype=np.float64).ravel() + theta = [] + link_prime = [] + for name, link_name in zip(self.channel_names, self.links, strict=True): + fn, dfn = _link_pair(link_name) + r = np.asarray(raw[name], dtype=np.float64).ravel() + theta.append(fn(r)) + link_prime.append(dfn(r)) + + f, jac_theta = _formula_and_jac(self.formula, theta, x, eps=self.fd_eps) + jac_raw = jac_theta * np.stack(link_prime, axis=1) + residual = f - y + direction = _ggn_step(residual, jac_raw, self.precond, self.damp) + + n = y.shape[0] + ones = np.ones(n, dtype=np.float32) + grads: GradHess = {} + for j, name in enumerate(self.channel_names): + grads[name] = (direction[:, j].astype(np.float32), ones.copy()) + return _apply_sample_weight(grads, sample_weight) + + def loss_value( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> float: + pred = self._predict_f(raw, extra) + y = np.asarray(y, dtype=np.float64).ravel() + return float(np.average(0.5 * (pred - y) ** 2, weights=sample_weight)) + + +# ============================================================================= +# Weibull AFT objective (survival / censored regression) +# ============================================================================= + +# NGBoost's built-in families have no censored likelihood, and XGBoost's +# survival:aft objective learns only the location while holding the +# distribution scale (shape) as a single global hyperparameter. WeibullAFT +# boosts BOTH the scale lambda(z) and the shape k(z) surfaces from a censored +# NLL, which is the capability neither can express. +# +# Parameterization (both channels live in log space; that is the raw score): +# u = raw["scale"] = log(lambda), lambda = exp(u) +# w = raw["shape"] = log(k), k = exp(w) +# For time t with event indicator delta in {0, 1} (1 = observed, 0 = right +# censored), with m = log(t) - u and z = (t/lambda)^k = exp(k * m): +# NLL = -delta * (w - u + (k - 1) * m) + z +# (i.e. -delta * log hazard + cumulative hazard). + + +class WeibullAFTObjective: + """Boost Weibull scale ``lambda(z)`` and shape ``k(z)`` under a censored NLL. + + ``extra['event']`` is the event indicator (1 observed, 0 right-censored); + defaults to all-observed. Preconditioning is a damped, PD-safeguarded + per-sample 2x2 observed-information solve (Newton-natural step), sharing + the trainer's unit-hessian / GPU-tree path with FormulaBoost. + """ + + _W_CLIP = 12.0 # clip log-shape to keep k = exp(w) finite + _Z_CLIP = 60.0 # clip k*m before exp + _EULER = 0.5772156649015329 + # Expected Fisher info of the log-shape channel for an observed Weibull + # event (constant in the log parameterization): (1-gamma)^2 + pi^2/6. + _I_WW = (1.0 - _EULER) ** 2 + (np.pi ** 2) / 6.0 + + def __init__(self, *, damp: float = 1.0): + self.channel_names = ["scale", "shape"] + self.damp = float(damp) + + @property + def device_capable(self) -> bool: + return False + + @property + def unit_hessian(self) -> bool: + return True + + def _event(self, y: NDArray, extra: dict[str, Any] | None) -> NDArray: + n = len(y) + if extra is None or extra.get("event") is None: + return np.ones(n, dtype=np.float64) + e = np.asarray(extra["event"], dtype=np.float64).ravel() + if e.shape[0] != n: + raise ValueError( + f"event has length {e.shape[0]}, expected {n} (matching y)." + ) + return e + + def _pieces(self, u, w, t): + """Shared intermediates for grad/hess/nll.""" + u = np.asarray(u, dtype=np.float64).ravel() + w = np.clip(np.asarray(w, dtype=np.float64).ravel(), -self._W_CLIP, self._W_CLIP) + k = np.exp(w) + m = np.log(t) - u + z = np.exp(np.clip(k * m, -self._Z_CLIP, self._Z_CLIP)) + return k, m, z + + def init_raw( + self, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> dict[str, float]: + from scipy.optimize import minimize + + t = np.asarray(y, dtype=np.float64).ravel() + if np.any(t <= 0): + raise ValueError("WeibullAFT requires strictly positive times y.") + delta = self._event(t, extra) + + def mean_nll(v): + u, w = np.full_like(t, v[0]), np.full_like(t, v[1]) + k, m, z = self._pieces(u, w, t) + nll = -delta * (w - u + (k - 1.0) * m) + z + if not np.all(np.isfinite(nll)): + return 1e12 + return float(np.average(nll, weights=sample_weight)) + + v0 = np.array([float(np.log(np.median(t))), 0.0]) + res = minimize(mean_nll, v0, method="Nelder-Mead", + options={"xatol": 1e-4, "fatol": 1e-6, "maxiter": 400}) + v = res.x if res.success else v0 + return {"scale": float(v[0]), "shape": float(v[1])} + + def constrain( + self, + raw: RawScores, + extra: dict[str, Any] | None = None, + ) -> dict[str, NDArray]: + u = np.asarray(raw["scale"], dtype=np.float64).ravel() + w = np.clip(np.asarray(raw["shape"], dtype=np.float64).ravel(), + -self._W_CLIP, self._W_CLIP) + return {"scale": np.exp(u), "shape": np.exp(w)} + + def step( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> GradHess: + t = np.asarray(y, dtype=np.float64).ravel() + delta = self._event(t, extra) + u = np.asarray(raw["scale"], dtype=np.float64).ravel() + w = np.asarray(raw["shape"], dtype=np.float64).ravel() + k, m, z = self._pieces(u, w, t) + + # Gradient of the censored NLL w.r.t. raw (u = log-scale, w = log-shape). + g_u = k * (delta - z) + g_w = -delta * (1.0 + m * k) + k * m * z + + # Damped natural gradient: precondition with the EXPECTED Fisher + # information (not the observed Hessian, whose scale term k^2 z blows + # up when lambda is wrong and freezes the update). In the log + # parameterization the per-observed-event Fisher is + # [[k^2, -k(1-gamma)], [-k(1-gamma), (1-gamma)^2 + pi^2/6]]. + off = -k * (1.0 - self._EULER) + a = k * k + self.damp + c = self._I_WW + self.damp + b = off + det = a * c - b * b + + d_u = (c * g_u - b * g_w) / det + d_w = (a * g_w - b * g_u) / det + + ones = np.ones(len(t), dtype=np.float32) + grads: GradHess = { + "scale": (d_u.astype(np.float32), ones.copy()), + "shape": (d_w.astype(np.float32), ones.copy()), + } + return _apply_sample_weight(grads, sample_weight) + + def loss_value( + self, + raw: RawScores, + y: NDArray, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + ) -> float: + t = np.asarray(y, dtype=np.float64).ravel() + delta = self._event(t, extra) + u = np.asarray(raw["scale"], dtype=np.float64).ravel() + w = np.asarray(raw["shape"], dtype=np.float64).ravel() + k, m, z = self._pieces(u, w, t) + nll = -delta * (w - u + (k - 1.0) * m) + z + return float(np.average(nll, weights=sample_weight)) diff --git a/src/openboost/_trainer.py b/src/openboost/_trainer.py new file mode 100644 index 0000000..6039885 --- /dev/null +++ b/src/openboost/_trainer.py @@ -0,0 +1,303 @@ +"""Unified multi-channel boosting trainer. + +One loop for every model: bin once, maintain raw scores F (n, K), ask the +objective for per-channel (grad, hess), fit one tree per channel, update F +and every eval set incrementally. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +from ._array import BinnedArray, array +from ._backends import is_cuda +from ._callbacks import ( + Callback, + CallbackManager, + EarlyStopping, + TrainingState, + warn_if_early_stopping_without_eval_set, +) +from ._core._growth import TreeStructure +from ._core._tree import fit_tree, fit_tree_gpu_native +from ._objectives import Objective, RawScores +from ._validation import validate_eval_set, validate_sample_weight + + +@dataclass +class TrainerConfig: + """Tree-building knobs shared by every facade.""" + + n_trees: int = 100 + max_depth: int = 6 + learning_rate: float = 0.1 + min_child_weight: float = 1.0 + reg_lambda: float = 1.0 + reg_alpha: float = 0.0 + subsample: float = 1.0 + colsample_bytree: float = 1.0 + n_bins: int = 254 + + +def _to_host(pred: NDArray) -> NDArray: + if hasattr(pred, "copy_to_host"): + return pred.copy_to_host() + if hasattr(pred, "get"): + return pred.get() + return np.asarray(pred) + + +def _is_device(arr) -> bool: + return hasattr(arr, "__cuda_array_interface__") + + +def _as_host_raw(raw: RawScores) -> RawScores: + return {name: _to_host(score) for name, score in raw.items()} + + +def _gpu_native_eligible(X_binned: BinnedArray, config: TrainerConfig) -> bool: + if not is_cuda(): + return False + if config.reg_alpha != 0.0 or config.colsample_bytree < 1.0 or config.subsample < 1.0: + return False + has_missing = ( + hasattr(X_binned, "has_missing") + and len(X_binned.has_missing) > 0 + and np.any(X_binned.has_missing) + ) + has_categorical = ( + hasattr(X_binned, "is_categorical") + and len(X_binned.is_categorical) > 0 + and np.any(X_binned.is_categorical) + ) + return not has_missing and not has_categorical + + +def _legacy_to_structure(legacy, n_features: int, max_depth: int) -> TreeStructure: + features, thresholds, values, left, right = legacy.to_arrays() + return TreeStructure( + features=features, + thresholds=thresholds, + left_children=left, + right_children=right, + values=values, + n_nodes=len(features), + depth=max_depth, + n_features=n_features, + ) + + +def _empty_raw(n: int, base_scores: dict[str, float]) -> RawScores: + return { + name: np.full(n, score, dtype=np.float32) for name, score in base_scores.items() + } + + +def _bin_features(X, reference: BinnedArray | None, n_bins: int) -> BinnedArray: + if isinstance(X, BinnedArray): + return X + if reference is not None: + return reference.transform(X) + return array(X, n_bins=n_bins) + + +def predict_raw(model: Any, X) -> RawScores: + """Accumulate raw scores from ``model.trees_`` / ``_base_scores``.""" + if not getattr(model, "trees_", None): + raise RuntimeError("Model not fitted. Call fit() first.") + + X_binned = _bin_features(X, getattr(model, "X_binned_", None), model.n_bins) + n = X_binned.n_samples + raw = _empty_raw(n, model._base_scores) + lr = model.learning_rate + for name, trees in model.trees_.items(): + pred = raw[name] + for tree in trees: + pred = pred + lr * _to_host(tree(X_binned)) + raw[name] = pred + return raw + + +def fit_boosting( + model: Any, + objective: Objective, + X, + y: NDArray, + *, + config: TrainerConfig, + sample_weight: NDArray | None = None, + extra: dict[str, Any] | None = None, + callbacks: list[Callback] | None = None, + early_stopping_rounds: int | None = None, + eval_sets: list[dict[str, Any]] | None = None, + eval_fn: Callable[[NDArray, RawScores, dict[str, Any] | None], float] | None = None, + eval_metric_name: str = "loss", +) -> Any: + """Fit ``model`` in place. Returns ``model``. + + ``eval_sets`` entries are dicts with keys ``X``, ``y``, optional ``extra``, + optional ``name`` (defaults to ``eval_0``, ``eval_1``, ...). + """ + y = np.asarray(y).ravel() + n_samples = len(y) + sample_weight = validate_sample_weight(sample_weight, n_samples) + extra = extra or {} + + model.X_binned_ = _bin_features(X, None, config.n_bins) + model.n_features_in_ = model.X_binned_.n_features + model.learning_rate = config.learning_rate + model.n_bins = config.n_bins + + base = objective.init_raw(y, sample_weight, extra) + model._base_scores = dict(base) + model.trees_ = {name: [] for name in objective.channel_names} + + use_gpu = is_cuda() + device_state = ( + use_gpu + and bool(getattr(objective, "device_capable", False)) + and extra.get("log_offset") is None + ) + use_native = _gpu_native_eligible(model.X_binned_, config) + unit_hess = bool(getattr(objective, "unit_hessian", False)) + + if use_gpu: + from numba import cuda + + from ._core._predict import _add_inplace_cuda + + binned_gpu = model.X_binned_.data + if not _is_device(binned_gpu): + binned_gpu = cuda.to_device(binned_gpu) + else: + cuda = None # type: ignore[assignment] + _add_inplace_cuda = None # type: ignore[assignment] + binned_gpu = model.X_binned_.data + + if device_state: + raw = { + name: cuda.to_device(np.full(n_samples, score, dtype=np.float32)) + for name, score in model._base_scores.items() + } + y_step = cuda.to_device(np.ascontiguousarray(y, dtype=np.float32)) + sw_step = ( + cuda.to_device(np.ascontiguousarray(sample_weight, dtype=np.float32)) + if sample_weight is not None + else None + ) + else: + raw = _empty_raw(n_samples, model._base_scores) + y_step = y + sw_step = sample_weight + + cb_list = list(callbacks) if callbacks else [] + if early_stopping_rounds is not None: + cb_list.append(EarlyStopping(patience=early_stopping_rounds, restore_best=True)) + cb_manager = CallbackManager(cb_list) + state = TrainingState(model=model, n_rounds=config.n_trees) + cb_manager.on_train_begin(state) + + prepared_eval: list[tuple[str, BinnedArray, NDArray, dict[str, Any], RawScores]] = [] + if eval_sets: + pairs = [(item["X"], item["y"]) for item in eval_sets] + validate_eval_set(pairs, model.X_binned_.n_features) + for i, item in enumerate(eval_sets): + name = item.get("name", f"eval_{i}") + X_e = _bin_features(item["X"], model.X_binned_, config.n_bins) + y_e = np.asarray(item["y"]).ravel() + extra_e = item.get("extra") or {} + raw_e = _empty_raw(X_e.n_samples, model._base_scores) + prepared_eval.append((name, X_e, y_e, extra_e, raw_e)) + + warn_if_early_stopping_without_eval_set(cb_list, prepared_eval or None) + model.evals_result_ = {name: {eval_metric_name: []} for name, *_ in prepared_eval} + + score_eval = eval_fn or ( + lambda y_e, raw_e, extra_e: objective.loss_value(raw_e, y_e, None, extra_e) + ) + + for round_idx in range(config.n_trees): + if device_state: + grads = objective.step(raw, y_step, sw_step, extra) + else: + grads = objective.step(_as_host_raw(raw), y, sample_weight, extra) + + for name in objective.channel_names: + grad, hess = grads[name] + if use_gpu and not _is_device(grad): + grad = cuda.to_device(np.ascontiguousarray(grad, dtype=np.float32)) + hess = cuda.to_device(np.ascontiguousarray(hess, dtype=np.float32)) + elif not use_gpu: + grad = np.ascontiguousarray(grad, dtype=np.float32) + hess = np.ascontiguousarray(hess, dtype=np.float32) + + if use_native: + pred_buf = raw[name] if device_state else None + legacy = fit_tree_gpu_native( + binned_gpu, + grad, + hess, + max_depth=config.max_depth, + min_child_weight=config.min_child_weight, + reg_lambda=config.reg_lambda, + pred_gpu=pred_buf, + learning_rate=config.learning_rate, + const_hess=1.0 if unit_hess else 0.0, + ) + tree = _legacy_to_structure( + legacy, model.n_features_in_, config.max_depth + ) + if not device_state: + raw[name] = raw[name] + config.learning_rate * _to_host( + tree(model.X_binned_) + ) + else: + tree = fit_tree( + model.X_binned_, + grad, + hess, + max_depth=config.max_depth, + min_child_weight=config.min_child_weight, + reg_lambda=config.reg_lambda, + reg_alpha=config.reg_alpha, + subsample=config.subsample, + colsample_bytree=config.colsample_bytree, + ) + update = tree(model.X_binned_) + if device_state: + if not _is_device(update): + update = cuda.to_device( + np.ascontiguousarray(_to_host(update), dtype=np.float32) + ) + _add_inplace_cuda(raw[name], update, config.learning_rate) + else: + raw[name] = _to_host(raw[name]) + config.learning_rate * _to_host( + update + ) + + model.trees_[name].append(tree) + for _n, X_e, _y_e, _extra_e, raw_e in prepared_eval: + raw_e[name] = raw_e[name] + config.learning_rate * _to_host(tree(X_e)) + + last_metric = None + for name, _X_e, y_e, extra_e, raw_e in prepared_eval: + last_metric = float(score_eval(y_e, raw_e, extra_e)) + model.evals_result_[name][eval_metric_name].append(last_metric) + + state.round_idx = round_idx + if cb_manager.callbacks: + state.train_loss = objective.loss_value( + _as_host_raw(raw), y, sample_weight, extra + ) + if last_metric is not None: + state.val_loss = last_metric + if not cb_manager.on_round_end(state): + break + + cb_manager.on_train_end(state) + return model From bcbbdbce11ad69879a5ceaac73e165a98e450d49 Mon Sep 17 00:00:00 2001 From: J Xu Date: Mon, 17 Aug 2026 09:39:25 -0700 Subject: [PATCH 2/3] test: cover FormulaBoost, WeibullAFT, and GPU parity Add tests/test_formula.py (extrapolation beats a black-box GBDT, full preconditioning beats plain, parameter-surface recovery, eval_set early stopping) and tests/test_survival.py (recovers the shape surface, censoring is actually used, quantile and survival-curve monotonicity, input guards). Extend the CUDA verification and Modal GPU tests to exercise the unified trainer, FormulaBoost, and NaturalBoost CPU/GPU parity, and pin numpy<2.5 on the Modal image to match numba-cuda. Co-authored-by: Cursor --- tests/modal_gpu_tests.py | 95 ++++++++++++++- tests/test_cuda_verification.py | 45 ++++++++ tests/test_formula.py | 199 ++++++++++++++++++++++++++++++++ tests/test_survival.py | 129 +++++++++++++++++++++ 4 files changed, 467 insertions(+), 1 deletion(-) create mode 100644 tests/test_formula.py create mode 100644 tests/test_survival.py diff --git a/tests/modal_gpu_tests.py b/tests/modal_gpu_tests.py index a9985f8..c7e58e9 100644 --- a/tests/modal_gpu_tests.py +++ b/tests/modal_gpu_tests.py @@ -16,7 +16,7 @@ image = ( modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12") .pip_install( - "numpy>=1.24", + "numpy>=1.24,<2.5", "numba>=0.60", "numba-cuda>=0.23", "pytest>=8.0", @@ -249,6 +249,25 @@ def test_all_models_gpu() -> dict: ("OpenBoostRegressor", ob.OpenBoostRegressor(n_estimators=10, max_depth=3), y), ("OpenBoostClassifier", ob.OpenBoostClassifier(n_estimators=10, max_depth=3), y_binary.astype(np.int32)), ] + + def _curve(theta, xx): + a, b = theta + return a * np.abs(xx) ** (1.0 / (1.0 + np.exp(-np.clip(b * xx, -20, 20)))) + + try: + x_struct = np.abs(X[:, 0]) + 0.3 + fb = ob.FormulaBoost( + formula=_curve, n_params=2, links=("log", "identity"), + n_trees=10, max_depth=3, + ) + fb.fit(X, y_positive, model_input=x_struct) + pred = fb.predict(X, model_input=x_struct) + results["FormulaBoost"] = { + "status": "passed", + "prediction_shape": pred.shape, + } + except Exception as e: + results["FormulaBoost"] = {"status": "failed", "error": str(e)} for name, model, target in models_to_test: try: @@ -267,6 +286,80 @@ def test_all_models_gpu() -> dict: return results +@app.function(gpu="T4", image=image, timeout=600) +def verify_unified_gpu() -> dict: + """End-to-end check of the unified trainer on a real GPU.""" + import sys + import traceback + + import numpy as np + + sys.path.insert(0, "/root/src") + import openboost as ob + + out = {"backend": None, "checks": {}, "error": None} + try: + ob.set_backend("cuda") + out["backend"] = ob.get_backend() + rng = np.random.default_rng(0) + X = rng.normal(size=(800, 6)).astype(np.float32) + y = (X[:, 0] * 1.5 + rng.normal(size=800) * 0.4).astype(np.float32) + + nb = ob.NaturalBoostNormal(n_trees=25, max_depth=3, learning_rate=0.1) + nb.fit(X, y) + nll = float(nb.nll(X, y)) + lo, hi = nb.predict_interval(X, alpha=0.1) + out["checks"]["naturalboost_nll"] = nll + out["checks"]["naturalboost_interval_ok"] = bool(np.all(lo <= hi)) + out["checks"]["naturalboost_trees"] = { + k: len(v) for k, v in nb.trees_.items() + } + + poiss_y = rng.poisson(np.exp(0.2 + 0.4 * X[:, 0]), 800).astype(np.float32) + po = ob.NaturalBoostPoisson(n_trees=15, max_depth=3) + po.fit(X, poiss_y) + out["checks"]["poisson_mean_finite"] = bool(np.all(np.isfinite(po.predict(X)))) + + def curve(theta, xx): + a, b = theta + return a * xx ** (1.0 / (1.0 + np.exp(-np.clip(b * xx, -20.0, 20.0)))) + + x = np.abs(X[:, 1]) + 0.4 + fb = ob.FormulaBoost( + formula=curve, n_params=2, links=("log", "identity"), + n_trees=20, max_depth=3, + ) + fb.fit(X, np.abs(y) + 0.2, model_input=x) + pred = fb.predict(X, model_input=x) + params = fb.predict_params(X) + out["checks"]["formula_pred_finite"] = bool(np.all(np.isfinite(pred))) + out["checks"]["formula_a_positive"] = bool(np.all(params["theta_0"] > 0)) + out["ok"] = all( + ( + out["backend"] == "cuda", + np.isfinite(nll), + out["checks"]["naturalboost_interval_ok"], + out["checks"]["poisson_mean_finite"], + out["checks"]["formula_pred_finite"], + out["checks"]["formula_a_positive"], + ) + ) + except Exception as exc: + out["ok"] = False + out["error"] = f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}" + return out + + +@app.local_entrypoint() +def unified(): + """Run the unified-trainer GPU check.""" + result = verify_unified_gpu.remote() + print(result) + if not result.get("ok"): + raise RuntimeError(f"Unified GPU check failed: {result.get('error')}") + print("Unified GPU check passed.") + + @app.local_entrypoint() def ci(): """Run the CUDA test suite for CI and propagate failures to the caller.""" diff --git a/tests/test_cuda_verification.py b/tests/test_cuda_verification.py index 8cead58..85184a4 100644 --- a/tests/test_cuda_verification.py +++ b/tests/test_cuda_verification.py @@ -248,6 +248,51 @@ def test_natural_boost_poisson_gpu(self): mean = model.predict(X) assert mean.shape == (500,) + def test_natural_boost_normal_cpu_gpu_parity(self, sample_data): + """Device-resident Normal path should match the CPU trainer closely.""" + X, y = sample_data + kwargs = dict(n_trees=30, max_depth=3, learning_rate=0.1, n_bins=64) + + ob.set_backend("cpu") + cpu = ob.NaturalBoostNormal(**kwargs) + cpu.fit(X, y) + nll_cpu = cpu.nll(X, y) + pred_cpu = cpu.predict(X) + + ob.set_backend("cuda") + gpu = ob.NaturalBoostNormal(**kwargs) + gpu.fit(X, y) + nll_gpu = gpu.nll(X, y) + pred_gpu = gpu.predict(X) + + assert np.isfinite(nll_gpu) + assert abs(nll_gpu - nll_cpu) < 0.15 + assert np.corrcoef(pred_cpu, pred_gpu)[0, 1] > 0.95 + + def test_formula_boost_gpu(self, sample_data): + """FormulaBoost builds trees on GPU (GGN stays on host).""" + X, y = sample_data + x = np.abs(X[:, 0]) + 0.3 + + def curve(theta, xx): + a, b = theta + return a * xx ** (1.0 / (1.0 + np.exp(-np.clip(b * xx, -20, 20)))) + + ob.set_backend("cuda") + model = ob.FormulaBoost( + formula=curve, + n_params=2, + links=("log", "identity"), + n_trees=15, + max_depth=3, + ) + model.fit(X, np.abs(y) + 0.2, model_input=x) + pred = model.predict(X, model_input=x) + params = model.predict_params(X) + assert pred.shape == (len(y),) + assert np.all(np.isfinite(pred)) + assert np.all(params["theta_0"] > 0) + @pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") class TestOtherModelsGPU: diff --git a/tests/test_formula.py b/tests/test_formula.py new file mode 100644 index 0000000..f9ac079 --- /dev/null +++ b/tests/test_formula.py @@ -0,0 +1,199 @@ +"""Tests for FormulaBoost and the unified trainer path it shares.""" + +from __future__ import annotations + +import numpy as np +import pytest + + +def _sigmoid(t): + return 1.0 / (1.0 + np.exp(-np.clip(t, -30.0, 30.0))) + + +def sales_curve(theta, x): + a, b = theta + return a * x ** _sigmoid(b * x) + + +def _sales_data(n=2_000, seed=0): + rng = np.random.default_rng(seed) + Z = rng.uniform(0, 1, (n, 4)) + x = rng.uniform(0.3, 2.5, n) + u = 0.4 + 0.8 * Z[:, 0] - 0.5 * Z[:, 1] + a = np.exp(u) + b = 0.6 + 1.5 * Z[:, 2] + f = sales_curve((a, b), x) + y = f + 0.05 * f.std() * rng.standard_normal(n) + return Z, x, y, a, b, f + + +class TestFormulaBoost: + def test_import(self): + import openboost as ob + + assert hasattr(ob, "FormulaBoost") + + def test_fit_predict_shapes(self): + from openboost import FormulaBoost + + Z, x, y, *_ = _sales_data(n=800, seed=1) + model = FormulaBoost( + formula=sales_curve, + n_params=2, + links=("log", "identity"), + param_names=("a", "b"), + n_trees=20, + max_depth=3, + learning_rate=0.1, + ) + model.fit(Z, y, model_input=x) + params = model.predict_params(Z) + assert set(params) == {"a", "b"} + assert params["a"].shape == (len(y),) + assert np.all(params["a"] > 0) + pred = model.predict(Z, model_input=x) + assert pred.shape == (len(y),) + assert np.all(np.isfinite(pred)) + + def test_beats_global_and_recovers_a(self): + """GGN boosting should recover the scale surface and beat a global curve.""" + from openboost import FormulaBoost + + Z, x, y, a, b, f = _sales_data(n=3_000, seed=2) + n_tr = 2_400 + model = FormulaBoost( + formula=sales_curve, + n_params=2, + links=("log", "identity"), + param_names=("a", "b"), + n_trees=80, + max_depth=3, + learning_rate=0.1, + precond="full", + ) + model.fit(Z[:n_tr], y[:n_tr], model_input=x[:n_tr]) + params = model.predict_params(Z[n_tr:]) + pred = model.predict(Z[n_tr:], model_input=x[n_tr:]) + rmse = float(np.sqrt(np.mean((pred - y[n_tr:]) ** 2))) + global_rmse = float(np.sqrt(np.mean((y[n_tr:] - y[:n_tr].mean()) ** 2))) + assert rmse < 0.5 * global_rmse + assert np.corrcoef(params["a"], a[n_tr:])[0, 1] > 0.9 + + def test_extrapolation_beats_blackbox(self): + from openboost import FormulaBoost, GradientBoosting + + Z, x, y, a, b, f = _sales_data(n=3_000, seed=3) + n_tr = 2_400 + model = FormulaBoost( + formula=sales_curve, + n_params=2, + links=("log", "identity"), + param_names=("a", "b"), + n_trees=80, + max_depth=3, + learning_rate=0.1, + ) + model.fit(Z[:n_tr], y[:n_tr], model_input=x[:n_tr]) + + rng = np.random.default_rng(3) + Z_ex = rng.uniform(0, 1, (600, 4)) + x_ex = rng.uniform(3.0, 5.0, 600) + u = 0.4 + 0.8 * Z_ex[:, 0] - 0.5 * Z_ex[:, 1] + a_ex = np.exp(u) + b_ex = 0.6 + 1.5 * Z_ex[:, 2] + f_ex = sales_curve((a_ex, b_ex), x_ex) + + fb_rmse = float( + np.sqrt(np.mean((model.predict(Z_ex, model_input=x_ex) - f_ex) ** 2)) + ) + + gb = GradientBoosting(n_trees=80, max_depth=6, learning_rate=0.1) + gb.fit(np.column_stack([Z[:n_tr], x[:n_tr]]), y[:n_tr]) + gb_rmse = float( + np.sqrt(np.mean((gb.predict(np.column_stack([Z_ex, x_ex])) - f_ex) ** 2)) + ) + assert fb_rmse < 0.5 * gb_rmse + + def test_plain_precond_worse_than_full(self): + from openboost import FormulaBoost + + Z, x, y, *_ = _sales_data(n=1_500, seed=4) + n_tr, n_te = 1_200, 300 + kwargs = dict( + formula=sales_curve, + n_params=2, + links=("log", "identity"), + n_trees=40, + max_depth=3, + learning_rate=0.1, + ) + full = FormulaBoost(precond="full", **kwargs) + plain = FormulaBoost(precond="plain", **kwargs) + full.fit(Z[:n_tr], y[:n_tr], model_input=x[:n_tr]) + plain.fit(Z[:n_tr], y[:n_tr], model_input=x[:n_tr]) + rmse_full = np.sqrt( + np.mean((full.predict(Z[-n_te:], model_input=x[-n_te:]) - y[-n_te:]) ** 2) + ) + rmse_plain = np.sqrt( + np.mean((plain.predict(Z[-n_te:], model_input=x[-n_te:]) - y[-n_te:]) ** 2) + ) + assert rmse_full < rmse_plain + + def test_eval_set_and_early_stopping(self): + from openboost import FormulaBoost + + Z, x, y, *_ = _sales_data(n=1_200, seed=5) + model = FormulaBoost( + formula=sales_curve, + n_params=2, + links=("log", "identity"), + n_trees=200, + max_depth=3, + learning_rate=0.2, + ) + model.fit( + Z[:800], + y[:800], + model_input=x[:800], + eval_set=[(Z[800:], y[800:], x[800:])], + early_stopping_rounds=15, + ) + assert "eval_0" in model.evals_result_ + assert len(model.evals_result_["eval_0"]["mse"]) < 200 + assert model.best_iteration_ < 200 + + def test_bad_eval_set_raises(self): + from openboost import FormulaBoost + + Z, x, y, *_ = _sales_data(n=100, seed=6) + model = FormulaBoost( + formula=sales_curve, n_params=2, links=("log", "identity"), n_trees=2 + ) + with pytest.raises(ValueError, match="eval_set"): + model.fit(Z, y, model_input=x, eval_set=[(Z, y)]) + + def test_unknown_link_raises(self): + from openboost import FormulaBoost + + with pytest.raises(ValueError, match="Unknown link"): + FormulaBoost( + formula=sales_curve, n_params=2, links=("log", "relu") + )._make_objective() + + +class TestTrainerNaturalBoostParity: + """NaturalBoost still trains and predicts after the trainer port.""" + + def test_naturalboost_fit_predict(self): + from openboost import NaturalBoostNormal + + rng = np.random.default_rng(0) + X = rng.normal(size=(300, 4)).astype(np.float32) + y = (X[:, 0] * 2 + rng.normal(size=300)).astype(np.float32) + model = NaturalBoostNormal(n_trees=15, max_depth=3) + model.fit(X, y) + pred = model.predict(X) + assert pred.shape == (300,) + params = model.predict_params(X) + assert params["scale"].min() > 0 + assert model.nll(X, y) < 3.0 diff --git a/tests/test_survival.py b/tests/test_survival.py new file mode 100644 index 0000000..6401e6f --- /dev/null +++ b/tests/test_survival.py @@ -0,0 +1,129 @@ +"""Tests for WeibullAFT: censored survival boosting on the unified trainer.""" + +from __future__ import annotations + +import numpy as np +import pytest + + +def _weibull_data(n=6000, d=5, censor_q=0.7, seed=0, vary_shape=True): + """Weibull AFT DGP with covariate-dependent scale (and optionally shape).""" + rng = np.random.default_rng(seed) + Z = rng.uniform(0, 1, (n, d)) + lam = np.exp(0.5 + 1.0 * Z[:, 0] - 0.8 * Z[:, 1]) + k = (np.exp(0.2 + 1.2 * Z[:, 2] - 0.6 * Z[:, 3]) if vary_shape + else np.full(n, 1.3)) + u = rng.uniform(1e-9, 1.0, n) + T = lam * (-np.log(u)) ** (1.0 / k) + C = rng.exponential(np.quantile(T, censor_q), n) + t = np.minimum(T, C) + event = (T <= C).astype(float) + return Z, t, event, lam, k + + +class TestWeibullAFT: + def test_import(self): + import openboost as ob + + assert hasattr(ob, "WeibullAFT") + + def test_fit_predict_shapes(self): + from openboost import WeibullAFT + + Z, t, ev, *_ = _weibull_data(n=800, seed=1) + m = WeibullAFT(n_trees=30, max_depth=3, learning_rate=0.1) + m.fit(Z, t, event=ev) + params = m.predict_params(Z) + assert set(params) == {"scale", "shape"} + assert params["scale"].shape == (len(t),) + assert np.all(params["scale"] > 0) + assert np.all(params["shape"] > 0) + med = m.predict(Z) + assert med.shape == (len(t),) + assert np.all(np.isfinite(med)) and np.all(med > 0) + + def test_global_init_recovers_constant_params(self): + """On covariate-free data the global MLE init should match truth.""" + from openboost._objectives import WeibullAFTObjective + + rng = np.random.default_rng(3) + n = 20000 + lam_true, k_true = 2.5, 1.7 + u = rng.uniform(1e-9, 1.0, n) + t = lam_true * (-np.log(u)) ** (1.0 / k_true) # uncensored + obj = WeibullAFTObjective() + init = obj.init_raw(t, None, {"event": np.ones(n)}) + assert np.exp(init["scale"]) == pytest.approx(lam_true, rel=0.05) + assert np.exp(init["shape"]) == pytest.approx(k_true, rel=0.05) + + def test_recovers_scale_and_shape_surfaces(self): + """The capability: boost BOTH lambda(z) and k(z) from censored data.""" + from openboost import WeibullAFT + + Z, t, ev, lam, k = _weibull_data(n=6000, seed=0, vary_shape=True) + m = WeibullAFT(n_trees=200, max_depth=3, learning_rate=0.1) + m.fit(Z, t, event=ev) + p = m.predict_params(Z) + corr_scale = np.corrcoef(np.log(p["scale"]), np.log(lam))[0, 1] + corr_shape = np.corrcoef(np.log(p["shape"]), np.log(k))[0, 1] + assert corr_scale > 0.85 + assert corr_shape > 0.80 # shape varying with covariates is recovered + + def test_censoring_is_used(self): + """Ignoring right-censoring biases survival time downward.""" + from openboost import WeibullAFT + + Z, t, ev, lam, k = _weibull_data(n=5000, seed=2, censor_q=0.5) + kwargs = dict(n_trees=120, max_depth=3, learning_rate=0.1) + with_cens = WeibullAFT(**kwargs).fit(Z, t, event=ev) + ignore_cens = WeibullAFT(**kwargs).fit(Z, t) # all treated as observed + med_with = np.median(with_cens.predict(Z)) + med_ignore = np.median(ignore_cens.predict(Z)) + # Treating censored points as events underestimates the time. + assert med_with > med_ignore + + def test_predict_quantile_and_survival_monotone(self): + from openboost import WeibullAFT + + Z, t, ev, *_ = _weibull_data(n=1500, seed=4) + m = WeibullAFT(n_trees=60, max_depth=3, learning_rate=0.1).fit(Z, t, event=ev) + q10 = m.predict_quantile(Z, 0.1) + q50 = m.predict_quantile(Z, 0.5) + q90 = m.predict_quantile(Z, 0.9) + assert np.all(q10 < q50) and np.all(q50 < q90) + # Survival is a decreasing function of time. + s_lo = m.predict_survival(Z, t=0.5) + s_hi = m.predict_survival(Z, t=5.0) + assert np.all(s_hi <= s_lo) + assert np.all((s_lo >= 0) & (s_lo <= 1)) + + def test_eval_set_and_early_stopping(self): + from openboost import WeibullAFT + + Z, t, ev, *_ = _weibull_data(n=2000, seed=5) + ntr = 1500 + m = WeibullAFT(n_trees=300, max_depth=3, learning_rate=0.2) + m.fit( + Z[:ntr], t[:ntr], event=ev[:ntr], + eval_set=[(Z[ntr:], t[ntr:], ev[ntr:])], + early_stopping_rounds=15, + ) + assert "eval_0" in m.evals_result_ + assert len(m.evals_result_["eval_0"]["nll"]) < 300 + assert m.best_iteration_ < 300 + + def test_positive_time_required(self): + from openboost import WeibullAFT + + Z, t, ev, *_ = _weibull_data(n=200, seed=6) + t = t.copy() + t[0] = 0.0 + with pytest.raises(ValueError, match="positive"): + WeibullAFT(n_trees=2).fit(Z, t, event=ev) + + def test_bad_event_length_raises(self): + from openboost import WeibullAFT + + Z, t, ev, *_ = _weibull_data(n=200, seed=7) + with pytest.raises(ValueError, match="event"): + WeibullAFT(n_trees=2).fit(Z, t, event=ev[:100]) From 1df6dcf11c687020dc65aa1bcc073a1b3a30549b Mon Sep 17 00:00:00 2001 From: J Xu Date: Mon, 17 Aug 2026 09:39:33 -0700 Subject: [PATCH 3/3] bench: probabilistic, formula, and survival capability suites Add the pivot yardsticks, each with a local CPU path and a Modal A100 entry point, writing JSON to benchmarks/results/: - bench_probabilistic.py: speed and quality vs NGBoost (A100: 1229x faster at 90K rows on a shared 500-tree budget, NLL tied). - bench_formula.py: sales-saturation-curve capability vs global fit, black-box GBDT, and a hand-rolled XGBoost multi-output custom objective (extrapolation ~21x better than black-box; full GGN recovers the saturation-speed surface a diagonal Hessian cannot). - bench_survival.py: Weibull AFT vs XGBoost survival:aft, which cannot vary the shape with covariates (OpenBoost recovers shape(z) at corr 0.997 with a better censored NLL). Also adds the GGN spike that grounds the approach and the unified-engine design document. Co-authored-by: Cursor --- benchmarks/bench_formula.py | 546 +++++++++++++++++ benchmarks/bench_probabilistic.py | 574 ++++++++++++++++++ benchmarks/bench_survival.py | 422 +++++++++++++ development/paramboost/spike_ggn.py | 315 ++++++++++ development/paramboost/spike_ggn_results.json | 71 +++ planning/unified-engine-design.md | 184 ++++++ 6 files changed, 2112 insertions(+) create mode 100644 benchmarks/bench_formula.py create mode 100644 benchmarks/bench_probabilistic.py create mode 100644 benchmarks/bench_survival.py create mode 100644 development/paramboost/spike_ggn.py create mode 100644 development/paramboost/spike_ggn_results.json create mode 100644 planning/unified-engine-design.md diff --git a/benchmarks/bench_formula.py b/benchmarks/bench_formula.py new file mode 100644 index 0000000..22a5687 --- /dev/null +++ b/benchmarks/bench_formula.py @@ -0,0 +1,546 @@ +"""Formula capability benchmark: the yardstick for "Boost every parameter". + +FormulaBoost fits a user formula ``y ~ f(theta(z), x)`` where each parameter +surface ``theta_k(z)`` is its own boosting ensemble and ``x`` enters only +through the formula. This suite proves the claim NGBoost and black-box GBDT +cannot express, on the sales/saturation curve from the GGN spike: + + sales = a(z) * x ** sigmoid(b(z) * x) + +Baselines (what a practitioner would actually reach for): + + global one (a, b) fit to all rows (formula, but no theta(z)). + blackbox GradientBoosting on [z, x] -> y (no formula; trees flatten + outside the training x-range, so extrapolation collapses). + xgb hand-rolled XGBoost multi-output custom objective with the SAME + FD Jacobian as FormulaBoost, but a DIAGONAL Hessian. XGBoost's + custom-objective API only accepts a diagonal Hessian ("the Hessian + for each row should be diagonal", XGBoost docs), so the off-diagonal + GGN term is inexpressible. This baseline is exactly FormulaBoost's + ``precond='diag'`` on xgb trees. + +FormulaBoost is run in all three preconditioning modes (plain / diag / full); +``full`` is the only one that uses the off-diagonal GGN coupling. + +Metrics: in-range test RMSE, extrapolation RMSE (x beyond training range, vs +the true noiseless curve), parameter recovery (corr + RMSE of a_hat, b_hat), +fit time. + +Usage: + # Local CPU + uv run --with xgboost python benchmarks/bench_formula.py + uv run --with xgboost python benchmarks/bench_formula.py --quick + + # Modal A100 (proves the formula path runs on GPU-built trees) + uv run modal run benchmarks/bench_formula.py + uv run modal run benchmarks/bench_formula.py --quick + +Output: benchmarks/results/formula_.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +PROJECT_ROOT = Path(__file__).parent.parent +RESULTS_DIR = Path(__file__).parent / "results" + +SEED = 42 +N_FEATURES = 6 +ROUNDS = 300 +LR = 0.1 +DEPTH = 3 +DAMP = 1.0 +BLACKBOX_DEPTH = 6 +X_TRAIN_RANGE = (0.25, 2.5) +X_EXTRAP_RANGE = (3.0, 5.0) +NOISE_FRAC = 0.05 + +# Acceptance gates (see planning/formula_capability_next plan). +GATE_EXTRAP_VS_BLACKBOX = 5.0 # full must beat black-box extrap by >= this +GATE_XGB_EXTRAP_SLACK = 1.15 # full extrap must be within this factor of xgb +GATE_XGB_CORRB_SLACK = 0.02 # full b-recovery must be >= xgb - this + + +# ============================================================================= +# Data generating process (ported from development/paramboost/spike_ggn.py) +# ============================================================================= + +def sigmoid(t): + return 1.0 / (1.0 + np.exp(-np.clip(t, -30.0, 30.0))) + + +def sales_curve(theta, x): + """formula(theta, x) = a * x ** sigmoid(b*x), stable in log space.""" + a, b = theta + return np.exp(np.log(np.clip(a, 1e-12, None)) + sigmoid(b * x) * np.log(x)) + + +def true_params(Z): + u = 0.5 + 1.0 * Z[:, 0] - 0.8 * Z[:, 1] + 0.5 * np.sin(2 * np.pi * Z[:, 2]) + a = np.exp(u) + b = 0.5 + 2.0 * Z[:, 3] + Z[:, 4] * Z[:, 0] + return a, b + + +def make_data(n, rng, x_range): + Z = rng.uniform(0, 1, (n, N_FEATURES)) + x = rng.uniform(*x_range, n) + a, b = true_params(Z) + f = sales_curve((a, b), x) + y = f + NOISE_FRAC * f.std() * rng.standard_normal(n) + return Z, x, y, a, b, f + + +# ============================================================================= +# Metrics +# ============================================================================= + +def rmse(pred, target): + return float(np.sqrt(np.mean((np.asarray(pred) - np.asarray(target)) ** 2))) + + +def _corr(pred, target): + """Pearson corr, or None when either side is constant (e.g. a global fit).""" + pred = np.asarray(pred, dtype=np.float64) + target = np.asarray(target, dtype=np.float64) + if pred.std() < 1e-12 or target.std() < 1e-12: + return None + return float(np.corrcoef(pred, target)[0, 1]) + + +def param_recovery(a_hat, b_hat, a_true, b_true): + return { + "rmse_a": rmse(a_hat, a_true), + "rmse_b": rmse(b_hat, b_true), + "corr_a": _corr(a_hat, a_true), + "corr_b": _corr(b_hat, b_true), + } + + +# ============================================================================= +# Shared formula math (identical FD Jacobian to FormulaObjective, so the only +# difference between the xgb baseline and FormulaBoost is diagonal-vs-full +# Hessian and xgb-vs-openboost trees). +# ============================================================================= + +def _raw_to_theta(raw): + """raw (n,2) -> (a, b): log link on a, identity on b.""" + return np.exp(np.clip(raw[:, 0], -30.0, 30.0)), raw[:, 1] + + +def _fd_jac_raw(raw, x, eps=1e-5): + """Forward-difference Jacobian df/d(raw), shape (n,2).""" + a, b = _raw_to_theta(raw) + f0 = sales_curve((a, b), x) + # bump raw_a + a1, _ = _raw_to_theta(np.column_stack([raw[:, 0] + eps, raw[:, 1]])) + f_a = sales_curve((a1, b), x) + # bump raw_b + _, b1 = _raw_to_theta(np.column_stack([raw[:, 0], raw[:, 1] + eps])) + f_b = sales_curve((a, b1), x) + jac = np.column_stack([(f_a - f0) / eps, (f_b - f0) / eps]) + return f0, jac + + +def fit_global_raw(x, y): + """Single global raw (u, b) via L-BFGS-B (matches FormulaObjective init).""" + from scipy.optimize import minimize + + v0 = np.zeros(2, dtype=np.float64) + if np.all(y > 0): + v0[0] = float(np.log(np.mean(y))) + + def loss(v): + a = np.exp(np.clip(v[0], -30.0, 30.0)) + pred = sales_curve((np.full_like(x, a), np.full_like(x, v[1])), x) + if not np.all(np.isfinite(pred)): + return 1e12 + return float(np.mean(0.5 * (pred - y) ** 2)) + + res = minimize(loss, v0, method="L-BFGS-B") + return res.x if res.success else v0 + + +# ============================================================================= +# Models +# ============================================================================= + +def run_formulaboost(precond, Z_tr, x_tr, y_tr, sets, rounds): + import openboost as ob + + model = ob.FormulaBoost( + formula=sales_curve, + n_params=2, + links=("log", "identity"), + param_names=("a", "b"), + precond=precond, + damp=DAMP, + n_trees=rounds, + max_depth=DEPTH, + learning_rate=LR, + ) + t0 = time.perf_counter() + model.fit(Z_tr, y_tr, model_input=x_tr) + fit_time = time.perf_counter() - t0 + + out = {"fit_time_s": round(fit_time, 2)} + for name, (Z, x, target, a_true, b_true) in sets.items(): + params = model.predict_params(Z) + pred = model.predict(Z, model_input=x) + entry = {"rmse": rmse(pred, target)} + entry.update(param_recovery(params["a"], params["b"], a_true, b_true)) + out[name] = entry + return out + + +def run_global(Z_tr, x_tr, y_tr, sets): + v = fit_global_raw(x_tr, y_tr) + a_g, b_g = float(np.exp(v[0])), float(v[1]) + out = {"a": a_g, "b": b_g} + for name, (_Z, x, target, a_true, b_true) in sets.items(): + pred = sales_curve((np.full_like(x, a_g), np.full_like(x, b_g)), x) + entry = {"rmse": rmse(pred, target)} + # a global constant "recovers" nothing; report corr for completeness. + entry.update(param_recovery( + np.full_like(a_true, a_g), np.full_like(b_true, b_g), + a_true, b_true)) + out[name] = entry + return out + + +def run_blackbox(Z_tr, x_tr, y_tr, sets, rounds): + import openboost as ob + + model = ob.GradientBoosting( + n_trees=rounds, max_depth=BLACKBOX_DEPTH, learning_rate=LR, + random_state=SEED, + ) + t0 = time.perf_counter() + model.fit(np.column_stack([Z_tr, x_tr]), y_tr) + fit_time = time.perf_counter() - t0 + + out = {"fit_time_s": round(fit_time, 2)} + for name, (Z, x, target, _a, _b) in sets.items(): + pred = model.predict(np.column_stack([Z, x])) + out[name] = {"rmse": rmse(pred, target)} # no theta(z) to recover + return out + + +def run_xgb_custom(Z_tr, x_tr, y_tr, sets, rounds): + """Hand-rolled XGBoost multi-output custom objective (diagonal Hessian). + + Same FD Jacobian and global init as FormulaBoost; the only difference is + that XGBoost's custom-objective API accepts a diagonal Hessian only, so + this is precond='diag' on xgb trees. Returns {"error": ...} if the + installed xgboost lacks multi-output custom objectives. + """ + try: + import xgboost as xgb + except ImportError as exc: + return {"error": f"xgboost unavailable: {exc}"} + + try: + v0 = fit_global_raw(x_tr, y_tr) + n = len(y_tr) + base_tr = np.tile(v0.astype(np.float64), (n, 1)) + + def obj(pred, dtrain): + pred = np.asarray(pred, dtype=np.float64).reshape(n, 2) + f0, jac = _fd_jac_raw(pred, x_tr) + residual = f0 - y_tr + grad = residual[:, None] * jac + hess = np.maximum(jac * jac, 1e-6) # diagonal only + return grad.astype(np.float32), hess.astype(np.float32) + + dtrain = xgb.DMatrix( + np.ascontiguousarray(Z_tr), + label=np.zeros((n, 2)), # unused: obj closes over y, x + base_margin=base_tr, + ) + params = { + "tree_method": "hist", + "num_target": 2, + "base_score": 0.0, + "disable_default_eval_metric": True, + "max_depth": DEPTH, + "eta": LR, + "lambda": DAMP, + "seed": SEED, + } + t0 = time.perf_counter() + booster = xgb.train(params, dtrain, num_boost_round=rounds, obj=obj) + fit_time = time.perf_counter() - t0 + + out = {"fit_time_s": round(fit_time, 2), "note": "diagonal Hessian only"} + for name, (Z, x, target, a_true, b_true) in sets.items(): + m = xgb.DMatrix( + np.ascontiguousarray(Z), + base_margin=np.tile(v0.astype(np.float64), (len(x), 1)), + ) + raw = np.asarray(booster.predict(m), dtype=np.float64).reshape(-1, 2) + a_hat, b_hat = _raw_to_theta(raw) + pred = sales_curve((a_hat, b_hat), x) + entry = {"rmse": rmse(pred, target)} + entry.update(param_recovery(a_hat, b_hat, a_true, b_true)) + out[name] = entry + return out + except Exception as exc: # noqa: BLE001 - baseline must never crash the suite + import traceback + return {"error": f"{type(exc).__name__}: {exc}", + "traceback": traceback.format_exc()} + + +# ============================================================================= +# Gates +# ============================================================================= + +def compute_gates(res): + """Boolean acceptance gates from one size's results.""" + full = res["formulaboost"]["full"] + plain = res["formulaboost"]["plain"] + blackbox = res["blackbox"] + glob = res["global"] + xgb = res.get("xgb", {}) + + gates = {} + + # 1. formula constrains the x-shape: extrapolation crushes black-box. + ratio_bb = blackbox["extrap"]["rmse"] / max(full["extrap"]["rmse"], 1e-12) + gates["extrap_vs_blackbox_ratio"] = round(ratio_bb, 2) + gates["extrap_vs_blackbox_pass"] = ratio_bb >= GATE_EXTRAP_VS_BLACKBOX + + # 2. off-diagonal GGN pays: full beats plain and recovers b better. + gates["full_vs_plain_test_ratio"] = round( + plain["test"]["rmse"] / max(full["test"]["rmse"], 1e-12), 2) + gates["full_corr_b"] = round(full["test"]["corr_b"], 3) + gates["plain_corr_b"] = round(plain["test"]["corr_b"], 3) + gates["full_beats_plain_pass"] = ( + full["test"]["rmse"] < plain["test"]["rmse"] + and full["test"]["corr_b"] > plain["test"]["corr_b"] + ) + + # 3. theta(z) is necessary: full beats the global constant curve. + gates["full_beats_global_pass"] = ( + full["test"]["rmse"] < glob["test"]["rmse"] + and full["extrap"]["rmse"] < glob["extrap"]["rmse"] + ) + + # 4. not worse than the hand-rolled XGBoost diagonal baseline. + if xgb and "error" not in xgb: + extrap_ok = (full["extrap"]["rmse"] + <= xgb["extrap"]["rmse"] * GATE_XGB_EXTRAP_SLACK) + corrb_ok = (full["test"]["corr_b"] + >= xgb["test"]["corr_b"] - GATE_XGB_CORRB_SLACK) + gates["full_vs_xgb_extrap_ratio"] = round( + xgb["extrap"]["rmse"] / max(full["extrap"]["rmse"], 1e-12), 2) + gates["full_corr_b_minus_xgb"] = round( + full["test"]["corr_b"] - xgb["test"]["corr_b"], 3) + gates["full_not_worse_than_xgb_pass"] = bool(extrap_ok and corrb_ok) + else: + gates["full_not_worse_than_xgb_pass"] = None # skipped (xgb missing) + + hard = [ + gates["extrap_vs_blackbox_pass"], + gates["full_beats_plain_pass"], + gates["full_beats_global_pass"], + ] + gates["passed"] = all(hard) + return gates + + +# ============================================================================= +# Runner +# ============================================================================= + +def run_size(n_train, n_test, rounds): + import openboost as ob + + rng = np.random.default_rng(SEED) + Z_tr, x_tr, y_tr, *_ = make_data(n_train, rng, X_TRAIN_RANGE) + Z_te, x_te, y_te, a_te, b_te, _ = make_data(n_test, rng, X_TRAIN_RANGE) + # Extrapolation: same z, x beyond the training range; target is the TRUE + # noiseless curve (did the model learn the functional form, not the noise?). + Z_ex, x_ex, _, a_ex, b_ex, f_ex = make_data(n_test, rng, X_EXTRAP_RANGE) + + sets = { + "test": (Z_te, x_te, y_te, a_te, b_te), + "extrap": (Z_ex, x_ex, f_ex, a_ex, b_ex), + } + + print(f"\n=== n_train={n_train:,} rounds={rounds} " + f"backend={ob.get_backend()} ===", flush=True) + + res = {"n_train": n_train, "n_test": n_test, "rounds": rounds, + "noise_floor_test_rmse": None, "formulaboost": {}} + + _, _, _, _, _, f_te = make_data(n_test, np.random.default_rng(SEED + 1), + X_TRAIN_RANGE) + res["noise_floor_test_rmse"] = round(float(NOISE_FRAC * f_te.std()), 4) + + for precond in ("plain", "diag", "full"): + out = run_formulaboost(precond, Z_tr, x_tr, y_tr, sets, rounds) + res["formulaboost"][precond] = out + print(f" FormulaBoost[{precond:>5}] test {out['test']['rmse']:.4f} " + f"extrap {out['extrap']['rmse']:.4f} " + f"corr_b {out['test']['corr_b']:.3f} ({out['fit_time_s']}s)", + flush=True) + + res["global"] = run_global(Z_tr, x_tr, y_tr, sets) + print(f" global test {res['global']['test']['rmse']:.4f} " + f"extrap {res['global']['extrap']['rmse']:.4f}", flush=True) + + res["blackbox"] = run_blackbox(Z_tr, x_tr, y_tr, sets, rounds) + print(f" blackbox GBDT test {res['blackbox']['test']['rmse']:.4f} " + f"extrap {res['blackbox']['extrap']['rmse']:.4f} " + f"({res['blackbox']['fit_time_s']}s)", flush=True) + + res["xgb"] = run_xgb_custom(Z_tr, x_tr, y_tr, sets, rounds) + if "error" in res["xgb"]: + print(f" xgb custom SKIPPED ({res['xgb']['error']})", flush=True) + else: + print(f" xgb custom(diag) test {res['xgb']['test']['rmse']:.4f} " + f"extrap {res['xgb']['extrap']['rmse']:.4f} " + f"corr_b {res['xgb']['test']['corr_b']:.3f} " + f"({res['xgb']['fit_time_s']}s)", flush=True) + + res["gates"] = compute_gates(res) + g = res["gates"] + print(f" GATES extrap_vs_blackbox={g['extrap_vs_blackbox_ratio']}x " + f"({_pf(g['extrap_vs_blackbox_pass'])}) " + f"full>plain={_pf(g['full_beats_plain_pass'])} " + f"full>global={_pf(g['full_beats_global_pass'])} " + f"full>=xgb={_pf(g['full_not_worse_than_xgb_pass'])} " + f"=> {'PASS' if g['passed'] else 'FAIL'}", flush=True) + return res + + +def _pf(v): + if v is None: + return "skip" + return "PASS" if v else "FAIL" + + +def run_suite(quick=False): + import openboost as ob + + sizes = ([(8_000, 3_000, 80)] if quick + else [(40_000, 10_000, ROUNDS), (200_000, 20_000, ROUNDS)]) + + # Warmup (JIT the tree kernels off the clock). + rng = np.random.default_rng(0) + Zw, xw, yw, *_ = make_data(512, rng, X_TRAIN_RANGE) + run_formulaboost("full", Zw, xw, yw, + {"test": (Zw, xw, yw, *true_params(Zw))}, rounds=5) + + report = { + "benchmark": "bench_formula", + "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "seed": SEED, + "quick_mode": quick, + "formula": "sales = a(z) * x ** sigmoid(b(z) * x)", + "budget": {"rounds": ROUNDS, "lr": LR, "depth": DEPTH, "damp": DAMP, + "blackbox_depth": BLACKBOX_DEPTH}, + "ranges": {"x_train": X_TRAIN_RANGE, "x_extrap": X_EXTRAP_RANGE, + "noise_frac": NOISE_FRAC}, + "platform": { + "python": platform.python_version(), + "system": f"{platform.system()} {platform.machine()}", + }, + "backend": ob.get_backend(), + "versions": {"openboost": ob.__version__, "numpy": np.__version__}, + "sizes": [], + } + try: + import xgboost + report["versions"]["xgboost"] = xgboost.__version__ + except ImportError: + report["versions"]["xgboost"] = None + + for n_train, n_test, rounds in sizes: + report["sizes"].append(run_size(n_train, n_test, rounds)) + + report["passed"] = all(s["gates"]["passed"] for s in report["sizes"]) + return report + + +def save_report(report): + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + out = RESULTS_DIR / f"formula_{stamp}.json" + out.write_text(json.dumps(report, indent=2) + "\n") + print(f"\nWrote {out}") + return out + + +# ============================================================================= +# Modal entry point +# ============================================================================= + +try: + import modal + + app = modal.App("openboost-formula-bench") + image = ( + modal.Image.from_registry( + "nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12") + .pip_install( + "numpy>=1.24,<2.5", "numba>=0.60", "numba-cuda>=0.23", + "scipy>=1.10", "scikit-learn>=1.0", "xgboost>=2.0", + "joblib>=1.2", + ) + .add_local_dir( + str(PROJECT_ROOT / "src" / "openboost"), + remote_path="/root/openboost", + copy=True, + ) + .env({"PYTHONPATH": "/root", "OPENBOOST_BACKEND": "cuda"}) + ) +except ImportError: + modal = None + app = None + image = None + +if modal is not None and app is not None: + + @app.function(gpu="A100", image=image, timeout=2 * 3600) + def _run_remote(quick: bool = False): + sys.path.insert(0, "/root") + import openboost as ob + + ob.set_backend("cuda") + print(f"backend={ob.get_backend()} quick={quick}", flush=True) + return run_suite(quick=quick) + + @app.local_entrypoint() + def main(quick: bool = False): + report = _run_remote.remote(quick=quick) + save_report(report) + print(f"\nOverall: {'PASS' if report['passed'] else 'FAIL'}") + + +# ============================================================================= +# Local execution +# ============================================================================= + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--quick", action="store_true", + help="Smoke-test sizes/rounds") + args = parser.parse_args() + + os.environ.setdefault("OPENBOOST_BACKEND", "cpu") + sys.path.insert(0, str(PROJECT_ROOT / "src")) + + report = run_suite(quick=args.quick) + save_report(report) + print(f"\nOverall: {'PASS' if report['passed'] else 'FAIL'}") diff --git a/benchmarks/bench_probabilistic.py b/benchmarks/bench_probabilistic.py new file mode 100644 index 0000000..3def88a --- /dev/null +++ b/benchmarks/bench_probabilistic.py @@ -0,0 +1,574 @@ +"""Probabilistic GBDT benchmark: the yardstick for the parametric-boosting pivot. + +Two suites (see tasks/todo.md and planning/unified-engine-design.md): + + quality Paired-split comparison vs NGBoost on real datasets. + Metrics: NLL, CRPS, RMSE, 90% interval coverage, pinball loss. + Paired Wilcoxon test on per-split NLL. CPU by default (NGBoost is + CPU-only; set OPENBOOST_BACKEND to control OpenBoost's side). + + speed Wall-clock fit time + time-to-same-NLL vs NGBoost (and PGBM if + installed) on large synthetic heteroscedastic data. Meant for + Modal A100; also runs locally at reduced sizes. + +Usage: + # Local quality suite (CPU, honest vs NGBoost) + uv run --with ngboost python benchmarks/bench_probabilistic.py --suite quality + uv run --with ngboost python benchmarks/bench_probabilistic.py --suite quality --quick + + # Local speed suite (small sizes, sanity only) + uv run --with ngboost python benchmarks/bench_probabilistic.py --suite speed --quick + + # Modal A100 (speed is the ≥10x gate; quality is CPU-honest vs NGBoost) + uv run modal run benchmarks/bench_probabilistic.py --suite speed + uv run modal run benchmarks/bench_probabilistic.py --suite all + +Output: benchmarks/results/probabilistic__.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +PROJECT_ROOT = Path(__file__).parent.parent +RESULTS_DIR = Path(__file__).parent / "results" + +SEED = 42 +# Shared budget, deliberately close to NGBoost defaults (its native regime). +N_TREES = 500 +LEARNING_RATE = 0.03 +MAX_DEPTH = 3 + +# quality suite: (n_splits_small, n_splits_large) — large datasets get fewer +# paired splits because each fit is expensive. +N_SPLITS_SMALL = 20 +N_SPLITS_LARGE = 5 +LARGE_THRESHOLD = 50_000 +# NGBoost-paper-style protocol: hold out a validation set from the training +# split; both libraries early-stop on the SAME validation set. +VAL_FRACTION = 0.2 +ES_PATIENCE = 50 + +# ============================================================================= +# Dataset registry (quality suite) +# ============================================================================= +# NGBoost-paper-style UCI datasets, fetched from OpenML by name (best effort: +# a fetch failure skips the dataset with a note instead of crashing the run). + +OPENML_DATASETS = [ + # (short name, openml name, version) + ("boston", "boston", 1), + ("concrete", "Concrete_Compressive_Strength", 1), + ("energy", "energy-efficiency", 1), + ("kin8nm", "kin8nm", 1), + ("naval", "naval_propulsion_plant", 1), + ("power", "combined_cycle_power_plant", 1), + ("protein", "physicochemical-protein", 1), + ("wine", "wine_quality", 1), + ("yacht", "yacht_hydrodynamics", 1), +] + + +def load_quality_datasets(quick: bool = False) -> tuple[list[dict], list[str]]: + import numpy as np + + datasets: list[dict] = [] + notes: list[str] = [] + + from sklearn.datasets import fetch_california_housing, fetch_openml + + if quick: + wanted = OPENML_DATASETS[:1] + else: + wanted = OPENML_DATASETS + + for short, name, version in wanted: + try: + bunch = fetch_openml(name=name, version=version, as_frame=True, + parser="auto") + frame = bunch.frame.select_dtypes(include=[np.number]).dropna() + target_col = bunch.target_names[0] if bunch.target_names else \ + frame.columns[-1] + if target_col not in frame.columns: + target_col = frame.columns[-1] + y = frame[target_col].to_numpy(dtype=np.float64) + X = frame.drop(columns=[target_col]).to_numpy(dtype=np.float64) + datasets.append({"name": short, "X": X, "y": y, + "source": f"openml:{name}/v{version}"}) + except Exception as exc: + notes.append(f"{short} skipped ({type(exc).__name__}: {exc})") + print(f"WARNING: {notes[-1]}") + + try: + housing = fetch_california_housing() + X, y = housing.data, housing.target + if quick: + X, y = X[:2_000], y[:2_000] + datasets.append({"name": "california", "X": X.astype("float64"), + "y": y.astype("float64"), + "source": "sklearn:california_housing"}) + except Exception as exc: + notes.append(f"california skipped ({type(exc).__name__}: {exc})") + print(f"WARNING: {notes[-1]}") + + if not quick: + try: + msd = fetch_openml(name="YearPredictionMSD", version=1, + as_frame=False, parser="auto") + X = msd.data.astype("float64") + y = msd.target.astype("float64") + # Subsample for the quality suite: NGBoost's exact-split trees at + # this budget take hours on the full 515K. Full size is covered + # by the speed suite. + rng = np.random.RandomState(SEED) + idx = rng.choice(len(y), 100_000, replace=False) + datasets.append({"name": "year_msd_100k", "X": X[idx], "y": y[idx], + "source": "openml:YearPredictionMSD/v1 (100K subsample)"}) + except Exception as exc: + notes.append(f"year_msd skipped ({type(exc).__name__}: {exc})") + print(f"WARNING: {notes[-1]}") + + return datasets, notes + + +def make_heteroscedastic(n_samples: int, n_features: int = 80, seed: int = SEED): + """Speed-suite synthetic: Friedman#1 signal + feature-dependent noise.""" + import numpy as np + + rng = np.random.RandomState(seed) + X = rng.uniform(0, 1, (n_samples, n_features)).astype(np.float32) + f = ( + 10 * np.sin(np.pi * X[:, 0] * X[:, 1]) + + 20 * (X[:, 2] - 0.5) ** 2 + + 10 * X[:, 3] + + 5 * X[:, 4] + ) + sigma = 1.0 + 2.0 * X[:, 5] + y = (f + sigma * rng.randn(n_samples)).astype(np.float64) + return X.astype(np.float64), y + + +# ============================================================================= +# Models +# ============================================================================= + +def make_openboost(): + import openboost as ob + + return ob.NaturalBoostNormal( + n_trees=N_TREES, learning_rate=LEARNING_RATE, max_depth=MAX_DEPTH, + n_bins=254, + ) + + +def make_ngboost(): + from ngboost import NGBRegressor + from ngboost.distns import Normal + from sklearn.tree import DecisionTreeRegressor + + return NGBRegressor( + Dist=Normal, + n_estimators=N_TREES, + learning_rate=LEARNING_RATE, + Base=DecisionTreeRegressor(criterion="friedman_mse", max_depth=MAX_DEPTH), + natural_gradient=True, + verbose=False, + random_state=SEED, + ) + + +def make_pgbm(): + """PGBM (torch, GPU-capable probabilistic GBDT). Optional.""" + from pgbm.sklearn import HistGradientBoostingRegressor as PGBMRegressor + + return PGBMRegressor(max_iter=N_TREES, learning_rate=LEARNING_RATE, + max_depth=MAX_DEPTH, random_state=SEED) + + +def predict_normal_params(model, X): + import numpy as np + + if hasattr(model, "predict_distribution"): # OpenBoost + # early_stopping_rounds restores the best iteration internally + params = model.predict_distribution(X).params + return (np.asarray(params["loc"], dtype=np.float64), + np.asarray(params["scale"], dtype=np.float64)) + if hasattr(model, "pred_dist"): # NGBoost + best_iter = getattr(model, "best_val_loss_itr", None) + dist = (model.pred_dist(X, max_iter=best_iter) + if best_iter is not None else model.pred_dist(X)) + params = dist.params + return (np.asarray(params["loc"], dtype=np.float64), + np.asarray(params["scale"], dtype=np.float64)) + # PGBM sklearn wrapper + mean, std = model.predict(X, return_std=True) + return np.asarray(mean, dtype=np.float64), np.asarray(std, dtype=np.float64) + + +# ============================================================================= +# Metrics (shared closed forms — identical for every library) +# ============================================================================= + +def gaussian_nll(y, mean, std): + import numpy as np + + std = np.clip(std, 1e-12, None) + return float(np.mean(0.5 * np.log(2 * np.pi * std**2) + + (y - mean) ** 2 / (2 * std**2))) + + +def coverage_90(y, mean, std): + import numpy as np + from scipy.stats import norm + + lo = mean + norm.ppf(0.05) * std + hi = mean + norm.ppf(0.95) * std + return float(np.mean((y >= lo) & (y <= hi))) + + +def pinball(y, mean, std, quantiles=(0.05, 0.5, 0.95)): + import numpy as np + from scipy.stats import norm + + losses = {} + for q in quantiles: + pred_q = mean + norm.ppf(q) * std + diff = y - pred_q + losses[str(q)] = float(np.mean(np.maximum(q * diff, (q - 1) * diff))) + return losses + + +def score_model(model, X_te, y_te): + import numpy as np + + import openboost as ob + + mean, std = predict_normal_params(model, X_te) + return { + "nll": gaussian_nll(y_te, mean, std), + "crps": float(ob.crps_gaussian(y_te, mean, std)), + "rmse": float(np.sqrt(np.mean((y_te - mean) ** 2))), + "coverage_90": coverage_90(y_te, mean, std), + "pinball": pinball(y_te, mean, std), + } + + +# ============================================================================= +# Quality suite +# ============================================================================= + +def run_quality(quick: bool = False) -> dict: + import numpy as np + from scipy.stats import wilcoxon + from sklearn.model_selection import train_test_split + + import openboost as ob + + datasets, notes = load_quality_datasets(quick=quick) + + # Untimed warmup (JIT compile both sides) + Xw, yw = make_heteroscedastic(512, n_features=10, seed=0) + make_openboost().fit(Xw[:256], yw[:256]) + ngb_w = make_ngboost() + ngb_w.set_params(n_estimators=5) + ngb_w.fit(Xw[:256], yw[:256]) + + results = [] + for ds in datasets: + X, y = ds["X"], ds["y"] + n_splits = (2 if quick else + N_SPLITS_LARGE if len(y) > LARGE_THRESHOLD else + N_SPLITS_SMALL) + per_split = {"openboost": [], "ngboost": []} + t_ds = time.perf_counter() + for split_seed in range(n_splits): + X_tr, X_te, y_tr, y_te = train_test_split( + X, y, test_size=0.2, random_state=split_seed) + # Both libraries early-stop on the SAME validation set. + X_fit, X_val, y_fit, y_val = train_test_split( + X_tr, y_tr, test_size=VAL_FRACTION, random_state=split_seed) + for lib, factory in (("openboost", make_openboost), + ("ngboost", make_ngboost)): + model = factory() + t0 = time.perf_counter() + if lib == "openboost": + model.fit(X_fit, y_fit, eval_set=[(X_val, y_val)], + early_stopping_rounds=ES_PATIENCE) + else: + model.fit(X_fit, y_fit, X_val=X_val, Y_val=y_val, + early_stopping_rounds=ES_PATIENCE) + fit_time = time.perf_counter() - t0 + metrics = score_model(model, X_te, y_te) + metrics["fit_time_s"] = round(fit_time, 3) + per_split[lib].append(metrics) + + def agg(lib, key): + vals = [m[key] for m in per_split[lib]] + return {"mean": float(np.mean(vals)), "std": float(np.std(vals))} + + ob_nll = np.array([m["nll"] for m in per_split["openboost"]]) + ngb_nll = np.array([m["nll"] for m in per_split["ngboost"]]) + if n_splits >= 5 and not np.allclose(ob_nll, ngb_nll): + stat = wilcoxon(ob_nll, ngb_nll) + p_value = float(stat.pvalue) + else: + p_value = None + + row = { + "dataset": ds["name"], + "source": ds["source"], + "n": int(len(y)), + "n_features": int(X.shape[1]), + "n_splits": n_splits, + "openboost": {k: agg("openboost", k) + for k in ("nll", "crps", "rmse", "coverage_90", + "fit_time_s")}, + "ngboost": {k: agg("ngboost", k) + for k in ("nll", "crps", "rmse", "coverage_90", + "fit_time_s")}, + "nll_delta_mean": float(np.mean(ob_nll - ngb_nll)), + "nll_paired_wilcoxon_p": p_value, + "per_split": per_split, + } + results.append(row) + print(f"{ds['name']:<12} splits={n_splits} " + f"OB NLL {row['openboost']['nll']['mean']:.4f} " + f"NGB NLL {row['ngboost']['nll']['mean']:.4f} " + f"delta {row['nll_delta_mean']:+.4f} p={p_value} " + f"OB cov90 {row['openboost']['coverage_90']['mean']:.3f} " + f"({time.perf_counter() - t_ds:.0f}s)") + + return {"suite": "quality", "results": results, "skipped": notes, + "backend": ob.get_backend()} + + +# ============================================================================= +# Speed suite +# ============================================================================= + +def _time_to_nll(model_name, nll_curve_fn, target_nll): + """Placeholder hook: time-to-same-NLL requires per-round eval; wired in + once the unified trainer exposes cheap incremental eval on GPU.""" + return None + + +def run_speed(quick: bool = False, use_gpu: bool = False) -> dict: + import numpy as np + + import openboost as ob + + # 1M hits the ≥1M acceptance gate. NGBoost exact-split trees took + # ~24 min at 45K; skip it past 100K. OpenBoost still times 500K/1M. + sizes = [50_000] if quick else [100_000, 500_000, 1_000_000] + sync = use_gpu + + # Warmup (few trees — JIT only) + Xw, yw = make_heteroscedastic(512, n_features=10, seed=0) + for _ in range(2): + ob.NaturalBoostNormal( + n_trees=5, learning_rate=LEARNING_RATE, max_depth=MAX_DEPTH, + n_bins=254, + ).fit(Xw, yw) + if sync: + from numba import cuda + cuda.synchronize() + + have_ngboost = True + try: + import ngboost # noqa: F401 + except ImportError: + have_ngboost = False + have_pgbm = True + try: + import pgbm # noqa: F401 + except ImportError: + have_pgbm = False + + results = [] + for n in sizes: + X, y = make_heteroscedastic(n) + split = int(0.9 * n) + X_tr, y_tr = X[:split], y[:split] + X_te, y_te = X[split:], y[split:] + row = {"n_train": split, "n_features": int(X.shape[1])} + print(f"size n={n:,} train={split:,} backend={ob.get_backend()}", + flush=True) + + # OpenBoost NaturalBoost + times = [] + n_trials = 1 if n >= 500_000 else 3 + for trial in range(n_trials): + model = make_openboost() + t0 = time.perf_counter() + model.fit(X_tr, y_tr) + if sync: + from numba import cuda + cuda.synchronize() + elapsed = time.perf_counter() - t0 + times.append(elapsed) + print(f" openboost trial {trial + 1}/{n_trials}: {elapsed:.1f}s", + flush=True) + times.sort() + row["openboost"] = {"fit_time_s": round(times[len(times) // 2], 2), + **score_model(model, X_te, y_te)} + + # NGBoost (CPU only) — skip past 100K: 45K already took ~24 min. + if have_ngboost and (quick or n <= 100_000): + print(f" ngboost starting (n={split:,})...", flush=True) + model = make_ngboost() + t0 = time.perf_counter() + model.fit(X_tr, y_tr) + row["ngboost"] = {"fit_time_s": round(time.perf_counter() - t0, 2), + **score_model(model, X_te, y_te)} + row["speedup_vs_ngboost"] = round( + row["ngboost"]["fit_time_s"] / row["openboost"]["fit_time_s"], 2) + + # PGBM (torch; GPU-capable) — optional + if have_pgbm: + try: + model = make_pgbm() + t0 = time.perf_counter() + model.fit(X_tr, y_tr) + row["pgbm"] = {"fit_time_s": round(time.perf_counter() - t0, 2), + **score_model(model, X_te, y_te)} + row["speedup_vs_pgbm"] = round( + row["pgbm"]["fit_time_s"] / row["openboost"]["fit_time_s"], 2) + except Exception as exc: + row["pgbm"] = {"error": f"{type(exc).__name__}: {exc}"} + + results.append(row) + parts = [f"n={split:,}", + f"OB {row['openboost']['fit_time_s']}s " + f"NLL {row['openboost']['nll']:.4f}"] + if "ngboost" in row: + parts.append(f"NGB {row['ngboost']['fit_time_s']}s " + f"({row['speedup_vs_ngboost']}x)") + if isinstance(row.get("pgbm"), dict) and "fit_time_s" in row["pgbm"]: + parts.append(f"PGBM {row['pgbm']['fit_time_s']}s " + f"({row['speedup_vs_pgbm']}x)") + print(" ".join(parts), flush=True) + + return {"suite": "speed", "results": results, + "backend": ob.get_backend(), + "ngboost_available": have_ngboost, + "pgbm_available": have_pgbm, + "ngboost_max_n": 100_000} + + +# ============================================================================= +# Runner / report +# ============================================================================= + +def run_suites(suite: str, quick: bool, use_gpu: bool) -> dict: + import numpy as np + + import openboost as ob + + report = { + "benchmark": "bench_probabilistic", + "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "seed": SEED, + "quick_mode": quick, + "use_gpu": use_gpu, + "budget": {"n_trees": N_TREES, "learning_rate": LEARNING_RATE, + "max_depth": MAX_DEPTH}, + "platform": { + "python": platform.python_version(), + "system": f"{platform.system()} {platform.machine()}", + "cpu_count": os.cpu_count(), + }, + "versions": {"openboost": ob.__version__, "numpy": np.__version__}, + "suites": {}, + } + try: + import ngboost + report["versions"]["ngboost"] = ngboost.__version__ + except ImportError: + pass + + if suite in ("quality", "all"): + report["suites"]["quality"] = run_quality(quick=quick) + if suite in ("speed", "all"): + report["suites"]["speed"] = run_speed(quick=quick, use_gpu=use_gpu) + return report + + +def save_report(report: dict, suite: str) -> Path: + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + out = RESULTS_DIR / f"probabilistic_{suite}_{stamp}.json" + out.write_text(json.dumps(report, indent=2) + "\n") + print(f"\nWrote {out}") + return out + + +# ============================================================================= +# Modal entry points +# ============================================================================= + +try: + import modal + + app = modal.App("openboost-probabilistic-bench") + image = ( + modal.Image.from_registry( + "nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12") + .pip_install( + "numpy>=1.24,<2.5", "numba>=0.60", "numba-cuda>=0.23", + "scipy>=1.10", "scikit-learn>=1.0", "ngboost>=0.5", + "pandas>=2.0", "joblib>=1.2", + ) + .add_local_dir( + str(PROJECT_ROOT / "src" / "openboost"), + remote_path="/root/openboost", + copy=True, + ) + .env({"PYTHONPATH": "/root", "OPENBOOST_BACKEND": "cuda"}) + ) +except ImportError: + modal = None + app = None + image = None + +if modal is not None and app is not None: + + @app.function(gpu="A100", image=image, timeout=4 * 3600) + def _run_remote(suite: str = "all", quick: bool = False): + sys.path.insert(0, "/root") + import openboost as ob + + ob.set_backend("cuda") + print(f"backend={ob.get_backend()} suite={suite} quick={quick}", + flush=True) + return run_suites(suite=suite, quick=quick, use_gpu=True) + + @app.local_entrypoint() + def main(suite: str = "all", quick: bool = False): + report = _run_remote.remote(suite=suite, quick=quick) + save_report(report, suite) + + +# ============================================================================= +# Local execution +# ============================================================================= + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--suite", choices=["quality", "speed", "all"], + default="quality") + parser.add_argument("--quick", action="store_true", + help="Smoke-test sizes/splits") + args = parser.parse_args() + + os.environ.setdefault("OPENBOOST_BACKEND", "cpu") + sys.path.insert(0, str(PROJECT_ROOT / "src")) + + report = run_suites(suite=args.suite, quick=args.quick, use_gpu=False) + save_report(report, args.suite) diff --git a/benchmarks/bench_survival.py b/benchmarks/bench_survival.py new file mode 100644 index 0000000..fa47e2d --- /dev/null +++ b/benchmarks/bench_survival.py @@ -0,0 +1,422 @@ +"""Survival capability benchmark: Weibull AFT with covariate-dependent shape. + +The second capability proof for "Boost every parameter". Data is Weibull AFT +where BOTH the scale lambda(z) AND the shape k(z) depend on covariates: + + T ~ Weibull(scale=lambda(z), shape=k(z)), right-censored at ~35%. + +Baselines: + + global constant (lambda, k) MLE (no covariate dependence). + xgb-aft XGBoost's built-in `survival:aft`. It learns only the location; + the distribution scale `aft_loss_distribution_scale` is a single + GLOBAL hyperparameter, so the Weibull shape k = 1/sigma is the SAME + for every row. It structurally cannot vary shape with covariates. + We set its global sigma to the MLE 1/k_global (its best setting). + +OpenBoost WeibullAFT boosts both lambda(z) and k(z) from the censored NLL, so +it is the only model that recovers the true shape surface. NGBoost has no +censored likelihood at all, so it cannot enter this benchmark. + +Metrics (identical closed forms for both libraries via the implied Weibull): + C-index (Harrell), censored NLL, 80% interval coverage, shape recovery + (corr of predicted vs true log k). + +Usage: + uv run --with xgboost python benchmarks/bench_survival.py + uv run --with xgboost python benchmarks/bench_survival.py --quick + uv run modal run benchmarks/bench_survival.py # A100 + +Output: benchmarks/results/survival_.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +PROJECT_ROOT = Path(__file__).parent.parent +RESULTS_DIR = Path(__file__).parent / "results" + +SEED = 42 +N_FEATURES = 6 +ROUNDS = 300 +LR = 0.1 +DEPTH = 3 +DAMP = 1.0 +CENSOR_Q = 0.65 # exponential censoring quantile -> ~35% censored +CINDEX_CAP = 4000 # subsample cap for the O(n^2) concordance + +# Gates +GATE_CINDEX_SLACK = 0.005 # OpenBoost C-index >= xgb - this +GATE_SHAPE_CORR = 0.5 # OpenBoost must recover the shape surface + + +# ============================================================================= +# Data generating process: varying scale AND shape +# ============================================================================= + +def make_data(n, rng): + Z = rng.uniform(0, 1, (n, N_FEATURES)) + lam = np.exp(0.5 + 1.0 * Z[:, 0] - 0.8 * Z[:, 1] + 0.4 * Z[:, 4]) + k = np.exp(0.2 + 1.2 * Z[:, 2] - 0.6 * Z[:, 3]) # shape varies with z + u = rng.uniform(1e-9, 1.0, n) + T = lam * (-np.log(u)) ** (1.0 / k) + C = rng.exponential(np.quantile(T, CENSOR_Q), n) + t = np.minimum(T, C) + event = (T <= C).astype(np.float64) + return Z, t, event, lam, k, T + + +# ============================================================================= +# Metrics (shared closed forms via the implied Weibull(lambda, k)) +# ============================================================================= + +def weibull_nll(t, event, lam, k): + lam = np.clip(lam, 1e-9, None) + k = np.clip(k, 1e-6, None) + m = np.log(t) - np.log(lam) + z = np.exp(np.clip(k * m, -60, 60)) + nll = -event * (np.log(k) - np.log(lam) + (k - 1.0) * m) + z + return float(np.mean(nll)) + + +def weibull_quantile(lam, k, q): + return lam * (-np.log(1.0 - q)) ** (1.0 / k) + + +def interval_coverage(true_T, lam, k, lo=0.1, hi=0.9): + """Coverage of the true (uncensored) event time by the [lo, hi] interval.""" + q_lo = weibull_quantile(lam, k, lo) + q_hi = weibull_quantile(lam, k, hi) + return float(np.mean((true_T >= q_lo) & (true_T <= q_hi))) + + +def c_index(time, event, risk, seed=SEED, cap=CINDEX_CAP): + """Harrell's concordance. Higher risk should mean earlier event.""" + n = len(time) + if n > cap: + idx = np.random.default_rng(seed).choice(n, cap, replace=False) + time, event, risk = time[idx], event[idx], risk[idx] + ti = time[:, None] + tj = time[None, :] + ei = event[:, None] + # comparable: i had an event and died strictly before j + comparable = (ti < tj) & (ei == 1) + ri = risk[:, None] + rj = risk[None, :] + concordant = (ri > rj) & comparable + tied = (ri == rj) & comparable + denom = comparable.sum() + if denom == 0: + return float("nan") + return float((concordant.sum() + 0.5 * tied.sum()) / denom) + + +# ============================================================================= +# Models +# ============================================================================= + +def run_openboost(Ztr, ttr, evtr, sets, rounds): + import openboost as ob + + m = ob.WeibullAFT(n_trees=rounds, max_depth=DEPTH, learning_rate=LR, damp=DAMP) + t0 = time.perf_counter() + m.fit(Ztr, ttr, event=evtr) + fit_time = time.perf_counter() - t0 + + out = {"fit_time_s": round(fit_time, 2)} + for name, (Z, t, ev, true_T, _lam_true, k_true) in sets.items(): + p = m.predict_params(Z) + lam, k = p["scale"], p["shape"] + med = weibull_quantile(lam, k, 0.5) + out[name] = { + "c_index": c_index(t, ev, -med), + "nll": weibull_nll(t, ev, lam, k), + "coverage_80": interval_coverage(true_T, lam, k), + "shape_corr": float(np.corrcoef(np.log(k), np.log(k_true))[0, 1]), + } + return out + + +def run_global(Ztr, ttr, evtr, sets): + from openboost._objectives import WeibullAFTObjective + + obj = WeibullAFTObjective() + init = obj.init_raw(ttr, None, {"event": evtr}) + lam_g, k_g = float(np.exp(init["scale"])), float(np.exp(init["shape"])) + out = {"lambda": lam_g, "shape": k_g} + for name, (_Z, t, ev, true_T, _lam, _k) in sets.items(): + lam = np.full(len(t), lam_g) + k = np.full(len(t), k_g) + med = weibull_quantile(lam, k, 0.5) + out[name] = { + "c_index": c_index(t, ev, -med), + "nll": weibull_nll(t, ev, lam, k), + "coverage_80": interval_coverage(true_T, lam, k), + "shape_corr": None, # constant shape: not defined + } + return out, k_g + + +def run_xgb_aft(Ztr, ttr, evtr, sets, rounds, k_global): + """XGBoost survival:aft with a single global scale sigma = 1/k_global.""" + try: + import xgboost as xgb + except ImportError as exc: + return {"error": f"xgboost unavailable: {exc}"} + + try: + sigma = float(1.0 / max(k_global, 1e-6)) + dtrain = xgb.DMatrix(np.ascontiguousarray(Ztr)) + upper = np.where(evtr == 1, ttr, np.inf) + dtrain.set_float_info("label_lower_bound", ttr) + dtrain.set_float_info("label_upper_bound", upper) + params = { + "objective": "survival:aft", + "eval_metric": "aft-nloglik", + # 'extreme' value log-time distribution == Weibull survival time. + "aft_loss_distribution": "extreme", + "aft_loss_distribution_scale": sigma, # GLOBAL shape, not per-row + "tree_method": "hist", + "max_depth": DEPTH, + "eta": LR, + "lambda": DAMP, + "seed": SEED, + } + t0 = time.perf_counter() + bst = xgb.train(params, dtrain, num_boost_round=rounds) + fit_time = time.perf_counter() - t0 + + k_const = 1.0 / sigma + out = {"fit_time_s": round(fit_time, 2), "global_shape": k_const, + "note": "single global scale; shape identical for all rows"} + for name, (Z, t, ev, true_T, _lam, _k_true) in sets.items(): + lam = np.asarray(bst.predict(xgb.DMatrix(np.ascontiguousarray(Z))), + dtype=np.float64) # AFT location = Weibull scale + k = np.full(len(t), k_const) + med = weibull_quantile(lam, k, 0.5) + out[name] = { + "c_index": c_index(t, ev, -med), + "nll": weibull_nll(t, ev, lam, k), + "coverage_80": interval_coverage(true_T, lam, k), + "shape_corr": None, # constant shape by construction + } + return out + except Exception as exc: # noqa: BLE001 - baseline must not crash the suite + import traceback + return {"error": f"{type(exc).__name__}: {exc}", + "traceback": traceback.format_exc()} + + +# ============================================================================= +# Gates +# ============================================================================= + +def compute_gates(res): + ob = res["openboost"]["test"] + glob = res["global"]["test"] + xgb = res.get("xgb", {}) + gates = { + "openboost_shape_corr": round(ob["shape_corr"], 3), + "recovers_shape_pass": ob["shape_corr"] >= GATE_SHAPE_CORR, + "beats_global_cindex_pass": ob["c_index"] > glob["c_index"], + "beats_global_nll_pass": ob["nll"] < glob["nll"], + } + if xgb and "error" not in xgb: + xt = xgb["test"] + gates["cindex_vs_xgb_delta"] = round(ob["c_index"] - xt["c_index"], 4) + gates["nll_vs_xgb_delta"] = round(ob["nll"] - xt["nll"], 4) + gates["cindex_not_worse_than_xgb_pass"] = ( + ob["c_index"] >= xt["c_index"] - GATE_CINDEX_SLACK) + gates["nll_better_than_xgb_pass"] = ob["nll"] < xt["nll"] + else: + gates["cindex_not_worse_than_xgb_pass"] = None + gates["nll_better_than_xgb_pass"] = None + + hard = [ + gates["recovers_shape_pass"], + gates["beats_global_cindex_pass"], + gates["beats_global_nll_pass"], + ] + if gates["cindex_not_worse_than_xgb_pass"] is not None: + hard.append(gates["cindex_not_worse_than_xgb_pass"]) + hard.append(gates["nll_better_than_xgb_pass"]) + gates["passed"] = all(hard) + return gates + + +def _pf(v): + return "skip" if v is None else ("PASS" if v else "FAIL") + + +# ============================================================================= +# Runner +# ============================================================================= + +def run_size(n_train, n_test, rounds): + import openboost as ob + + rng = np.random.default_rng(SEED) + Ztr, ttr, evtr, *_ = make_data(n_train, rng) + Zte, tte, evte, lam_te, k_te, T_te = make_data(n_test, rng) + sets = {"test": (Zte, tte, evte, T_te, lam_te, k_te)} + + print(f"\n=== n_train={n_train:,} rounds={rounds} " + f"backend={ob.get_backend()} censor={1 - evtr.mean():.2f} ===", + flush=True) + + res = {"n_train": n_train, "n_test": n_test, "rounds": rounds} + res["openboost"] = run_openboost(Ztr, ttr, evtr, sets, rounds) + o = res["openboost"]["test"] + print(f" OpenBoost WeibullAFT C-index {o['c_index']:.4f} " + f"NLL {o['nll']:.4f} cov80 {o['coverage_80']:.3f} " + f"shape_corr {o['shape_corr']:.3f} ({res['openboost']['fit_time_s']}s)", + flush=True) + + res["global"], k_global = run_global(Ztr, ttr, evtr, sets) + g = res["global"]["test"] + print(f" global C-index {g['c_index']:.4f} " + f"NLL {g['nll']:.4f} cov80 {g['coverage_80']:.3f} shape_corr n/a", + flush=True) + + res["xgb"] = run_xgb_aft(Ztr, ttr, evtr, sets, rounds, k_global) + if "error" in res["xgb"]: + print(f" xgb survival:aft SKIPPED ({res['xgb']['error']})", flush=True) + else: + x = res["xgb"]["test"] + print(f" xgb survival:aft C-index {x['c_index']:.4f} " + f"NLL {x['nll']:.4f} cov80 {x['coverage_80']:.3f} " + f"shape_corr n/a (global k={res['xgb']['global_shape']:.2f}) " + f"({res['xgb']['fit_time_s']}s)", flush=True) + + res["gates"] = compute_gates(res) + gt = res["gates"] + print(f" GATES recovers_shape({gt['openboost_shape_corr']})=" + f"{_pf(gt['recovers_shape_pass'])} " + f"NLL=xgb={_pf(gt['cindex_not_worse_than_xgb_pass'])} " + f"beats_global={_pf(gt['beats_global_nll_pass'])} " + f"=> {'PASS' if gt['passed'] else 'FAIL'}", flush=True) + return res + + +def run_suite(quick=False): + import openboost as ob + + sizes = ([(8_000, 4_000, 80)] if quick + else [(40_000, 10_000, ROUNDS), (200_000, 20_000, ROUNDS)]) + + # Warmup (JIT the tree kernels off the clock). + rng = np.random.default_rng(0) + Zw, tw, evw, *_ = make_data(512, rng) + ob.WeibullAFT(n_trees=5, max_depth=DEPTH).fit(Zw, tw, event=evw) + + report = { + "benchmark": "bench_survival", + "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "seed": SEED, + "quick_mode": quick, + "model": "Weibull AFT, scale(z) and shape(z) both boosted", + "budget": {"rounds": ROUNDS, "lr": LR, "depth": DEPTH, "damp": DAMP, + "censor_q": CENSOR_Q}, + "platform": { + "python": platform.python_version(), + "system": f"{platform.system()} {platform.machine()}", + }, + "backend": ob.get_backend(), + "versions": {"openboost": ob.__version__, "numpy": np.__version__}, + "sizes": [], + } + try: + import xgboost + report["versions"]["xgboost"] = xgboost.__version__ + except ImportError: + report["versions"]["xgboost"] = None + + for n_train, n_test, rounds in sizes: + report["sizes"].append(run_size(n_train, n_test, rounds)) + + report["passed"] = all(s["gates"]["passed"] for s in report["sizes"]) + return report + + +def save_report(report): + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + out = RESULTS_DIR / f"survival_{stamp}.json" + out.write_text(json.dumps(report, indent=2) + "\n") + print(f"\nWrote {out}") + return out + + +# ============================================================================= +# Modal entry point +# ============================================================================= + +try: + import modal + + app = modal.App("openboost-survival-bench") + image = ( + modal.Image.from_registry( + "nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12") + .pip_install( + "numpy>=1.24,<2.5", "numba>=0.60", "numba-cuda>=0.23", + "scipy>=1.10", "scikit-learn>=1.0", "xgboost>=2.0", + "joblib>=1.2", + ) + .add_local_dir( + str(PROJECT_ROOT / "src" / "openboost"), + remote_path="/root/openboost", + copy=True, + ) + .env({"PYTHONPATH": "/root", "OPENBOOST_BACKEND": "cuda"}) + ) +except ImportError: + modal = None + app = None + image = None + +if modal is not None and app is not None: + + @app.function(gpu="A100", image=image, timeout=2 * 3600) + def _run_remote(quick: bool = False): + sys.path.insert(0, "/root") + import openboost as ob + + ob.set_backend("cuda") + print(f"backend={ob.get_backend()} quick={quick}", flush=True) + return run_suite(quick=quick) + + @app.local_entrypoint() + def main(quick: bool = False): + report = _run_remote.remote(quick=quick) + save_report(report) + print(f"\nOverall: {'PASS' if report['passed'] else 'FAIL'}") + + +# ============================================================================= +# Local execution +# ============================================================================= + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--quick", action="store_true", + help="Smoke-test sizes/rounds") + args = parser.parse_args() + + os.environ.setdefault("OPENBOOST_BACKEND", "cpu") + sys.path.insert(0, str(PROJECT_ROOT / "src")) + + report = run_suite(quick=args.quick) + save_report(report) + print(f"\nOverall: {'PASS' if report['passed'] else 'FAIL'}") diff --git a/development/paramboost/spike_ggn.py b/development/paramboost/spike_ggn.py new file mode 100644 index 0000000..addba44 --- /dev/null +++ b/development/paramboost/spike_ggn.py @@ -0,0 +1,315 @@ +"""Spike: generalized Gauss-Newton (GGN) parametric boosting on a custom formula. + +Model: sales y = a(z) * x ** sigmoid(b(z) * x) + noise + a = exp(F_a(z)) (log link, positive ceiling parameter) + b = F_b(z) (identity link, saturation-speed parameter) + +Each parameter channel F_k(z) is a boosting ensemble over feature vector z; +x (spend/price) enters only through the user formula. + +Question this spike answers: is GGN preconditioning + tree pooling stable and +necessary? We compare three per-sample update rules for d = "step direction +the trees are fit to": + + plain : d = g (raw gradient, single learning rate) + diag : d = g / (diag(G) + lam) (per-channel GGN scaling) + full : d = (G + lam*I)^{-1} g (full 2x2 GGN solve, off-diagonal) + +where g = dL/dF and G = J^T J (MSE loss => GGN middle matrix is identity). + +Baselines: + global : one (a, b) fit to all samples by Adam (no per-sample params) + blackbox: ob.GradientBoosting on [z, x] -> y (no formula) + +Metrics: + - in-range test RMSE (new z, x within training range) + - extrapolation RMSE (new z, x beyond training range: [3, 5]) + - parameter recovery (RMSE of a_hat vs a_true, b_hat vs b_true) + +Derivatives (analytic; JAX unavailable on this machine, same math): + s = sigmoid(b*x) + f = exp(u + s*ln x) with u = F_a, b = F_b + df/du = f + df/db = f * ln x * s*(1-s) * x + +Usage: uv run python development/paramboost/spike_ggn.py [--quick] +Output: development/paramboost/spike_ggn_results.json +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np + +import openboost as ob + +SEED = 42 +N_TRAIN = 40_000 +N_TEST = 10_000 +N_FEATURES = 6 +ROUNDS = 300 +LR = 0.1 +DEPTH = 3 +DAMP = 1.0 +X_TRAIN_RANGE = (0.25, 2.5) +X_EXTRAP_RANGE = (3.0, 5.0) +NOISE_FRAC = 0.05 # additive noise sd as fraction of signal sd + + +# ============================================================================= +# Data generating process +# ============================================================================= + +def sigmoid(t): + return 1.0 / (1.0 + np.exp(-t)) + + +def true_params(Z): + """Ground-truth parameter surfaces a(z), b(z).""" + u = 0.5 + 1.0 * Z[:, 0] - 0.8 * Z[:, 1] + 0.5 * np.sin(2 * np.pi * Z[:, 2]) + a = np.exp(u) + b = 0.5 + 2.0 * Z[:, 3] + Z[:, 4] * Z[:, 0] + return a, b + + +def formula(a, b, x): + """sales = a * x ** sigmoid(b*x), computed stably in log space.""" + s = sigmoid(b * x) + return np.exp(np.log(a) + s * np.log(x)) + + +def make_data(n, rng, x_range): + Z = rng.uniform(0, 1, (n, N_FEATURES)) + x = rng.uniform(*x_range, n) + a, b = true_params(Z) + f = formula(a, b, x) + y = f + NOISE_FRAC * f.std() * rng.standard_normal(n) + return Z, x, y, a, b, f + + +# ============================================================================= +# Per-sample gradients / GGN for the formula under MSE loss +# ============================================================================= + +def grads_and_ggn(u, b, x, y): + """Return g (n,2), G entries (n,3): [G_uu, G_bb, G_ub].""" + lnx = np.log(x) + s = sigmoid(b * x) + f = np.exp(np.clip(u + s * lnx, -30.0, 30.0)) + fu = f + fb = f * lnx * s * (1.0 - s) * x + r = f - y # dL/df for L = 0.5*(f-y)^2 + g = np.stack([r * fu, r * fb], axis=1) + G_uu = fu * fu + G_bb = fb * fb + G_ub = fu * fb + return g, G_uu, G_bb, G_ub, f + + +def step_direction(mode, g, G_uu, G_bb, G_ub, damp): + """Compute d per sample for the chosen preconditioning mode.""" + if mode == "plain": + return g + if mode == "diag": + return g / np.stack([G_uu + damp, G_bb + damp], axis=1) + # full 2x2 solve, analytic inverse with damping + det = (G_uu + damp) * (G_bb + damp) - G_ub * G_ub + d_u = ((G_bb + damp) * g[:, 0] - G_ub * g[:, 1]) / det + d_b = ((G_uu + damp) * g[:, 1] - G_ub * g[:, 0]) / det + return np.stack([d_u, d_b], axis=1) + + +# ============================================================================= +# Global curve fit (baseline + boosting init) +# ============================================================================= + +def fit_global(x, y, iters=2000, lr=0.05): + """Adam on scalar (u, b) minimizing mean 0.5*(f-y)^2.""" + u, b = float(np.log(np.mean(y))), 1.0 + m = np.zeros(2) + v = np.zeros(2) + for t in range(1, iters + 1): + g, *_ = grads_and_ggn(np.full_like(x, u), np.full_like(x, b), x, y) + gm = g.mean(axis=0) + m = 0.9 * m + 0.1 * gm + v = 0.999 * v + 0.001 * gm * gm + mh = m / (1 - 0.9**t) + vh = v / (1 - 0.999**t) + u -= lr * mh[0] / (np.sqrt(vh[0]) + 1e-8) + b -= lr * mh[1] / (np.sqrt(vh[1]) + 1e-8) + return u, b + + +# ============================================================================= +# Parametric booster +# ============================================================================= + +def fit_paramboost(mode, Z_tr, x_tr, y_tr, eval_sets, rounds=ROUNDS, + lr=LR, depth=DEPTH, damp=DAMP): + """Boost F_u, F_b on binned Z. Returns dict with curves and final F fns. + + eval_sets: {name: (Z_binned_data, x, y)} evaluated every 10 rounds. + """ + ba = ob.array(Z_tr) + Zb = np.ascontiguousarray(ba.data) + + u0, b0 = fit_global(x_tr, y_tr) + F_u = np.full(len(y_tr), u0, dtype=np.float64) + F_b = np.full(len(y_tr), b0, dtype=np.float64) + + # Per-eval-set running ensemble predictions (base + sum of tree outputs) + eval_state = { + name: { + "Zb": np.ascontiguousarray(ba.transform(Z).data), + "F_u": np.full(len(y), u0, dtype=np.float64), + "F_b": np.full(len(y), b0, dtype=np.float64), + "x": x, "y": y, + } + for name, (Z, x, y) in eval_sets.items() + } + + curves = {name: {} for name in eval_sets} + diverged = False + t0 = time.perf_counter() + for r in range(rounds): + g, G_uu, G_bb, G_ub, f = grads_and_ggn(F_u, F_b, x_tr, y_tr) + if not np.isfinite(g).all() or f.max() > 1e12: + diverged = True + break + d = step_direction(mode, g, G_uu, G_bb, G_ub, damp) + + ones = np.ones(len(y_tr), dtype=np.float64) + for ch, F in ((0, F_u), (1, F_b)): + # fit_tree leaf value = -sum(grad)/(sum(hess)+reg); grad=d, hess=1 + # gives leaf = -mean(d) so F += lr*tree is a descent step. + tree = ob.fit_tree(ba, np.ascontiguousarray(d[:, ch]), ones, + max_depth=depth) + F += lr * tree(Zb) + for st in eval_state.values(): + if ch == 0: + st["F_u"] += lr * tree(st["Zb"]) + else: + st["F_b"] += lr * tree(st["Zb"]) + + if (r + 1) % 10 == 0 or r == 0: + for name, st in eval_state.items(): + pred = formula(np.exp(st["F_u"]), st["F_b"], st["x"]) + curves[name][r + 1] = float( + np.sqrt(np.mean((pred - st["y"]) ** 2))) + + fit_time = time.perf_counter() - t0 + return { + "curves": curves, + "eval_state": eval_state, + "diverged": diverged, + "rounds_done": r + 1 if not diverged else r, + "fit_time_s": fit_time, + "init": (u0, b0), + } + + +# ============================================================================= +# Runner +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--quick", action="store_true") + args = parser.parse_args() + + global N_TRAIN, N_TEST, ROUNDS + if args.quick: + N_TRAIN, N_TEST, ROUNDS = 5_000, 2_000, 60 + + rng = np.random.default_rng(SEED) + Z_tr, x_tr, y_tr, _, _, _ = make_data(N_TRAIN, rng, X_TRAIN_RANGE) + Z_te, x_te, y_te, a_te, b_te, f_te = make_data(N_TEST, rng, X_TRAIN_RANGE) + # Extrapolation: same z distribution, x beyond training range, NOISELESS + # target (we compare against the true curve). + Z_ex, x_ex, _, a_ex, b_ex, f_ex = make_data(N_TEST, rng, X_EXTRAP_RANGE) + + eval_sets = {"test": (Z_te, x_te, y_te), "extrap": (Z_ex, x_ex, f_ex)} + results = {} + + # --- Parametric boosting variants --- + for mode in ["plain", "diag", "full"]: + out = fit_paramboost(mode, Z_tr, x_tr, y_tr, eval_sets, rounds=ROUNDS) + st_te = out["eval_state"]["test"] + a_hat = np.exp(st_te["F_u"]) + b_hat = st_te["F_b"] + test_curve = out["curves"]["test"] + best_r = min(test_curve, key=test_curve.get) if test_curve else None + results[mode] = { + "diverged": out["diverged"], + "rounds_done": out["rounds_done"], + "fit_time_s": round(out["fit_time_s"], 2), + "test_rmse_final": test_curve.get(ROUNDS), + "test_rmse_best": test_curve[best_r] if best_r else None, + "best_round": best_r, + "extrap_rmse_final": out["curves"]["extrap"].get(ROUNDS), + "param_rmse_a": float(np.sqrt(np.mean((a_hat - a_te) ** 2))), + "param_rmse_b": float(np.sqrt(np.mean((b_hat - b_te) ** 2))), + "param_corr_a": float(np.corrcoef(a_hat, a_te)[0, 1]), + "param_corr_b": float(np.corrcoef(b_hat, b_te)[0, 1]), + } + r = results[mode] + print(f"{mode:>6}: diverged={r['diverged']} " + f"test RMSE best {r['test_rmse_best']}@{r['best_round']} " + f"final {r['test_rmse_final']} extrap {r['extrap_rmse_final']} " + f"a: rmse {r['param_rmse_a']:.3f} corr {r['param_corr_a']:.3f} " + f"b: rmse {r['param_rmse_b']:.3f} corr {r['param_corr_b']:.3f} " + f"({r['fit_time_s']}s)") + + # --- Baseline: global curve fit --- + u_g, b_g = fit_global(x_tr, y_tr) + pred_te = formula(np.exp(u_g), b_g, x_te) + pred_ex = formula(np.exp(u_g), b_g, x_ex) + a_true_mean, b_true_mean = true_params(Z_te) + results["global"] = { + "test_rmse": float(np.sqrt(np.mean((pred_te - y_te) ** 2))), + "extrap_rmse": float(np.sqrt(np.mean((pred_ex - f_ex) ** 2))), + "param_rmse_a": float(np.sqrt(np.mean((np.exp(u_g) - a_true_mean) ** 2))), + "param_rmse_b": float(np.sqrt(np.mean((b_g - b_true_mean) ** 2))), + } + print(f"global: test RMSE {results['global']['test_rmse']:.4f} " + f"extrap {results['global']['extrap_rmse']:.4f}") + + # --- Baseline: black-box GBDT on [z, x] --- + X_tr = np.column_stack([Z_tr, x_tr]) + gb = ob.GradientBoosting(n_trees=ROUNDS, max_depth=6, learning_rate=0.1, + random_state=SEED) + t0 = time.perf_counter() + gb.fit(X_tr, y_tr) + gb_time = time.perf_counter() - t0 + pred_te = gb.predict(np.column_stack([Z_te, x_te])) + pred_ex = gb.predict(np.column_stack([Z_ex, x_ex])) + results["blackbox"] = { + "fit_time_s": round(gb_time, 2), + "test_rmse": float(np.sqrt(np.mean((pred_te - y_te) ** 2))), + "extrap_rmse": float(np.sqrt(np.mean((pred_ex - f_ex) ** 2))), + } + print(f"blackbox: test RMSE {results['blackbox']['test_rmse']:.4f} " + f"extrap {results['blackbox']['extrap_rmse']:.4f}") + + # Reference: irreducible noise floor on the test set + results["noise_floor_test_rmse"] = float(NOISE_FRAC * f_te.std()) + results["config"] = { + "n_train": N_TRAIN, "n_test": N_TEST, "rounds": ROUNDS, "lr": LR, + "depth": DEPTH, "damp": DAMP, "seed": SEED, + "x_train_range": X_TRAIN_RANGE, "x_extrap_range": X_EXTRAP_RANGE, + "noise_frac": NOISE_FRAC, + } + print(f"noise floor (test RMSE lower bound): " + f"{results['noise_floor_test_rmse']:.4f}") + + out_path = Path(__file__).parent / "spike_ggn_results.json" + out_path.write_text(json.dumps(results, indent=2) + "\n") + print(f"wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/development/paramboost/spike_ggn_results.json b/development/paramboost/spike_ggn_results.json new file mode 100644 index 0000000..3006b55 --- /dev/null +++ b/development/paramboost/spike_ggn_results.json @@ -0,0 +1,71 @@ +{ + "plain": { + "diverged": false, + "rounds_done": 300, + "fit_time_s": 66.41, + "test_rmse_final": 1.5626822535612688, + "test_rmse_best": 1.0519128296422544, + "best_round": 10, + "extrap_rmse_final": 3.902738643156501, + "param_rmse_a": 1.0204866560925787, + "param_rmse_b": 0.6026993198569472, + "param_corr_a": 0.5629250077210485, + "param_corr_b": 0.6496257259312742 + }, + "diag": { + "diverged": false, + "rounds_done": 300, + "fit_time_s": 65.37, + "test_rmse_final": 0.13401843837542263, + "test_rmse_best": 0.13401843837542263, + "best_round": 300, + "extrap_rmse_final": 0.2002377135310078, + "param_rmse_a": 0.05037521823023994, + "param_rmse_b": 0.5327583587955843, + "param_corr_a": 0.998967848285634, + "param_corr_b": 0.7020771854857124 + }, + "full": { + "diverged": false, + "rounds_done": 300, + "fit_time_s": 66.4, + "test_rmse_final": 0.14511242548544248, + "test_rmse_best": 0.14511242548544248, + "best_round": 300, + "extrap_rmse_final": 0.20271665109258358, + "param_rmse_a": 0.05542024184891472, + "param_rmse_b": 0.6462846398404533, + "param_corr_a": 0.9987569217788413, + "param_corr_b": 0.8397744392167091 + }, + "global": { + "test_rmse": 1.6725387930519824, + "extrap_rmse": 4.374675453926745, + "param_rmse_a": 1.1030746499952404, + "param_rmse_b": 0.6702363041288179 + }, + "blackbox": { + "fit_time_s": 48.09, + "test_rmse": 0.14461229974048906, + "extrap_rmse": 3.9102367405351384 + }, + "noise_floor_test_rmse": 0.10361939551132492, + "config": { + "n_train": 40000, + "n_test": 10000, + "rounds": 300, + "lr": 0.1, + "depth": 3, + "damp": 1.0, + "seed": 42, + "x_train_range": [ + 0.25, + 2.5 + ], + "x_extrap_range": [ + 3.0, + 5.0 + ], + "noise_frac": 0.05 + } +} diff --git a/planning/unified-engine-design.md b/planning/unified-engine-design.md new file mode 100644 index 0000000..b2d9782 --- /dev/null +++ b/planning/unified-engine-design.md @@ -0,0 +1,184 @@ +# Unified GPU Boosting Engine — Design + +Status: Phase A+B implemented; A100 speed gate passed (1229x vs NGBoost at 90K) · 2026-08-17 +Evidence: `development/paramboost/spike_ggn.py` (GGN spike), CustomDistribution/loop +survey, NaturalBoost GPU-boundary survey. + +## 1. Goal + +One training core in which every model is a configuration, not a fork: + +``` +z --trees--> F (n,K raw scores) --link--> theta --formula f--> eta --loss L--> scalar +``` + +Per-round update rule (the only one in the engine), generalized Gauss-Newton: + +``` +d_i = (J_i^T M_i J_i + damp*I)^{-1} g_i # K x K per sample, K small +``` + +| Model | f | L | M | K | +|---|---|---|---|---| +| GradientBoosting | identity | mse/logloss/... | scalar hessian | 1 | +| NaturalBoost | identity | NLL | analytic Fisher | n_params | +| CustomDistribution | identity | user NLL (JAX/FD) | GGN (upgrade from empirical diag) | n_params | +| FormulaBoost (new) | user formula f(x;theta) | mse or NLL | GGN, modes plain/diag/full | n_formula_params | +| Weibull AFT (new) | identity | censored NLL | Fisher/GGN | 2 | + +Spike evidence (sales formula `a*x^sigmoid(bx)`, 40K samples): +plain gradient unusable (RMSE 1.05, degrades); diag GGN 0.134; full GGN 0.145 +with much better parameter recovery (b corr 0.84 vs 0.70); extrapolation 19x +better than black-box GBDT (0.20 vs 3.91). Preconditioning is load-bearing. + +## 2. Current state (survey summary) + +- Two disjoint training loops: `_boosting.py` (1 channel, full GPU story: + `_fit_gpu`, device-resident buffers, `fit_tree_gpu_native`, growth/GOSS) and + `_distributional.py` (K channels, CPU-only gradients AND host `fit_tree`, + richer eval: multi-set, incremental raw scores, `evals_result_`). +- `CustomDistribution`: user NLL + links, JAX autodiff w/ FD fallback, + empirical **diagonal** Fisher only, no formula layer, quantile/sample assume + Normal. +- No survival/censoring anywhere. CRPS is eval-only. No line search. +- JAX is optional (`[jax]` extra); guarded import; macOS-Intel local dev has no + working jaxlib -> FD fallback must remain first-class. + +## 3. Target architecture + +### 3.1 Objective protocol (new, internal) + +```python +class Objective(Protocol): + n_channels: int + def init_state(y, sample_weight) -> RawState # base scores (n,K) + def step(F, y, sw) -> tuple[d, h] # per-channel step dirs + def loss_value(F, y) -> float # for eval/early stop + device_capable: bool # can consume device arrays +``` + +Implementations wrap what exists today: + +- `LossObjective` — K=1, delegates to `_loss.py` losses (incl. GPU in-place + kernels). `d = g`, `h = hess` (leaf Newton step unchanged; this is the K=1 + degenerate case of GGN). +- `DistributionObjective` — wraps `Distribution.nll_gradient` + + `natural_gradient` (analytic Fisher). CUDA kernels per family, ported + incrementally (Normal -> LogNormal -> Poisson; digamma families stay CPU + until worth it). +- `CustomLikelihoodObjective` — JAX (vmap grad + per-sample GGN) or FD + fallback. Upgrade from batch-mean diagonal Fisher to per-sample GGN with + damping (spike-validated). JAX-on-GPU interops with numba-cuda via dlpack. +- `FormulaObjective` — FormulaBoost: user `f(theta, x)` + loss; J via + autodiff or user-supplied; modes `precond='diag'|'full'` (default full, + spike-validated), damping `lam`, global-fit initialization for base scores. + +### 3.2 Unified trainer (one loop) + +Single `fit_boosting(objective, X, y, ...)` used by all model facades: + +- bin once (`ob.array`), upload once when backend=cuda +- state `F`: (n, K) — device-resident when backend=cuda and + `objective.device_capable` +- per round: `(d, h) = objective.step(F, y)`; then per channel k: + `fit_tree_gpu_native` (cuda) / `fit_tree` (cpu); `F[:,k] += lr * tree(X)` + in-place on device +- eval: incremental raw-score updates for all eval sets (adopt the better + `_distributional.py` pattern), `evals_result_`, callbacks, early stopping, + `early_stopping_rounds` sugar for every model +- D2H only at eval/checkpoint boundaries + +Non-goals for the shared loop v1: multi-GPU/Ray (stays on the K=1 fast path), +GOSS for K>1, sample_weight on CUDA (unchanged limitation). + +### 3.3 Facades (no backward-compat constraint — zero users today, design the +ideal API and delete what it replaces) + +- `GradientBoosting` — keeps its current fast paths in Phase A for perf + reasons only (gate: no regression on `benchmarks/check_performance.py`), + not for API stability; rename/reshape freely if the unified trainer wins. +- `NaturalBoost*` / `DistributionalGBDT` — rebuilt on the trainer; the old + `_distributional.py` loop is deleted, not kept alongside. API may change + where the unified design is cleaner. +- `FormulaBoost` — new. (Naming: "formula in, boosted parameters out"; the + statistical anchor for docs is varying-coefficient / semi-parametric + modeling, since x follows the user formula while theta(z) is non-parametric + via trees.) + +```python +def curve(theta, x): # jax.numpy inside for autodiff + a, b = theta + return a * x ** jax.nn.sigmoid(b * x) + +m = ob.FormulaBoost(model=curve, n_params=2, links=("log", "identity"), + loss="mse", precond="full", damp=1.0) +m.fit(Z, y, model_input=x, eval_set=[...], callbacks=[...]) +m.predict_params(Z) # (a, b) per sample — the deliverable +m.predict(Z, model_input=x_new) # counterfactual / extrapolation +``` + +- `register_loss` / `register_distribution` / `register_growth_strategy` + unchanged; add `register_objective` later only if needed. + +### 3.4 Vector-leaf (research track, structural option) + +The trainer's "per channel k: fit one tree" block becomes a strategy; a +vector-leaf strategy (one tree, K-dim leaves) slots in without touching +objectives. Blocked on closing the onetree research gap (alpha-tempering +showed promise); not on any v1 critical path. + +## 4. NGBoost feature-parity decisions + +| Gap | Decision | +|---|---| +| CRPS as training score | Defer. NLL-only training; CRPS stays an eval metric. Revisit if users ask. | +| Line search | Won't do. Damped GGN + lr replaces it (spike: stable without). Document the difference. | +| Survival / censored | Do. Weibull AFT as `DistributionObjective` (censored NLL, event indicator via `sample_weight`-style arg). Needed for capability benchmark #1. | +| Multivariate distributions | Defer (post-1.x). | + +## 5. Phasing (eval-first) + +- **0. Eval harness before any engine work** — the yardstick every later + phase is measured against: + - `benchmarks/bench_probabilistic.py`: quality suite (UCI 9 + California + + YearMSD; NLL/CRPS/RMSE/coverage/pinball vs NGBoost, paired splits) + + speed suite (Modal A100; wall-clock + time-to-same-NLL vs NGBoost/PGBM), + one command, JSON results in `benchmarks/results/` + - Capability evals defined *before* FormulaBoost exists: sales-formula + benchmark (generalize the spike: known-truth param recovery, in-range + + extrapolation RMSE, vs global-fit / black-box / hand-rolled XGBoost) and + Weibull AFT dataset + metric prep (C-index, calibration) + - Golden parity fixtures: current NaturalBoost/DistributionalGBDT outputs + at fixed seeds, so the Phase A refactor has a regression baseline + - **Run it once on current main** -> committed "before" numbers; every + later phase must move these numbers, not anecdotes +- **A. Trainer + objectives, CPU-correct** — extract unified loop; port + DistributionalGBDT/NaturalBoost onto it; `FormulaObjective` + + `FormulaBoost` facade (FD fallback works everywhere); behavioral parity + tests vs current models (same seeds, same NLL within 1e-6). +- **B. GPU end-to-end for K>1** — Normal/LogNormal/Poisson gradient+Fisher + CUDA kernels; device-resident F; `fit_tree_gpu_native` per channel; A100 + benchmark vs NGBoost (target >=10x at >=1M rows). +- **C. JAX-on-GPU custom objectives** — dlpack zero-copy bridge; FormulaBoost + and CustomDistribution get the GPU path. +- **D. Weibull AFT + capability benchmarks** — survival objective; the two + capability benchmarks (AFT vs XGBoost AFT; sales-formula vs hand-rolled). + +Each phase lands independently with a numeric gate from Phase 0: +A = golden parity (same seeds, NLL within 1e-6) + no perf regression; +B = >=10x vs NGBoost at >=1M rows, quality suite unchanged; +C = custom-objective GPU path beats its own CPU path, results identical; +D = capability benchmarks won (AFT vs XGBoost AFT; formula vs baselines). + +## 6. Risks / open questions + +- FD fallback cost for K-param GGN (2K+1 NLL evals/sample/round) — acceptable + for small K; document JAX as the fast path. +- Numerical parity CPU vs GPU for Fisher math (float32 on device vs float64 + numpy today) — decide dtype policy in Phase B; eval NLL stays CPU/float64. +- `GradientBoosting` K=1 fast path: keep only as long as it is measurably + faster than the unified trainer; fold in and delete once the trainer + matches it. No compat reason to keep two loops. +- Golden parity fixtures (Phase 0) are a correctness tool for the refactor, + not an API-stability promise; intentional behavior changes just update the + fixtures.