| """Synthetic prior v2 (Modular-TS v3): new generator families + augmentation cascade. |
| |
| Motivated by the run2/run3 GIFT-Eval diagnosis (v3/INSIGHTS.md) and TempoPFN's |
| leave-one-out prior ablation (arXiv 2510.25502, Table 3: removing the regime-switching |
| OU/SDE prior costs the most, then CauKer, spikes, stochastic rhythms, step functions, |
| sawtooth; the augmentation cascade is worth ~5% CRPS): |
| |
| New univariate families (all vectorised torch, run on the GPU): |
| * ou_regime - Ornstein-Uhlenbeck with Markov regime switching of mean / vol / speed |
| * waveform - non-sinusoidal periodic waves: sawtooth, square (random duty), triangle, |
| spike trains, with optional slow period drift (phase-drifting cycles) |
| * template - "load-curve" seasonality: a random smooth+spiky daily profile tiled at a |
| bank period, optionally modulated by a weekly profile (24x7, 96x7, 288x7) |
| and multiplicative trend; near-deterministic, low-noise (m4_hourly-like) |
| * steps - piecewise-constant levels with Poisson change points (+ small noise) |
| * growth - integrated / cumulative processes: cumsum of a positive rate (covid-like), |
| logistic / Gompertz growth with noise, compound-interest trends |
| Augmentations (applied after generation, each with its own probability): |
| time reversal, integration (cumsum), differencing, quantisation to a random grid, |
| censoring (clip at a random quantile), random smoothing conv filter, sample-and-hold |
| resampling artefacts, amplitude modulation, TSMixup of two series. |
| |
| `SynthGen2.univariate` returns a mixture of the v1 prior (`SynthGen.univariate`) and the new |
| families; everything downstream (multivariatizers, covariates, masks) is inherited. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| from torch import Tensor |
|
|
| from modular_ts.synth import PERIOD_BANK, SynthGen, SynthStream, _z |
|
|
|
|
| def _fft_filter(e: Tensor, phi: Tensor, L: int) -> Tensor: |
| """y_t = sum_k phi^k e_{t-k} for each row (AR(1) filter via FFT), phi (n,).""" |
| n, T = e.shape |
| L = min(T, L) |
| k = phi.view(n, 1) ** torch.arange(L, device=e.device, dtype=torch.float32) |
| nfft = 1 << (T + L - 1).bit_length() |
| return torch.fft.irfft(torch.fft.rfft(e, nfft) * torch.fft.rfft(k, nfft), nfft)[:, :T] |
|
|
|
|
| class SynthGen2(SynthGen): |
| def __init__(self, *a, p_new: float = 0.55, p_aug: float = 0.5, **k): |
| super().__init__(*a, **k) |
| self.p_new, self.p_aug = p_new, p_aug |
|
|
| |
| def _regimes(self, n, rate_lo=1e-4, rate_hi=3e-3): |
| """Piecewise-constant regime index field (n,T) in [0,1) via Poisson change points.""" |
| T = self.T |
| rate = self.logU(n, 1, lo=rate_lo, hi=rate_hi) |
| jumps = (self.U(n, T) < rate).float() |
| seg = torch.cumsum(jumps, -1).long() % 64 |
| table = self.U(n, 64) |
| return torch.gather(table, 1, seg) |
|
|
| def ou_regime(self, n): |
| r = self._regimes(n) |
| mu = (r - 0.5) * self.U(n, 1, lo=0.0, hi=6.0) |
| vol = torch.exp(self.U(n, 1, lo=-1.0, hi=1.0) * (self._regimes(n) - 0.5) * 2) |
| theta = self.logU(n, lo=1e-3, hi=0.3) |
| e = theta.view(n, 1) * mu + 0.3 * vol * self.N(n, self.T) |
| return _fft_filter(e, 1 - theta, 2048) |
|
|
| def _period(self, n): |
| idx = torch.randint(0, len(PERIOD_BANK), (n,), generator=self.g, device=self.dev) |
| per = PERIOD_BANK.to(self.dev)[idx] |
| rnd = self.logU(n, lo=3.0, hi=max(4.0, self.T / 4)) |
| return torch.where(self.U(n) < 0.3, rnd, per) |
|
|
| def _phase(self, n, per): |
| """Cycle phase in [0,1) with optional slow period drift (phase-drifting cycles).""" |
| drift = torch.where(self.U(n) < 0.3, self.N(n) * 0.02, torch.zeros(n, device=self.dev)) |
| inst = 1.0 / per.view(n, 1) * (1 + drift.view(n, 1) * torch.tanh(self.rbf(n))) |
| return torch.frac(torch.cumsum(inst.expand(n, self.T), -1) + self.U(n, 1)) |
|
|
| def waveform(self, n): |
| per = self._period(n) |
| ph = self._phase(n, per) |
| kind = torch.randint(0, 4, (n, 1), generator=self.g, device=self.dev) |
| duty = self.U(n, 1, lo=0.1, hi=0.9) |
| saw = ph * 2 - 1 |
| sq = torch.where(ph < duty, 1.0, -1.0) |
| tri = 1 - 4 * (ph - 0.5).abs() |
| spk = torch.exp(-((ph - 0.5) ** 2) / (2 * self.logU(n, 1, lo=1e-4, hi=1e-2))) |
| w = torch.where(kind == 0, saw, torch.where(kind == 1, sq, torch.where(kind == 2, tri, spk))) |
| w = torch.where(self.U(n, 1) < 0.5, -w, w) |
| return w |
|
|
| def template(self, n): |
| """Tiled random daily profile (smooth + bumps) x optional weekly profile.""" |
| T = self.T |
| per = PERIOD_BANK.to(self.dev)[torch.randint(4, len(PERIOD_BANK), (n,), generator=self.g, device=self.dev)] |
| per = torch.where(self.U(n) < 0.6, torch.tensor([24.0, 96.0, 288.0, 48.0, 12.0, 7.0], device=self.dev)[ |
| torch.randint(0, 6, (n,), generator=self.g, device=self.dev)], per) |
| ph = self._phase(n, per) |
| |
| K = 4 |
| c = self.U(n, K, 1) |
| wdt = self.logU(n, K, 1, lo=0.02, hi=0.25) |
| a = self.N(n, K, 1).abs() * torch.where(self.U(n, K, 1) < 0.2, -1.0, 1.0) |
| d = torch.minimum((ph[:, None] - c).abs(), 1 - (ph[:, None] - c).abs()) |
| prof = (a * torch.exp(-0.5 * (d / wdt) ** 2)).sum(1) |
| h = torch.arange(1, 4, device=self.dev, dtype=torch.float32).view(1, 3, 1) |
| prof = prof + (0.3 * self.N(n, 3, 1) / h * torch.sin(2 * math.pi * h * ph[:, None] + self.U(n, 3, 1, hi=6.3))).sum(1) |
| |
| wk = self.U(n, 7, lo=0.5, hi=1.2) |
| wk = torch.where(self.U(n, 1) < 0.5, wk, torch.ones_like(wk)) |
| day = (torch.cumsum((ph[:, 1:] < ph[:, :-1]).float(), -1)).long() % 7 |
| day = torch.cat([torch.zeros(n, 1, dtype=torch.long, device=self.dev), day], -1) |
| mod = torch.gather(wk, 1, day) |
| sig = prof * mod |
| |
| sig = sig + self.B(0.5, n, 1) * 0.3 * _z(self.trend(n)) + self.logU(n, 1, lo=1e-3, hi=0.2) * self.N(n, T) |
| return sig |
|
|
| def steps(self, n): |
| r = self._regimes(n, 2e-4, 1e-2) |
| lv = torch.where(self.U(n, 1) < 0.5, torch.round(r * self.U(n, 1, lo=2, hi=20)), r) |
| return lv + self.logU(n, 1, lo=1e-3, hi=0.1) * self.N(n, self.T) |
|
|
| def growth(self, n): |
| T = self.T |
| tt = self.t / T |
| kind = torch.randint(0, 3, (n, 1), generator=self.g, device=self.dev) |
| |
| rate = torch.exp(0.8 * _z(self.rbf(n)).clamp(-3, 3)) * (1 + 0.3 * self.N(n, T).abs() * self.B(0.5, n, 1)) |
| cum = torch.cumsum(rate, -1) |
| |
| t0 = self.U(n, 1, lo=-0.2, hi=1.2) |
| k = self.logU(n, 1, lo=2.0, hi=30.0) |
| logi = torch.sigmoid(k * (tt - t0)) |
| gomp = torch.exp(-torch.exp(-k * (tt - t0))) |
| sig = torch.where(kind == 0, cum, torch.where(kind == 1, logi, gomp)) |
| sig = sig * (1 + self.logU(n, 1, lo=1e-3, hi=0.05) * self.N(n, T)) |
| return sig |
|
|
| |
| def augment(self, x: Tensor) -> Tensor: |
| n, T = x.shape |
| on = lambda p: self.U(n, 1) < p * self.p_aug * 2 |
| |
| x = torch.where(on(0.15), x.flip(-1), x) |
| |
| xi = torch.cumsum(x - x.mean(-1, keepdim=True), -1) |
| x = torch.where(on(0.06), xi, x) |
| xd = torch.cat([x[:, :1] * 0, x[:, 1:] - x[:, :-1]], -1) |
| x = torch.where(on(0.04), xd, x) |
| |
| x = torch.where(on(0.1), x * (1 + 0.5 * torch.tanh(self.rbf(n))), x) |
| |
| w = torch.randint(2, 16, (n, 1), generator=self.g, device=self.dev) |
| cs = torch.cumsum(torch.nn.functional.pad(x, (1, 0)), -1) |
| idx = torch.arange(1, T + 1, device=self.dev).view(1, T) |
| lo = (idx - w).clamp_min(0) |
| ma = (cs[:, 1:] - torch.gather(cs, 1, lo)) / (idx - lo).float() |
| x = torch.where(on(0.1), ma, x) |
| |
| k = torch.randint(2, 8, (n, 1), generator=self.g, device=self.dev) |
| hold = torch.gather(x, 1, (torch.arange(T, device=self.dev).view(1, T) // k) * k) |
| x = torch.where(on(0.05), hold, x) |
| |
| lam = self.U(n, 1, lo=0.2, hi=0.8) |
| xm = lam * _z(x) + (1 - lam) * _z(x.roll(1, 0)) |
| x = torch.where(on(0.1), xm, x) |
| |
| qv = torch.quantile(x[:, :: max(1, T // 512)], 0.8 + 0.19 * torch.rand(1, device=self.dev).item(), dim=-1, keepdim=True) |
| x = torch.where(on(0.05), torch.minimum(x, qv), x) |
| return x |
|
|
| def quantize(self, x: Tensor) -> Tensor: |
| """Round some series to a coarse grid (integer-valued / low-resolution sensors).""" |
| n = x.shape[0] |
| step = x.std(-1, keepdim=True).clamp_min(1e-6) * self.logU(n, 1, lo=0.02, hi=0.5) |
| q = torch.round(x / step) * step |
| return torch.where(self.U(n, 1) < 0.08, q, x) |
|
|
| |
| def univariate(self, n) -> Tensor: |
| base = super().univariate(n) |
| fam = torch.randint(0, 5, (n, 1), generator=self.g, device=self.dev) |
| new = torch.zeros(n, self.T, device=self.dev) |
| gens = (self.ou_regime, self.waveform, self.template, self.steps, self.growth) |
| for i, g in enumerate(gens): |
| new = torch.where(fam == i, _z(g(n)), new) |
| |
| comb = self.U(n, 1) < 0.35 |
| new = torch.where(comb, new + self.logU(n, 1, lo=0.05, hi=1.0) * _z(base), new) |
| use_new = self.U(n, 1) < self.p_new |
| x = torch.where(use_new, new, base) |
| x = self.augment(x) |
| x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) |
| return x |
|
|
| def affine(self, x: Tensor) -> Tensor: |
| return self.quantize(super().affine(x)) |
|
|
|
|
| class SynthStream2(SynthStream): |
| def __init__(self, T, tokens_per_batch, device="cpu", seed=0, patch=32, block=32, min_ctx=4, chunk=8, p_new=0.55): |
| super().__init__(T, tokens_per_batch, device, seed, patch, block, min_ctx, chunk) |
| self.gen = SynthGen2(T, device, seed, p_new=p_new) |
|
|
