mlsauce

 19class ConformalBayesianRegressor(BaseEstimator, RegressorMixin):
 20    def __init__(
 21        self,
 22        obj=Ridge,
 23        level=95,
 24        hyperparameter_bounds=None,
 25        n_samples=20,
 26        calibration_fraction=0.2,
 27        scaling_method="standard",
 28        random_state=None,
 29        show_progress=False,
 30        verbose=True,
 31        n_jobs=-1,
 32    ):
 33        self.obj = obj
 34        self.level = level
 35        self.alpha_ = 1 - self.level / 100
 36        self.hyperparameter_bounds = hyperparameter_bounds
 37        self.n_samples = n_samples
 38        self.calibration_fraction = calibration_fraction
 39        self.scaling_method = scaling_method
 40        self.random_state = random_state
 41        self.verbose = verbose
 42        self.show_progress = show_progress
 43        self.n_jobs = n_jobs
 44
 45        self.is_fitted_ = False
 46
 47    def _sample_hyperparameters(self):
 48        # Simple uniform sampling or use fixed bounds
 49        configs = []
 50        for _ in range(self.n_samples):
 51            if self.hyperparameter_bounds:
 52                cfg = {}
 53                for k, v in self.hyperparameter_bounds.items():
 54                    if isinstance(v, list) and len(v) == 2:
 55                        # Always sample as float first
 56                        sampled_value = np.random.uniform(v[0], v[1])
 57
 58                        # If both bounds are integers, assume integer parameter
 59                        if isinstance(v[0], (int, np.integer)) and isinstance(
 60                            v[1], (int, np.integer)
 61                        ):
 62                            cfg[k] = int(sampled_value)
 63                        else:
 64                            cfg[k] = sampled_value
 65                    else:
 66                        # Assume fixed value
 67                        cfg[k] = v
 68            else:
 69                cfg = {}
 70            configs.append(cfg)
 71        return configs
 72
 73    def _train_models_parallel(self, X, y, configs):
 74        def train_model(cfg):
 75            self.obj.set_params(**cfg)
 76            self.obj.fit(X, y)
 77            return self.obj
 78
 79        if self.show_progress == False:
 80            models = joblib.Parallel(n_jobs=self.n_jobs)(
 81                joblib.delayed(train_model)(cfg)
 82                for cfg in tqdm(configs, disable=not self.verbose)
 83            )
 84        else:
 85            models = joblib.Parallel(n_jobs=self.n_jobs)(
 86                joblib.delayed(train_model)(cfg) for cfg in configs
 87            )
 88        return models
 89
 90    def _predict_models(self, models, X):
 91        if self.show_progress == False:
 92            preds = [m.predict(X) for m in models]
 93        else:
 94            preds = [
 95                m.predict(X) for m in tqdm(models, disable=not self.verbose)
 96            ]
 97        return np.column_stack(preds)
 98
 99    def fit(self, X, y):
