"""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 # ------------------------------------------------------------------ families 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 # segment id 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) # (n,T) in [0,1) # profile = sum of K random bumps on the circle + low-order Fourier terms 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) # weekly modulation for intraday periods 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 # slowly varying level and tiny noise (near-deterministic series) 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) # cumulative counts: cumsum of positive, slowly varying rate (+ noise) 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) # logistic / gompertz 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 # ------------------------------------------------------------------ augmentations def augment(self, x: Tensor) -> Tensor: n, T = x.shape on = lambda p: self.U(n, 1) < p * self.p_aug * 2 # noqa: E731 # time reversal x = torch.where(on(0.15), x.flip(-1), x) # integration / differencing 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) # amplitude modulation x = torch.where(on(0.1), x * (1 + 0.5 * torch.tanh(self.rbf(n))), x) # random smoothing conv filter (moving average of random width via cumsum trick) 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) # sample-and-hold resampling artefacts 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) # TSMixup with a rolled partner 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) # censoring at a random quantile (one side) 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) # ------------------------------------------------------------------ mixture 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) # optional combination of the new family with a v1 component (additive) 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)