| """Test-time ensembling for Modular-TS (v3): context-length ensembles, temporal-hierarchy |
| (multi-resolution) reconciliation and in-context backtest weighting. Inference only. |
| |
| Members (spec string, comma separated): |
| cN - the base forecaster with max_context N (e.g. c8192, c2048, c512) |
| dK - downsampled member: context mean-pooled by K (aligned to the forecast origin, so it |
| sees up to K x the history at the same token budget), forecast ceil(H/K) coarse |
| steps, step-upsampled back. Used as a *level corrector* for the fine forecasts |
| (temporal hierarchies / THieF, Athanasopoulos et al. 2017): fine + a*(coarse - pool(fine)) |
| instead of being averaged in directly, so intra-period shape comes from full res. |
| Combination of the c-members: quantile averaging (Vincentisation), uniform or, with |
| backtest=True, per-series softmax weights from each member's scaled MAE on the last H |
| observed steps of the context (a zero-shot, in-context model selection - it only uses the |
| series' own history). |
| |
| fc = EnsembleForecaster(TimestralForecaster.from_pretrained(exp), "c8192,c2048,d4", alpha=0.5) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| @dataclass |
| class _F: |
| quantiles: torch.Tensor |
|
|
|
|
| def _pool(x: np.ndarray | None, k: int) -> np.ndarray | None: |
| """(B,V,T) NaN-aware mean pool by k aligned to the end.""" |
| if x is None: |
| return None |
| B, V, T = x.shape |
| n = T // k |
| if n == 0: |
| return None |
| y = x[..., T - n * k:].reshape(B, V, n, k) |
| with np.errstate(all="ignore"): |
| cnt = np.isfinite(y).sum(-1) |
| s = np.nansum(y, -1) |
| out = np.where(cnt > 0, s / np.maximum(cnt, 1), np.nan) |
| return out.astype(np.float32) |
|
|
|
|
| class EnsembleForecaster: |
| def __init__(self, fc, spec: str = "c8192", alpha: float = 0.5, backtest: bool = False, tau: float = 0.5, bt_windows: int = 1, prior: float = 0.0, |
| min_coarse_ctx: int = 256, min_extra: float = 1.5, min_ratio: float = 8.0): |
| self.fc = fc |
| self.members = [m.strip() for m in spec.split(",") if m.strip()] |
| cl = fc.model.config.context_length |
| self.ctxs = [min(int(m[1:]), cl) for m in self.members if m[0] == "c"] or [fc.max_context] |
| self.ks = [int(m[1:]) for m in self.members if m[0] == "d"] |
| self.alpha, self.backtest, self.tau = alpha, backtest, tau |
| self.bt_windows, self.prior = bt_windows, prior |
| self.min_coarse_ctx, self.min_extra, self.min_ratio = min_coarse_ctx, min_extra, min_ratio |
| self.base_ctx = fc.max_context |
| self.max_context = max(self.ctxs + [self.base_ctx * k for k in self.ks] + [self.base_ctx]) |
| self.quantile_levels = fc.quantile_levels |
| self.mi = list(fc.quantile_levels).index(0.5) if 0.5 in list(fc.quantile_levels) else len(fc.quantile_levels) // 2 |
|
|
| def _one(self, ctx, h, past, mc): |
| self.fc.max_context = mc |
| try: |
| return self.fc.forecast(ctx, h, past_covariates=past).quantiles.float().cpu() |
| finally: |
| self.fc.max_context = self.base_ctx |
|
|
| def _members(self, ctx, h, past): |
| T = ctx.shape[-1] |
| outs, used = [], set() |
| for i, c in enumerate(self.ctxs): |
| if i > 0 and c < self.min_ratio * h: |
| continue |
| eff = min(c, T) |
| if eff in used: |
| continue |
| used.add(eff) |
| outs.append(self._one(ctx[..., -c:], h, None if past is None else past[..., -c:], c)) |
| return outs |
|
|
| def forecast(self, context, prediction_length, past_covariates=None): |
| ctx = np.asarray(context, np.float32) |
| if ctx.ndim == 2: |
| ctx = ctx[:, None] |
| past = None if past_covariates is None else np.asarray(past_covariates, np.float32) |
| h = prediction_length |
| T = ctx.shape[-1] |
| outs = self._members(ctx, h, past) |
| if len(outs) == 1: |
| q = outs[0] |
| elif self.backtest and T - h >= 2 * 32: |
| E, ok = 0.0, True |
| nw = 0 |
| for wdw in range(1, self.bt_windows + 1): |
| if T - wdw * h < 2 * 32: |
| break |
| cut = T - wdw * h |
| bt = self._members(ctx[..., :cut], h, None if past is None else past[..., :cut]) |
| if len(bt) != len(outs): |
| ok = False |
| break |
| y = torch.from_numpy(ctx[..., cut:cut + h]) |
| hist = torch.from_numpy(ctx[..., :cut]) |
| dd = hist[..., 1:] - hist[..., :-1] |
| fin = torch.isfinite(dd) |
| scale = (torch.where(fin, dd.abs(), torch.zeros_like(dd)).sum(-1) / fin.sum(-1).clamp_min(1)).clamp_min(1e-8) |
| errs = [] |
| for b in bt: |
| e = (b[..., self.mi] - y).abs() |
| fe = torch.isfinite(e) |
| e = torch.where(fe, e, torch.zeros_like(e)).sum(-1) / fe.sum(-1).clamp_min(1) / scale |
| errs.append(e.mean(-1)) |
| E = E + torch.stack(errs, 0) |
| nw += 1 |
| if not ok or nw == 0: |
| q = torch.stack(outs).mean(0) |
| else: |
| E = E / nw |
| logit = -E / (self.tau * E.min(0, keepdim=True).values.clamp_min(1e-3)) |
| logit[0] = logit[0] + self.prior |
| w = torch.softmax(logit, 0) |
| q = (torch.stack(outs) * w[:, :, None, None, None]).sum(0) |
| else: |
| q = torch.stack(outs).mean(0) |
| |
| for k in self.ks: |
| hk = math.ceil(h / k) |
| pooled = _pool(ctx, k) |
| if pooled is None or pooled.shape[-1] < self.min_coarse_ctx: |
| continue |
| |
| if T < self.min_extra * self.base_ctx and hk < 2: |
| continue |
| pp = _pool(past, k) if past is not None else None |
| cq = self._one(pooled[..., -self.base_ctx:], hk, None if pp is None else pp[..., -self.base_ctx:], self.base_ctx) |
| cmed = cq[..., self.mi] |
| fmed = q[..., self.mi] |
| pad = hk * k - h |
| fm = torch.cat([fmed, fmed[..., -1:].expand(*fmed.shape[:-1], pad)], -1) if pad else fmed |
| fpool = fm.view(*fm.shape[:-1], hk, k).mean(-1) |
| delta = (cmed - fpool).repeat_interleave(k, -1)[..., :h] |
| q = q + self.alpha / len(self.ks) * delta[..., None] |
| return _F(q) |
|
|
