Gradient Boosted Trees

"""Gradient boosting with depth-1 trees (stumps) - the loop, written out.

Each round: compute the residual, fit a stump to it, add a shrunken version to
the running prediction. That is the entire algorithm.
"""


def best_stump(X, residuals):
    """Find the (feature, threshold) split minimizing squared error on the
    residuals. Brute force over every candidate threshold."""
    n, d = len(X), len(X[0])
    best = (None, None, 0.0, 0.0, float("inf"))

    for feature in range(d):
        thresholds = sorted({row[feature] for row in X})
        for threshold in thresholds:
            left = [r for x, r in zip(X, residuals) if x[feature] <= threshold]
            right = [r for x, r in zip(X, residuals) if x[feature] > threshold]
            if not left or not right:
                continue

            left_value = sum(left) / len(left)
            right_value = sum(right) / len(right)
            error = sum((r - left_value) ** 2 for r in left) + sum(
                (r - right_value) ** 2 for r in right
            )
            if error < best[4]:
                best = (feature, threshold, left_value, right_value, error)

    return best[:4]


def fit(X, y, rounds=100, learning_rate=0.1):
    """Returns the base prediction plus the list of stumps."""
    base = sum(y) / len(y)
    predictions = [base] * len(y)
    trees = []

    for _ in range(rounds):
        # Negative gradient of squared error is just the residual.
        residuals = [target - pred for target, pred in zip(y, predictions)]

        feature, threshold, left_value, right_value = best_stump(X, residuals)
        if feature is None:
            break
        trees.append((feature, threshold, left_value, right_value))

        # Shrinkage: add only a fraction, so later trees still have work left.
        for i, row in enumerate(X):
            step = left_value if row[feature] <= threshold else right_value
            predictions[i] += learning_rate * step

    return base, trees


def predict(row, base, trees, learning_rate=0.1):
    total = base
    for feature, threshold, left_value, right_value in trees:
        total += learning_rate * (left_value if row[feature] <= threshold else right_value)
    return total

"""Gradient boosting - typed, early-stopped, with the loss chosen explicitly."""

from dataclasses import dataclass
from typing import Literal

import numpy as np
from numpy.typing import NDArray
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.model_selection import TimeSeriesSplit

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]

Objective = Literal["squared_error", "absolute_error", "poisson", "quantile"]


@dataclass(frozen=True)
class BoostingResult:
    model: HistGradientBoostingRegressor
    best_iteration: int
    n_estimators_requested: int

    @property
    def stopped_early(self) -> bool:
        """If it never stopped early, the round budget was the binding
        constraint - which means the search never found the optimum."""
        return self.best_iteration < self.n_estimators_requested - 1


def fit_forecaster(
    X: Matrix,
    y: Vector,
    objective: Objective = "squared_error",
    quantile: float | None = None,
    max_rounds: int = 2_000,
    learning_rate: float = 0.05,
) -> BoostingResult:
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"X has {X.shape[0]} rows but y has {y.shape[0]}")
    if objective == "quantile" and quantile is None:
        raise ValueError("quantile objective requires a quantile value")

    model = HistGradientBoostingRegressor(
        loss=objective,
        quantile=quantile,
        learning_rate=learning_rate,
        max_iter=max_rounds,
        # Early stopping is mandatory, not optional: training loss falls
        # monotonically forever and gives no signal about when to stop.
        early_stopping=True,
        n_iter_no_change=50,
        validation_fraction=0.15,
        random_state=0,
    )
    model.fit(X, y)

    return BoostingResult(
        model=model,
        best_iteration=int(model.n_iter_),
        n_estimators_requested=max_rounds,
    )


def timeseries_cv_score(X: Matrix, y: Vector, folds: int = 5) -> float:
    """Rolling-origin CV.

    TimeSeriesSplit, never KFold: a random split trains on the future and
    validates on the past, which inflates the score and picks far too many
    boosting rounds.
    """
    splitter = TimeSeriesSplit(n_splits=folds)
    errors = []

    for train_idx, test_idx in splitter.split(X):
        result = fit_forecaster(X[train_idx], y[train_idx])
        predictions = result.model.predict(X[test_idx])
        errors.append(float(np.mean(np.abs(predictions - y[test_idx]))))

    return float(np.mean(errors))

"""Gradient boosting - pre-binned features, native categoricals, one Dataset reuse."""

import lightgbm as lgb
import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


def build_dataset(
    X: Matrix, y: Vector, categorical_indices: list[int] | None = None
) -> lgb.Dataset:
    """Bin once, reuse everywhere.

    Binning is the expensive preprocessing step, and a Dataset built once can
    back every model in a hyperparameter sweep. Rebuilding it per trial - the
    default when passing raw arrays - repeats that cost on every fit.
    """
    return lgb.Dataset(
        np.ascontiguousarray(X, dtype=np.float32),   # float32 halves the bin scan traffic
        label=y,
        categorical_feature=categorical_indices or [],
        free_raw_data=True,   # drop the raw copy once binned
    )


def fit_quantiles(
    train: lgb.Dataset,
    valid: lgb.Dataset,
    quantiles: tuple[float, ...] = (0.1, 0.5, 0.9),
    rounds: int = 5_000,
) -> dict[float, lgb.Booster]:
    """One model per quantile, all sharing the same binned Dataset.

    A decision that is asymmetric in over- and under-supply needs the
    distribution, not the mean - and the binning work is done once for all
    three fits rather than three times.
    """
    boosters: dict[float, lgb.Booster] = {}

    for q in quantiles:
        boosters[q] = lgb.train(
            {
                "objective": "quantile",
                "alpha": q,
                "learning_rate": 0.03,
                "num_leaves": 63,
                "feature_fraction": 0.8,
                "bagging_fraction": 0.8,
                "bagging_freq": 1,
                "max_bin": 255,      # fits a bin index in one byte
                "verbosity": -1,
                "num_threads": 0,    # all cores
            },
            train,
            num_boost_round=rounds,
            valid_sets=[valid],
            callbacks=[lgb.early_stopping(100, verbose=False)],
        )

    return boosters


def predict_intervals(
    boosters: dict[float, lgb.Booster], X: Matrix
) -> dict[float, Vector]:
    """Score every quantile in one pass over the feature matrix, then enforce
    monotonicity - independently fitted quantiles can cross, which produces an
    incoherent distribution a downstream optimizer will happily exploit."""
    X = np.ascontiguousarray(X, dtype=np.float32)
    raw = {q: booster.predict(X, num_iteration=booster.best_iteration)
           for q, booster in boosters.items()}

    ordered = sorted(raw)
    stacked = np.vstack([raw[q] for q in ordered])
    np.maximum.accumulate(stacked, axis=0, out=stacked)   # in-place isotonic fix

    return {q: stacked[i] for i, q in enumerate(ordered)}