File size: 5,562 Bytes
02d27c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """KHALO's reusable uncertainty engine — the shared layer every surrogate wraps.
Two heads (GP posterior for regression, deep-ensemble for classification) plus
the canonical calibration report from `metrics`. sklearn is imported lazily so
`import kronos_ml` stays clean in a bare env.
"""
from __future__ import annotations
import numpy as np
from . import metrics
SEED = 20260726 # the corpus-wide fixed seed
class GPHead:
"""Gaussian-process regressor (Matern-5/2 x Constant + White), optionally
delta-learning on a cheap baseline. Returns calibrated (mean, std)."""
def __init__(self, baseline=None, length_scale=0.5, alpha=1e-8, n_restarts=2):
self.baseline = baseline
self.length_scale = length_scale
self.alpha = alpha
self.n_restarts = n_restarts
self._gp = None
self._xmu = self._xsd = None
def _std_x(self, X):
X = np.atleast_2d(np.asarray(X, float))
if self._xmu is None:
self._xmu = X.mean(0); self._xsd = X.std(0) + 1e-12
return (X - self._xmu) / self._xsd
def fit(self, X, y):
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import (
Matern, ConstantKernel, WhiteKernel)
X = np.atleast_2d(np.asarray(X, float)); y = np.asarray(y, float)
resid = y - self._base(X) if self.baseline is not None else y
k = (ConstantKernel(1.0, (1e-3, 1e3))
* Matern(length_scale=self.length_scale, length_scale_bounds=(1e-2, 1e1), nu=2.5)
+ WhiteKernel())
self._gp = GaussianProcessRegressor(
kernel=k, alpha=self.alpha, normalize_y=True,
n_restarts_optimizer=self.n_restarts, random_state=SEED)
self._gp.fit(self._std_x(X), resid)
return self
def _base(self, X):
return np.asarray([self.baseline(xi) for xi in np.atleast_2d(X)], float)
def predict(self, X):
X = np.atleast_2d(np.asarray(X, float))
mean, std = self._gp.predict(self._std_x(X), return_std=True)
if self.baseline is not None:
mean = mean + self._base(X)
return mean, std
class EnsembleHead:
"""Soft-vote deep ensemble (HistGB + RandomForest + MLP) for the disruption /
classifier path. Returns mean predict_proba."""
def __init__(self):
self.models = None
self.scaler = None
def fit(self, X, y):
from sklearn.ensemble import (HistGradientBoostingClassifier,
RandomForestClassifier)
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
X = np.asarray(X, float); y = np.asarray(y, int)
self.scaler = StandardScaler().fit(X)
Xs = self.scaler.transform(X)
self.models = [
HistGradientBoostingClassifier(max_iter=120, max_depth=6,
learning_rate=0.12, random_state=SEED).fit(Xs, y),
RandomForestClassifier(n_estimators=120, max_depth=10,
random_state=SEED).fit(Xs, y),
MLPClassifier(hidden_layer_sizes=(48, 24), max_iter=180,
early_stopping=True, random_state=SEED).fit(Xs, y),
]
return self
def predict_proba(self, X):
Xs = self.scaler.transform(np.asarray(X, float))
return np.mean([m.predict_proba(Xs)[:, 1] for m in self.models], axis=0)
class CalibratedClassifier:
"""Deep ensemble + isotonic probability calibration (the net-new KODEX layer
on the disruption path). Fits the ensemble on a train split and the isotonic
map on a held-out calibration split."""
def __init__(self, cal_frac=0.3):
self.cal_frac = cal_frac
self.ens = None
self.iso = None
def fit(self, X, y):
from sklearn.isotonic import IsotonicRegression
X = np.asarray(X, float); y = np.asarray(y, int)
n = len(y); ntr = int((1 - self.cal_frac) * n)
self.ens = EnsembleHead().fit(X[:ntr], y[:ntr])
p = self.ens.predict_proba(X[ntr:])
self.iso = IsotonicRegression(out_of_bounds="clip").fit(p, y[ntr:])
self._cal_X, self._cal_y = X[ntr:], y[ntr:]
return self
def predict_proba(self, X):
return self.iso.predict(self.ens.predict_proba(np.asarray(X, float)))
class MahalanobisGate:
"""Per-sample out-of-domain gate: Mahalanobis distance (LedoitWolf-shrunk
precision) to the training distribution, thresholded to a target in-dist
abstention rate. in_domain(x) is False when x is anomalous vs training."""
def __init__(self, abstain_rate=0.05):
self.abstain_rate = abstain_rate
self.mu = self.prec = self.thresh = None
def fit(self, X):
from sklearn.covariance import LedoitWolf
X = np.asarray(X, float)
self.mu = X.mean(0)
self.prec = LedoitWolf().fit(X).precision_
d = self._d(X)
self.thresh = float(np.quantile(d, 1 - self.abstain_rate))
return self
def _d(self, X):
D = np.atleast_2d(np.asarray(X, float)) - self.mu
return np.sqrt(np.clip(np.einsum("ij,jk,ik->i", D, self.prec, D), 0, None))
def in_domain(self, x):
return bool(self._d(x)[0] <= self.thresh)
def calibrate(y_true, mean, std, levels=metrics.DEFAULT_LEVELS):
"""The first-class calibration object KHALO stamps on any surrogate output."""
return metrics.calibration_report(y_true, mean, std, levels)
|