100        if hasattr(X, "values"):  # keep DataFrame if possible
101            X = X.copy()
102        else:
103            X = np.asarray(X)
104        X = np.asarray(X)
105        y = np.asarray(y)
106        rng = check_random_state(self.random_state)
107        X, y = shuffle(X, y, random_state=rng)
108        # 1. Cluster with GMM (full covariance)
109        n_clusters = min(10, X.shape[0] // 30)
110        gmm = GaussianMixture(
111            n_components=n_clusters,
112            covariance_type="full",
113            random_state=self.random_state,
114        )
115        clusters = gmm.fit_predict(X)
116        # 2. Stratified train/calibration split
117        X_train, X_calib, y_train, y_calib = train_test_split(
118            X,
119            y,
120            test_size=self.calibration_fraction,
121            random_state=self.random_state,
122            stratify=clusters,
123        )
124        # 3. Scale features
125        if self.scaling_method == "standard":
126            scaler = StandardScaler().fit(X_train)
127        else:
128            raise ValueError("Scaling method must be 'standard'")
129        self.scaler_ = scaler
130        X_train_s = scaler.transform(X_train)
131        X_calib_s = scaler.transform(X_calib)
132        # 4. Train ensemble
133        configs = self._sample_hyperparameters()
134        self.models_ = self._train_models_parallel(X_train_s, y_train, configs)
135        # 5. Calibration residuals
136        preds_calib = self._predict_models(self.models_, X_calib_s)
137        self.calibration_residuals_ = y_calib - np.median(preds_calib, axis=1)
138        self.is_fitted_ = True
139        return self
140
141    def predict(self, X, return_pi=False):
142        """Obtain predictions and prediction intervals
143
144        Args:
145
146            X: array-like, shape = [n_samples, n_features];
147                Testing set vectors, where n_samples is the number
148                of samples and n_features is the number of features.
149
150            return_pi: boolean
151                Whether the prediction interval is returned or not.
152                Default is False, for compatibility with other _estimators_.
153                If True, a tuple containing the predictions + lower and upper
154                bounds is returned.
155
156        """
157        if not self.is_fitted_:
158            raise RuntimeError("Fit the model first")
159        X_s = self.scaler_.transform(X)
160        preds = self._predict_models(self.models_, X_s)
161        self.mean_ = np.median(preds, axis=1)
162        if return_pi == False:
163            return self.mean_
164        DescribeResult = namedtuple(
165            "DescribeResult", ("mean", "lower", "upper")
166        )
167        q = np.quantile(self.calibration_residuals_, q=self.alpha_ / 200)
168        return DescribeResult(self.mean_, self.mean_ + q, self.mean_ - q)
169
170    def get_coverage(self, y_true, lower, upper):
171        return np.mean((y_true >= lower) & (y_true <= upper))

 21class ConformalBayesianClassifier(BaseEstimator, ClassifierMixin):
 22
 23    # construct the object -----
 24    _estimator_type = "classifier"
 25
 26    def __init__(
 27        self,
 28        obj=Ridge(),
 29        hyperparameter_bounds=None,
 30        n_samples=20,
 31        calibration_fraction=0.2,
 32        scaling_method="standard",
 33        random_state=None,
 34        verbose=True,
 35        n_jobs=-1,
 36        calibrated=False,
 37        calibration_method="sigmoid",
 38        calibration_cv=3,
 39        use_classifier=True,
 40    ):
 41        """
 42        Conformal Bayesian Classifier with optional probability calibration.
 43
 44        Parameters
 45        ----------
 46        obj : estimator object
 47            Base estimator (default: Ridge() for regression-based, or classifier if use_classifier=True)
 48        hyperparameter_bounds : dict, optional
 49            Bounds for hyperparameter sampling
 50        n_samples : int
 51            Number of ensemble models to train
 52        calibration_fraction : float
 53            Fraction of data to use for conformal calibration
 54        scaling_method : str
 55            Feature scaling method (default: "standard")
 56        random_state : int, optional
 57            Random seed for reproducibility
 58        verbose : bool
 59            Whether to show progress bars
 60        n_jobs : int
 61            Number of parallel jobs (-1 for all cores)
 62        calibrated : bool
 63            Whether to apply probability calibration (default: False)
 64        calibration_method : str
 65            Method for calibration: "sigmoid" or "isotonic" (default: "sigmoid")
 66        calibration_cv : int
 67            Number of CV folds for calibration (default: 3)
 68        use_classifier : bool
 69            If True, use sklearn-like classifier with probability averaging.
 70            If False, use regression-based approach with SimpleMultitaskClassifier (default: False)
 71        """
 72        self.obj_base = obj
 73        self.hyperparameter_bounds = hyperparameter_bounds
 74        self.n_samples = n_samples
 75        self.calibration_fraction = calibration_fraction
 76        self.scaling_method = scaling_method
 77        self.random_state = random_state
 78        self.verbose = verbose
 79        self.n_jobs = n_jobs
 80        self.calibrated = calibrated
 81        self.calibration_method = calibration_method
 82        self.calibration_cv = calibration_cv
 83        self.use_classifier = use_classifier
 84
 85        self.is_fitted_ = False
 86
 87    def _sample_hyperparameters(self):
 88        """Sample hyperparameters for ensemble members"""
 89        configs = []
 90
 91        # Get parameter constraints from the model if available
 92        param_constraints = {}
 93        if hasattr(self.obj_base, "_parameter_constraints"):
 94            param_constraints = self.obj_base._parameter_constraints
 95
 96        for _ in range(self.n_samples):
 97            if self.hyperparameter_bounds:
 98                cfg = {}
 99                for k, v in self.hyperparameter_bounds.items():
100                    if isinstance(v, list) and len(v) == 2:
101                        # If both bounds are integers, assume integer parameter
102                        if isinstance(v[0], (int, np.integer)) and isinstance(
103                            v[1], (int, np.integer)
104                        ):
105                            cfg[k] = np.random.randint(v[0], v[1] + 1)
106                        else:
107                            cfg[k] = np.random.uniform(v[0], v[1])
108                    else:
109                        # Assume fixed value
110                        cfg[k] = v
111            else:
112                cfg = {}
113            configs.append(cfg)
114        return configs
115
116    def _train_classifiers_parallel(self, X, y, configs):
117        """Train ensemble of classifiers in parallel"""
118
119        def train_classifier(cfg):
120            clf = clone(self.obj_base)
121            clf.set_params(**cfg)
122            clf.fit(X, y)
123            return clf
124
125        if self.verbose:
126            models = joblib.Parallel(n_jobs=self.n_jobs)(
127                joblib.delayed(train_classifier)(cfg)
128                for cfg in tqdm(configs, desc="Training classifiers")
129            )
130        else:
131            models = joblib.Parallel(n_jobs=self.n_jobs)(
132                joblib.delayed(train_classifier)(cfg) for cfg in configs
133            )
134        return models
135
136    def fit(self, X, y):
137        """
138        Fit the conformal Bayesian classifier.
139
140        Parameters
141        ----------
142        X : array-like of shape (n_samples, n_features)
143            Training data
144        y : array-like of shape (n_samples,)
145            Target values
146
147        Returns
148        -------
149        self : object
150            Fitted classifier
151        """
152        X = np.asarray(X)
153        y = np.asarray(y)
154
155        # Store classes
156        self.classes_ = np.unique(y)
157        self.n_classes_ = len(self.classes_)
158
159        if self.use_classifier:
160            # Simpler approach: use sklearn-like classifier with probability averaging
161            if not is_classifier(self.obj_base):
162                raise ValueError(
163                    "use_classifier=True requires obj to be a classifier "
164                    "(e.g., LogisticRegression, RandomForestClassifier)"
165                )
166
167            rng = check_random_state(self.random_state)
168            X, y = shuffle(X, y, random_state=rng)
169
170            # Stratified train/calibration split
171            X_train, X_calib, y_train, y_calib = train_test_split(
172                X,
173                y,
174                test_size=self.calibration_fraction,
175                random_state=self.random_state,
176                stratify=y,
177            )
178
179            # Scale features
180            if self.scaling_method == "standard":
181                self.scaler_ = StandardScaler().fit(X_train)
182            else:
183                raise ValueError("Scaling method must be 'standard'")
184
185            X_train_s = self.scaler_.transform(X_train)
186            X_calib_s = self.scaler_.transform(X_calib)
187
188            # Train ensemble of classifiers
189            configs = self._sample_hyperparameters()
190            self.classifiers_ = self._train_classifiers_parallel(
191                X_train_s, y_train, configs
192            )
193
194            # Store calibration data for potential conformal prediction
195            self.X_calib_ = X_calib_s
196            self.y_calib_ = y_calib
197
198            if self.calibrated:
199                # Calibrate each classifier
200                if self.verbose:
201                    print("Calibrating classifiers...")
202
203                self.calibrated_classifiers_ = []
204                for clf in tqdm(
205                    self.classifiers_,
206                    disable=not self.verbose,
207                    desc="Calibrating",
208                ):
209                    cal_clf = CalibratedClassifierCV(
210                        clf,
211                        method=self.calibration_method,
212                        cv="prefit",
213                        ensemble=False,
214                    )
215                    cal_clf.fit(X_calib_s, y_calib)
216                    self.calibrated_classifiers_.append(cal_clf)
217
218        else:
219            # Original approach: use regression-based with SimpleMultitaskClassifier
220            base_regressor = ConformalBayesianRegressor(
221                self.obj_base,
222                hyperparameter_bounds=self.hyperparameter_bounds,
223                n_samples=self.n_samples,
224                calibration_fraction=self.calibration_fraction,
225                scaling_method=self.scaling_method,
226                random_state=self.random_state,
227                verbose=self.verbose,
228                n_jobs=self.n_jobs,
229            )
230
231            # Wrap in multitask classifier
232            self.obj = ns.SimpleMultitaskClassifier(base_regressor)
233
234            if self.calibrated:
235                # Fit the base classifier first
236                self.obj.fit(X, y)
237
238                # Then wrap with calibration
239                self.calibrated_obj_ = CalibratedClassifierCV(
240                    self.obj,
241                    method=self.calibration_method,
242                    cv="prefit",
243                    ensemble=False,
244                )
245
246                # Fit calibration on the same data
247                self.calibrated_obj_.fit(X, y)
248            else:
249                self.obj.fit(X, y)
250
251        self.is_fitted_ = True
252        return self
253
254    def predict_proba(self, X):
255        """
256        Predict class probabilities.
257
258        Parameters
259        ----------
260        X : array-like of shape (n_samples, n_features)
261            Test data
262
263        Returns
264        -------
265        proba : array of shape (n_samples, n_classes)
266            Predicted probabilities for each class
267        """
268        if not self.is_fitted_:
269            raise RuntimeError("Fit the model first")
270
271        X = np.asarray(X)
272
273        if self.use_classifier:
274            # Average probabilities from ensemble
275            X_s = self.scaler_.transform(X)
276
277            if self.calibrated:
278                # Use calibrated classifiers
279                probas = np.array(
280                    [
281                        clf.predict_proba(X_s)
282                        for clf in self.calibrated_classifiers_
283                    ]
284                )
285            else:
286                # Use uncalibrated classifiers
287                probas = np.array(
288                    [clf.predict_proba(X_s) for clf in self.classifiers_]
289                )
290
291            # Average probabilities across ensemble
292            mean_proba = np.mean(probas, axis=0)
293
294            return mean_proba
295        else:
296            # Use regression-based approach
297            if self.calibrated:
298                return self.calibrated_obj_.predict_proba(X)
299            else:
300                return self.obj.predict_proba(X)
301
302    def predict(self, X):
303        """
304        Predict class labels.
305
306        Parameters
307        ----------
308        X : array-like of shape (n_samples, n_features)
309            Test data
310
311        Returns
312        -------
313        y_pred : array of shape (n_samples,)
314            Predicted class labels
315        """
316        proba = self.predict_proba(X)
317        return self.classes_[np.argmax(proba, axis=1)]
318
319    @property
320    def _estimator_type(self):
321        return "classifier"

 16class ContextAwareThetaForecaster:
 17    """
 18    Unified Theta Method with multiple estimation modes.
 19
 20    Variants:
 21    - 'classic': Standard Theta (SES + drift)
 22    - 'cox': Context-aware with Cox partial likelihood
 23    - 'ridge': Context-aware with Ridge regression
 24    - 'ml': ML-enhanced with sklearn estimator
 25    - 'rslope': R-style slopes via numerical differentiation (context-free)
 26
 27    Parameters
 28    ----------
 29    mode : {'classic', 'cox', 'ridge', 'ml', 'rslope'}
 30        Estimation mode
 31    theta : float, default=0.5
 32        Drift intensity (0=no drift, 0.5=classical, 1=full)
 33    estimator : sklearn estimator, optional
 34        For 'ml' and 'rslope' modes
 35    tau : float, default=12
 36        Temporal attention decay
 37    sigma_val : float, optional
 38        Value-based kernel bandwidth
 39    kernel : {'temporal', 'value', 'hybrid'}
 40        Attention kernel type
 41    seasonal_period : int, optional
 42        Seasonal period (auto-detected if None)
 43    risk_set_size : int, default=15
 44        Risk set size for Cox PL
 45    stability_factor : float, default=0.8
 46        Gamma clipping safety factor (0,1]
 47    random_state : int, optional
 48        Random seed for ML mode
 49    """
 50
 51    def __init__(
 52        self,
 53        mode: Literal["classic", "cox", "ridge", "ml", "rslope"] = "cox",
 54        theta: float = 0.5,
 55        estimator: Optional[Any] = None,
 56        tau: float = 12.0,
 57        sigma_val: Optional[float] = None,
 58        kernel: Literal["temporal", "value", "hybrid"] = "temporal",
 59        seasonal_period: Optional[int] = None,
 60        risk_set_size: int = 15,
 61        stability_factor: float = 0.8,
 62        random_state: Optional[int] = None,
 63    ):
 64        self.mode = mode
 65        self.theta = theta
 66        self.estimator = estimator
 67        self.tau = tau
 68        self.sigma_val = sigma_val
 69        self.kernel = kernel
 70        self.seasonal_period = seasonal_period
 71        self.risk_set_size = risk_set_size
 72        self.stability_factor = stability_factor
 73        self.random_state = random_state
 74
 75        # Fitted params
 76        self.alpha_ = None
 77        self.l_n_ = None
 78        self.b0_ = None
 79        self.gamma_ = None
 80        self.mu_z_ = None
 81        self.sigma_z_ = None
 82        self.sigma2_ = None
 83        self.seasonal_indices_ = None
 84        self._y_train = None
 85        self._fitted = False
 86
 87    # ============ SEASONAL DECOMPOSITION ============
 88    def _decompose(self, y: np.ndarray, period: int):
 89        """Multiplicative seasonal decomposition"""
 90        n = len(y)
 91        if n < 2 * period:
 92            return y, np.ones(n), y
 93
 94        # Centered MA trend
 95        trend = np.full(n, np.nan)
 96        half = period // 2
 97        for i in range(half, n - half):
 98            window = y[i - half : i + half + (period % 2)]
 99            trend[i] = np.mean(window)
100
101        # Fill edges
102        valid = np.where(~np.isnan(trend))[0]
103        if len(valid) > 0:
104            trend[: valid[0]] = trend[valid[0]]
105            trend[valid[-1] + 1 :] = trend[valid[-1]]
106
107        # Seasonal indices
108        detrended = y / (trend + 1e-10)
109        seasonal = np.zeros(period)
110        for i in range(period):
111            seasonal[i] = np.nanmean(detrended[i::period])
112        seasonal /= seasonal.mean() + 1e-10
113
114        seasonal_full = np.tile(seasonal, n // period + 1)[:n]
115        adjusted = y / (seasonal_full + 1e-10)
116
117        return adjusted, seasonal_full, trend
118
119    # ============ SES ============
120    def _ses_level(self, y: np.ndarray, alpha: float) -> np.ndarray:
121        """Compute SES level array"""
122        level = np.zeros(len(y))
123        level[0] = y[0]
124        for t in range(1, len(y)):
125            level[t] = alpha * y[t] + (1 - alpha) * level[t - 1]
126        return level
127
128    def _ses_nll(self, alpha: float, y: np.ndarray) -> float:
129        """SES negative log-likelihood"""
130        if alpha <= 0 or alpha >= 1:
131            return 1e10
132        level = self._ses_level(y, alpha)
133        resid = y[1:] - level[:-1]
134        sigma2 = np.var(resid) + 1e-10
135        return 0.5 * len(resid) * (np.log(2 * np.pi * sigma2) + 1)
136
137    def _fit_ses(self, y: np.ndarray):
138        """Estimate alpha via MLE"""
139        res = minimize(
140            lambda a: self._ses_nll(a[0], y),
141            [0.3],
142            bounds=[(0.01, 0.99)],
143            method="L-BFGS-B",
144        )
145        alpha = res.x[0]
146        level_array = self._ses_level(y, alpha)
147        return alpha, level_array[-1], level_array
148
149    # ============ DRIFT ============
150    def _estimate_drift(self, y: np.ndarray) -> float:
151        """Baseline drift: b0 = beta_OLS / 2"""
152        t = np.arange(len(y))
153        beta = np.sum((t - t.mean()) * (y - y.mean())) / (
154            np.sum((t - t.mean()) ** 2) + 1e-10
155        )
156        return beta / 2.0
157
158    # ============ ATTENTION CONTEXT ============
159    def _attention_kernel(self, Xj: float, Xt: float, j: int, t: int) -> float:
160        """Compute kernel weight"""
161        if self.kernel == "temporal":
162            return np.exp(-(t - j) / (self.tau + 1e-12))
163        elif self.kernel == "value":
164            sigma = self.sigma_val if self.sigma_val else 1.0
165            return np.exp(-((Xj - Xt) ** 2) / (2 * sigma**2 + 1e-12))
166        else:  # hybrid
167            sigma = self.sigma_val if self.sigma_val else 1.0
168            return np.exp(
169                -(t - j) / (self.tau + 1e-12)
170                - ((Xj - Xt) ** 2) / (2 * sigma**2 + 1e-12)
171            )
172
173    def _compute_context(self, y: np.ndarray) -> np.ndarray:
174        """Compute attention-weighted context z_t"""
175        n = len(y)
176        z = np.zeros(n)
177        for t in range(n):
178            weights = np.array(
179                [self._attention_kernel(y[j], y[t], j, t) for j in range(t + 1)]
180            )
181            weights /= weights.sum() + 1e-12
182            z[t] = np.dot(weights, y[: t + 1])
183        return z
184
185    # ============ GAMMA ESTIMATION ============
186    def _partial_nll(self, gamma: float, z_star: np.ndarray) -> float:
187        """Cox partial negative log-likelihood (stable)"""
188        n = len(z_star)
189        k = min(self.risk_set_size, n // 2)
190        nll = 0.0
191        for t in range(k, n):
192            z_risk = z_star[max(0, t - k) : t + 1]
193            nll -= gamma * z_star[t] - logsumexp(gamma * z_risk)
194        return nll
195
196    def _estimate_gamma_cox(self, z_star: np.ndarray) -> float:
197        """Estimate gamma via Cox PL"""
198        res = minimize(
199            lambda g: self._partial_nll(g[0], z_star), [0.0], method="BFGS"
200        )
201        return res.x[0]
202
203    def _estimate_gamma_ridge(
204        self,
205        y: np.ndarray,
206        z_star: np.ndarray,
207        level_array: np.ndarray,
208        alpha: float,
209        b0: float,
210    ) -> float:
211        """Estimate gamma via Ridge regression"""
212        try:
213            from sklearn.linear_model import Ridge
214        except ImportError:
215            warnings.warn("sklearn unavailable, using Cox method")
216            return self._estimate_gamma_cox(z_star)
217
218        n = len(y)
219        residuals = y - level_array
220
221        # Build design matrix
222        X_design = []
223        y_resid = []
224        for t in range(20, n):
225            D_h = self._D_n(1, alpha, t)
226            x_t = 0.5 * b0 * z_star[t - 1] * D_h
227            X_design.append(x_t)
228            y_resid.append(residuals[t])
229
230        X_design = np.array(X_design).reshape(-1, 1)
231        y_resid = np.array(y_resid)
232
233        ridge = Ridge(alpha=10.0, fit_intercept=False)
234        ridge.fit(X_design, y_resid)
235        return ridge.coef_[0]
236
237    def _estimate_gamma_ml(self, y: np.ndarray, h: int) -> float:
238        """Estimate gamma via ML numerical differentiation"""
239        if self.estimator is None:
240            from sklearn.linear_model import LinearRegression
241
242            self.estimator = LinearRegression()
243
244        n = len(y)
245        if self.random_state is not None:
246            np.random.seed(self.random_state)
247
248        # Features
249        time_idx = np.arange(n + h)
250        time_norm = time_idx / n
251        n_random = 3
252        random_cov = np.random.randn(n + h, n_random)
253        X_all = np.column_stack([time_norm, random_cov])
254
255        # Scale
256        from sklearn.preprocessing import StandardScaler
257
258        scaler = StandardScaler()
259        X_train = scaler.fit_transform(X_all[:n])
260        X_all = scaler.transform(X_all)
261
262        # Fit
263        self.estimator.fit(X_train, y)
264
265        # Numerical differentiation
266        eps = 1e-4 ** (1 / 3)
267        h_eps = np.maximum(eps * np.abs(time_norm), eps / n)
268
269        t_plus = np.clip(time_norm + h_eps, 0, 2.0)
270        t_minus = np.clip(time_norm - h_eps, 0, 2.0)
271
272        X_plus = scaler.transform(np.column_stack([t_plus, random_cov]))
273        X_minus = scaler.transform(np.column_stack([t_minus, random_cov]))
274
275        fx_plus = self.estimator.predict(X_plus)
276        fx_minus = self.estimator.predict(X_minus)
277
278        slopes = (fx_plus - fx_minus) / (2 * h_eps) / n
279
280        # Approximate gamma from slopes
281        return self.theta * slopes.mean() / (0.5 * self.b0_ + 1e-12)
282
283    def _estimate_slopes_rslope(self, y: np.ndarray, h: int) -> np.ndarray:
284        """
285        R-style slope estimation via numerical differentiation.
286        Returns slopes for forecast horizon (not gamma).
287        Similar to estimate_theta_slope() in R code.
288        """
289        if self.estimator is None:
290            from sklearn.linear_model import LinearRegression
291
292            self.estimator = LinearRegression()
293
294        n = len(y)
295        if self.random_state is not None:
296            np.random.seed(self.random_state)
297
298        # Create features: time + random noise
299        time_idx = np.arange(n + h)
300        time_norm = time_idx / n
301        n_random = 3
302        random_cov = np.random.randn(n + h, n_random)
303        X_all = np.column_stack([time_norm, random_cov])
304
305        # Scale features
306        from sklearn.preprocessing import StandardScaler
307
308        scaler = StandardScaler()
309        X_train = scaler.fit_transform(X_all[:n])
310        X_all_scaled = scaler.transform(X_all)
311
312        # Fit model
313        self.estimator.fit(X_train, y)
314
315        # Numerical differentiation for ALL points (historical + forecast)
316        eps = 1e-4 ** (1 / 3)
317        h_eps = np.maximum(eps * np.abs(time_norm), eps / n)
318
319        t_plus = np.clip(time_norm + h_eps, 0, 2.0)
320        t_minus = np.clip(time_norm - h_eps, 0, 2.0)
321
322        X_plus = scaler.transform(np.column_stack([t_plus, random_cov]))
323        X_minus = scaler.transform(np.column_stack([t_minus, random_cov]))
324
325        fx_plus = self.estimator.predict(X_plus)
326        fx_minus = self.estimator.predict(X_minus)
327
328        # Slopes at each time point
329        slopes = (fx_plus - fx_minus) / (2 * h_eps) / n
330
331        # Return ONLY the forecast horizon slopes (last h values)
332        return slopes[-h:] * self.theta
333
334    # ============ DRIFT MULTIPLIER ============
335    def _D_n(self, h: int, alpha: float, n: int) -> float:
336        """Drift multiplier D_n(h)"""
337        return (h - 1) + (1 - (1 - alpha) ** n) / (alpha + 1e-12)
338
339    # ============ FIT ============
340    def fit(self, y: np.ndarray):
341        """Fit the model"""
342        y = np.asarray(y, dtype=float).ravel()
343        self._y_train = y.copy()
344        n = len(y)
345
346        # Detect seasonality
347        period = self.seasonal_period
348        if period is None:
349            period = 12 if n >= 24 else 1
350
351        # Decompose
352        if period > 1 and n >= 2 * period:
353            y_adj, seasonal_full, _ = self._decompose(y, period)
354            self.seasonal_indices_ = seasonal_full[-period:]
355        else:
356            y_adj = y.copy()
357            self.seasonal_indices_ = None
358
359        # SES
360        self.alpha_, self.l_n_, level_array = self._fit_ses(y_adj)
361
362        # Drift
363        self.b0_ = self._estimate_drift(y_adj)
364
365        # Context & Gamma
366        if self.mode == "classic":
367            self.gamma_ = 0.0
368            self.mu_z_ = 0.0
369            self.sigma_z_ = 1.0
370        elif self.mode == "rslope":
371            # R-style: no gamma, slopes computed during prediction
372            self.gamma_ = 0.0
373            self.mu_z_ = 0.0
374            self.sigma_z_ = 1.0
375        else:
376            # Compute context
377            z_raw = self._compute_context(y_adj)
378            self.mu_z_ = z_raw.mean()
379            self.sigma_z_ = z_raw.std() + 1e-12
380            z_star = (z_raw - self.mu_z_) / self.sigma_z_
381
382            # Estimate gamma
383            if self.mode == "cox":
384                gamma_raw = self._estimate_gamma_cox(z_star)
385            elif self.mode == "ridge":
386                gamma_raw = self._estimate_gamma_ridge(
387                    y_adj, z_star, level_array, self.alpha_, self.b0_
388                )
389            else:  # ml
390                gamma_raw = self._estimate_gamma_ml(y_adj, 12)
391
392            # Stability constraint
393            D_max = self._D_n(36, self.alpha_, n)
394            stability_bound = 2.0 / (abs(self.b0_) * D_max + 1e-12)
395            self.gamma_ = np.clip(
396                gamma_raw,
397                -self.stability_factor * stability_bound,
398                self.stability_factor * stability_bound,
399            )
400
401        # Innovation variance
402        residuals = y_adj[1:] - level_array[:-1]
403        self.sigma2_ = np.var(residuals, ddof=1)
404
405        self._fitted = True
406        return self
407
408    # ============ PREDICT ============
409    def predict(
410        self, h: int, return_pi: bool = True, alpha_ci: float = 0.05
411    ) -> Dict[str, np.ndarray]:
412        """Generate forecasts"""
413        if not self._fitted:
414            raise RuntimeError("Call fit() first")
415
416        n = len(self._y_train)
417
418        # For rslope mode, compute slopes once
419        if self.mode == "rslope":
420            y_adj = (
421                self._y_train
422                if self.seasonal_indices_ is None
423                else self._y_train
424                / np.tile(
425                    self.seasonal_indices_, n // len(self.seasonal_indices_) + 1
426                )[:n]
427            )
428            rslope_slopes = self._estimate_slopes_rslope(y_adj, h)
429
430        # Deseasonalized forecast (recursive for non-rslope, direct for rslope)
431        if self.mode == "rslope":
432            # R-style: direct application of slopes
433            fc = np.zeros(h)
434            for step in range(h):
435                D_h = self._D_n(step + 1, self.alpha_, n)
436                fc[step] = self.l_n_ + rslope_slopes[step] * D_h
437        else:
438            # Original recursive logic
439            fc = []
440            history = list(
441                self._y_train
442                if self.seasonal_indices_ is None
443                else self._y_train
444                / np.tile(
445                    self.seasonal_indices_, n // len(self.seasonal_indices_) + 1
446                )[:n]
447            )
448
449            for step in range(1, h + 1):
450                # Recompute context
451                if self.mode != "classic":
452                    t_now = len(history) - 1
453                    weights = np.array(
454                        [
455                            self._attention_kernel(
456                                history[j], history[t_now], j, t_now
457                            )
458                            for j in range(t_now + 1)
459                        ]
460                    )
461                    weights /= weights.sum() + 1e-12
462                    z_h = np.dot(weights, history)
463                    z_h_star = (z_h - self.mu_z_) / self.sigma_z_
464                else:
465                    z_h_star = 0.0
466
467                # Forecast
468                D_h = self._D_n(step, self.alpha_, n)
469                context_factor = 1.0 + self.gamma_ * z_h_star
470                fc_val = (
471                    self.l_n_
472                    + 0.5 * self.b0_ * self.theta * context_factor * D_h
473                )
474
475                fc.append(fc_val)
476                history.append(fc_val)
477
478            fc = np.array(fc)
479
480        # Reseasonalize
481        if self.seasonal_indices_ is not None:
482            seasonal_fc = np.tile(
483                self.seasonal_indices_, (h // len(self.seasonal_indices_)) + 1
484            )[:h]
485            fc *= seasonal_fc
486
487        result = {"mean": fc}
488
489        # Prediction intervals
490        if return_pi:
491            z_score = norm.ppf(1 - alpha_ci / 2)
492            lower = []
493            upper = []
494
495            for step in range(1, h + 1):
496                D_h = self._D_n(step, self.alpha_, n)
497                var_ses = self.sigma2_ * ((step - 1) * self.alpha_**2 + 1)
498                var_ctx = (
499                    (0.5 * self.gamma_ * self.b0_ * self.sigma_z_ * D_h) ** 2
500                ) / n
501                se = np.sqrt(var_ses + var_ctx)
502
503                lower.append(fc[step - 1] - z_score * se)
504                upper.append(fc[step - 1] + z_score * se)
505
506            result["lower"] = np.array(lower)
507            result["upper"] = np.array(upper)
508
509        return result
510
511    # ============ UTILITIES ============
512    def get_params(self) -> Dict[str, Any]:
513        """Get fitted parameters"""
514        return {
515            "mode": self.mode,
516            "alpha": self.alpha_,
517            "b0": self.b0_,
518            "gamma": self.gamma_,
519            "l_n": self.l_n_,
520            "theta": self.theta,
521            "sigma2": self.sigma2_,
522            "seasonal": self.seasonal_indices_ is not None,
523        }
524
525    def plot(self, forecast: Dict[str, np.ndarray], title: str = None):
526        """Rich visualization of forecasts"""
527        import matplotlib.pyplot as plt
528
529        n = len(self._y_train)
530        h = len(forecast["mean"])
531
532        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
533        fig.suptitle(
534            title or f"Unified Theta: {self.mode.upper()} Mode",
535            fontsize=14,
536            fontweight="bold",
537        )
538
539        # Plot 1: Forecasts with PI
540        ax1 = axes[0, 0]
541        train_idx = np.arange(n)
542        fc_idx = np.arange(n, n + h)
543
544        ax1.plot(
545            train_idx,
546            self._y_train,
547            "o-",
548            color="black",
549            label="Train",
550            linewidth=1.5,
551            markersize=3,
552            alpha=0.7,
553        )
554        ax1.plot(
555            fc_idx,
556            forecast["mean"],
557            "s-",
558            color="steelblue",
559            label="Forecast",
560            linewidth=2.5,
561            markersize=5,
562        )
563        if "lower" in forecast:
564            ax1.fill_between(
565                fc_idx,
566                forecast["lower"],
567                forecast["upper"],
568                color="lightblue",
569                alpha=0.3,
570                label="95% PI",
571            )
572        ax1.axvline(
573            n - 0.5, color="red", linestyle="--", alpha=0.5, linewidth=2
574        )
575        ax1.set_title("Forecasts with Prediction Intervals", fontweight="bold")
576        ax1.set_xlabel("Time")
577        ax1.set_ylabel("Value")
578        ax1.legend(loc="upper left")
579        ax1.grid(alpha=0.3)
580
581        # Plot 2: Context variable (if not classic/rslope)
582        ax2 = axes[0, 1]
583        if self.mode not in ["classic", "rslope"]:
584            y_adj = (
585                self._y_train
586                if self.seasonal_indices_ is None
587                else self._y_train
588                / np.tile(
589                    self.seasonal_indices_, n // len(self.seasonal_indices_) + 1
590                )[:n]
591            )
592            z = self._compute_context(y_adj)
593            z_star = (z - self.mu_z_) / self.sigma_z_
594
595            ax2.plot(
596                train_idx,
597                z_star,
598                color="purple",
599                linewidth=2,
600                label="z* (standardized)",
601            )
602            ax2.axhline(0, color="black", linestyle="--", alpha=0.5)
603            ax2.fill_between(
604                train_idx,
605                0,
606                z_star,
607                where=(z_star > 0),
608                color="green",
609                alpha=0.2,
610                label="Above trend",
611            )
612            ax2.fill_between(
613                train_idx,
614                0,
615                z_star,
616                where=(z_star < 0),
617                color="red",
618                alpha=0.2,
619                label="Below trend",
620            )
621            ax2.set_title(
622                f"Context Signal (γ={self.gamma_:.4f})", fontweight="bold"
623            )
624            ax2.set_xlabel("Time")
625            ax2.set_ylabel("z* (std dev)")
626            ax2.legend()
627        elif self.mode == "rslope":
628            ax2.text(
629                0.5,
630                0.5,
631                "R-Slope Mode\n(Direct ML Slopes)",
632                ha="center",
633                va="center",
634                transform=ax2.transAxes,
635                fontsize=14,
636                bbox=dict(boxstyle="round", facecolor="lightblue"),
637            )
638            ax2.axis("off")
639        else:
640            ax2.text(
641                0.5,
642                0.5,
643                "Classic Mode\n(No Context)",
644                ha="center",
645                va="center",
646                transform=ax2.transAxes,
647                fontsize=14,
648                bbox=dict(boxstyle="round", facecolor="wheat"),
649            )
650            ax2.axis("off")
651        ax2.grid(alpha=0.3)
652
653        # Plot 3: Residuals
654        ax3 = axes[1, 0]
655        y_adj = (
656            self._y_train
657            if self.seasonal_indices_ is None
658            else self._y_train
659            / np.tile(
660                self.seasonal_indices_, n // len(self.seasonal_indices_) + 1
661            )[:n]
662        )
663        level_array = self._ses_level(y_adj, self.alpha_)
664        residuals = y_adj[1:] - level_array[:-1]
665
666        ax3.scatter(train_idx[1:], residuals, alpha=0.6, s=20, color="coral")
667        ax3.axhline(0, color="black", linestyle="--", linewidth=1.5)
668        ax3.set_title(f"Residuals (σ²={self.sigma2_:.3f})", fontweight="bold")
669        ax3.set_xlabel("Time")
670        ax3.set_ylabel("Residual")
671        ax3.grid(alpha=0.3)
672
673        # Histogram
674        ax3_inset = ax3.inset_axes([0.65, 0.65, 0.3, 0.3])
675        ax3_inset.hist(
676            residuals, bins=15, color="coral", alpha=0.7, edgecolor="black"
677        )
678        ax3_inset.axvline(0, color="black", linestyle="--", linewidth=1)
679        ax3_inset.set_title("Distribution", fontsize=8)
680        ax3_inset.tick_params(labelsize=7)
681
682        # Plot 4: Summary table
683        ax4 = axes[1, 1]
684        ax4.axis("off")
685
686        params = self.get_params()
687        summary = f"""
688╔═══════════════════════════════════════╗
689║        MODEL PARAMETERS               ║
690╠═══════════════════════════════════════╣
691║                                       ║
692║  Mode:          {self.mode.upper():<20}  ║
693║  Theta:         {self.theta:<20.4f}  ║
694║  Alpha (α):     {params['alpha']:<20.4f}  ║
695║  Drift (b₀):    {params['b0']:<20.4f}  ║
696║  Gamma (γ):     {params['gamma']:<20.6f}  ║
697║  Level (ℓₙ):    {params['l_n']:<20.2f}  ║
698║  σ²:            {params['sigma2']:<20.4f}  ║
699║  Seasonal:      {str(params['seasonal']):<20}  ║
700║                                       ║
701╠═══════════════════════════════════════╣
702║        FORECAST SUMMARY               ║
703╠═══════════════════════════════════════╣
704║                                       ║
705║  Horizon:       {h:<20} steps     ║
706║  Final FC:      {forecast['mean'][-1]:<20.2f}  ║
707"""
708        if "lower" in forecast:
709            summary += f"║  95% PI:        [{forecast['lower'][-1]:>6.2f}, {forecast['upper'][-1]:>6.2f}]   ║\n"
710
711        summary += "║                                       ║\n"
712        summary += "╚═══════════════════════════════════════╝"
713
714        ax4.text(
715            0.05,
716            0.95,
717            summary,
718            fontsize=9,
719            family="monospace",
720            verticalalignment="top",
721            transform=ax4.transAxes,
722            bbox=dict(
723                boxstyle="round",
724                facecolor="lightyellow",
725                alpha=0.8,
726                edgecolor="black",
727                linewidth=1.5,
728            ),
729        )
730
731        plt.tight_layout()
732        return fig