"""Streamed, fully vectorised synthetic time-series generator (torch; runs on GPU). Univariate bases (each switched on per series with its own probability, summed or multiplied): * KernelSynth-style GP components via random Fourier features: RBF (log-uniform length-scale), rational-quadratic (scale-mixture of RBFs), periodic kernels at frequency-typical periods with random harmonics, linear (trend), white noise; * TSI-style trends: linear / quadratic / piecewise-linear with change points / damped; * AR(1) / random walk (FFT convolution with phi^k kernels); * heteroscedastic noise (log-GP volatility, GARCH-like clustering), student-t tails; * level shifts, spikes / outliers; * intermittent (Croston-like zero-inflated count) and bursty (Poisson bursts with exponential decay) observability-style signals; * non-negativity transforms, random offset / scale over many orders of magnitude. Multivariatizers (Chronos-2 / CauKer style): latent univariate drivers mixed into each row by a random sparse matrix through random nonlinearities with lead-lag delays, plus idiosyncratic noise; optional covariate tasks where some rows are past-only or known-future covariates (incl. 0/1 event / calendar indicators) and targets respond to them through random lagged nonlinear effects. Masks: iid missing points, contiguous gaps, late series starts (left padding). `SynthStream(...).next_batch()` returns dict(values, observed, roles, s) ready for `ModularTS.forward_train`. """ from __future__ import annotations import math import torch from torch import Tensor PERIOD_BANK = torch.tensor([4, 5, 7, 12, 24, 48, 52, 60, 96, 144, 168, 288, 336, 365, 672, 1008, 1440, 2016], dtype=torch.float32) def _std(x: Tensor) -> Tensor: return x.std(-1, keepdim=True).clamp_min(1e-6) def _z(x: Tensor) -> Tensor: return (x - x.mean(-1, keepdim=True)) / _std(x) class SynthGen: def __init__(self, T: int, device: str | torch.device = "cpu", seed: int = 0, n_rff: int = 32): self.T = T self.dev = torch.device(device) self.g = torch.Generator(device=self.dev) self.g.manual_seed(seed) self.n_rff = n_rff self.t = torch.arange(T, device=self.dev, dtype=torch.float32) # -- random helpers def U(self, *shape, lo=0.0, hi=1.0): return lo + (hi - lo) * torch.rand(*shape, generator=self.g, device=self.dev) def N(self, *shape): return torch.randn(*shape, generator=self.g, device=self.dev) def B(self, p, *shape): return (self.U(*shape) < p).float() def logU(self, *shape, lo, hi): return torch.exp(self.U(*shape, lo=math.log(lo), hi=math.log(hi))) # -- components (n, T) def rff(self, n, lengthscale: Tensor) -> Tensor: M = self.n_rff w = self.N(n, M, 1) / lengthscale.view(n, 1, 1) b = self.U(n, M, 1, hi=2 * math.pi) a = self.N(n, M, 1) return (a * torch.cos(w * self.t + b)).sum(1) * math.sqrt(2.0 / M) def rbf(self, n): return self.rff(n, self.logU(n, lo=self.T / 400, hi=self.T / 1.5)) def rq(self, n): # scale mixture: sum of RBFs at several length-scales out = 0 for _ in range(3): out = out + self.U(n, 1) * self.rff(n, self.logU(n, lo=self.T / 1000, hi=self.T / 4)) return out def periodic(self, n): k = len(PERIOD_BANK) idx = torch.randint(0, k, (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 / 3)) per = torch.where(self.U(n) < 0.25, rnd, per) per = per * self.U(n, lo=0.97, hi=1.03) * torch.where(self.U(n) < 0.1, 1.0 + 0.2 * self.N(n), torch.ones(n, device=self.dev)) H = 6 h = torch.arange(1, H + 1, device=self.dev, dtype=torch.float32).view(1, H, 1) decay = self.U(n, 1, 1, lo=0.3, hi=1.5) amp = self.N(n, H, 1) / h.pow(decay) * self.B(0.8, n, H, 1).clamp_min((h == 1).float()) ph = self.U(n, H, 1, hi=2 * math.pi) sig = (amp * torch.sin(2 * math.pi * h * self.t / per.view(n, 1, 1) + ph)).sum(1) # optional slowly varying amplitude (multiplicative modulation) mod = 1 + 0.3 * self.B(0.3, n, 1) * torch.tanh(self.rbf(n)) return sig * mod def trend(self, n): tt = self.t / self.T lin = self.N(n, 1) * tt quad = self.N(n, 1) * 0.5 * tt**2 * self.B(0.3, n, 1) # piecewise linear with up to 4 change points K = 4 cps = self.U(n, K, 1) slopes = self.N(n, K, 1) * self.B(0.5, n, K, 1) pw = (slopes * torch.relu(tt.view(1, 1, -1) - cps)).sum(1) damp = self.N(n, 1) * (1 - torch.exp(-tt / self.U(n, 1, lo=0.05, hi=0.5))) * self.B(0.2, n, 1) return lin + quad + pw + damp def ar(self, n): phi = torch.where(self.U(n) < 0.3, torch.ones(n, device=self.dev), self.U(n, lo=0.3, hi=0.995)) e = self.N(n, self.T) L = min(self.T, 1024) k = phi.view(n, 1) ** torch.arange(L, device=self.dev, dtype=torch.float32) nfft = 1 << (self.T + L - 1).bit_length() y = torch.fft.irfft(torch.fft.rfft(e, nfft) * torch.fft.rfft(k, nfft), nfft)[:, : self.T] rw = torch.cumsum(e, -1) return torch.where((phi == 1).view(n, 1), rw, y) def noise(self, n): z = self.N(n, self.T) # heavy tails: scale mixture of normals (generator-friendly stand-in for student-t) heavy = self.B(0.3, n, 1) tz = z * (1 + heavy * (self.N(n, self.T).abs() ** 2) * 0.5) vol = torch.where(self.U(n, 1) < 0.25, torch.exp(0.7 * self.rbf(n).clamp(-3, 3)), torch.ones(n, 1, device=self.dev)) return tz * vol def level_shifts(self, n): K = 3 loc = self.U(n, K, 1) * self.T mag = self.N(n, K, 1) * self.B(0.6, n, K, 1) return (mag * (self.t.view(1, 1, -1) > loc).float()).sum(1) def spikes(self, n): p = self.logU(n, 1, lo=1e-3, hi=2e-2) return (self.U(n, self.T) < p).float() * self.N(n, self.T) * self.U(n, 1, lo=3, hi=10) def bursts(self, n): rate = self.logU(n, 1, lo=1e-3, hi=2e-2) ev = (self.U(n, self.T) < rate).float() * self.logU(n, self.T, lo=1.0, hi=20.0) tau = self.logU(n, lo=2.0, hi=60.0) L = min(self.T, 512) k = torch.exp(-torch.arange(L, device=self.dev).float().view(1, -1) / tau.view(n, 1)) nfft = 1 << (self.T + L - 1).bit_length() return torch.fft.irfft(torch.fft.rfft(ev, nfft) * torch.fft.rfft(k, nfft), nfft)[:, : self.T] # -- univariate def univariate(self, n) -> Tensor: z = torch.zeros(n, self.T, device=self.dev) on = lambda p: self.B(p, n, 1) # noqa: E731 w = lambda: self.logU(n, 1, lo=0.1, hi=2.0) # noqa: E731 trend = on(0.6) * w() * _z(self.trend(n)) seas = on(0.7) * w() * _z(self.periodic(n)) + on(0.25) * w() * _z(self.periodic(n)) smooth = on(0.5) * w() * _z(self.rbf(n)) + on(0.25) * w() * _z(self.rq(n)) arp = on(0.4) * w() * _z(self.ar(n)) sig = trend + seas + smooth + arp + z mult = on(0.25) sig = torch.where(mult > 0, (1 + 0.4 * torch.tanh(trend + smooth)) * seas + trend + arp, sig) sig = sig + on(0.15) * w() * self.level_shifts(n) noise_lvl = self.logU(n, 1, lo=0.01, hi=1.0) sig = sig + noise_lvl * self.noise(n) sig = sig + on(0.12) * self.spikes(n) # intermittent / bursty regimes (non-negative count-like) inter = on(0.10) lam = torch.exp(0.8 * _z(sig).clamp(-4, 4) + self.U(n, 1, lo=-3, hi=1.5)) p0 = self.U(n, 1, lo=0.2, hi=0.9) cnt = torch.poisson(lam, generator=self.g) * (self.U(n, self.T) > p0).float() sig = torch.where(inter > 0, cnt, sig) burst = on(0.07) sig = torch.where(burst > 0, self.U(n, 1, lo=0, hi=2) + 0.1 * self.noise(n).abs() + self.bursts(n), sig) # non-negativity transforms on some continuous series nn_ = on(0.2) * (1 - inter) * (1 - burst) sig = torch.where(nn_ > 0, torch.nn.functional.softplus(sig * self.U(n, 1, lo=0.5, hi=3)), sig) return sig def affine(self, x: Tensor) -> Tensor: n = x.shape[0] scale = self.logU(n, 1, lo=1e-2, hi=1e4) off = self.N(n, 1) * scale * self.logU(n, 1, lo=0.01, hi=20.0) * self.B(0.6, n, 1) sgn = torch.where(self.U(n, 1) < 0.5, -1.0, 1.0) pos = (x.min(-1, keepdim=True).values >= 0).float() # keep count-like series non-negative return torch.where(pos > 0, x * scale, sgn * x * scale + off) # -- multivariate group batch def group_batch(self, Bg: int, V: int, cov_prob: float = 0.35): """Returns values (Bg,V,T), roles (Bg,V).""" T = self.T R = Bg * V if V == 1: vals = self.affine(self.univariate(Bg)).view(Bg, 1, T) return vals, torch.zeros(Bg, 1, dtype=torch.long, device=self.dev) K = max(1, V // 2) lat = _z(self.univariate(Bg * K)).view(Bg, K, T) # lead-lag delays per (row, latent) lag = torch.randint(0, 48, (Bg, V, K), generator=self.g, device=self.dev) lag = lag * self.B(0.5, Bg, V, K).long() idx = (self.t.long().view(1, 1, 1, T) - lag[..., None]).clamp_min(0) shifted = torch.gather(lat[:, None].expand(Bg, V, K, T), 3, idx) W = self.N(Bg, V, K, 1) * self.B(0.6, Bg, V, K, 1) W[:, :, 0] = W[:, :, 0] + 0.5 mixed = (W * shifted).sum(2) f = torch.randint(0, 5, (Bg, V, 1), generator=self.g, device=self.dev) nl = torch.stack([mixed, torch.tanh(mixed), mixed.abs(), torch.relu(mixed), mixed + 0.3 * mixed**2], 0) mixed = torch.gather(nl, 0, f[None].expand(1, Bg, V, T))[0] own = _z(self.univariate(R)).view(Bg, V, T) a = self.U(Bg, V, 1, lo=0.0, hi=1.0) indep = self.B(0.15, Bg, 1, 1) rows = torch.where(indep > 0, own, (1 - a) * _z(mixed) + a * own) roles = torch.zeros(Bg, V, dtype=torch.long, device=self.dev) # covariate tasks cov_g = self.U(Bg) < cov_prob if cov_g.any(): n_cov = torch.randint(1, V, (Bg,), generator=self.g, device=self.dev) ridx = torch.arange(V, device=self.dev).view(1, V) is_cov = (ridx >= (V - n_cov).view(Bg, 1)) & cov_g.view(Bg, 1) known = is_cov & (self.U(Bg, V) < 0.6) # some covariates are 0/1 event indicators (holiday / promo / calendar) ev_rate = self.logU(Bg, V, 1, lo=0.005, hi=0.15) per = PERIOD_BANK.to(self.dev)[torch.randint(0, len(PERIOD_BANK), (Bg, V, 1), generator=self.g, device=self.dev)] cal = ((self.t.view(1, 1, T) % per) < torch.clamp(per * self.U(Bg, V, 1, lo=0.05, hi=0.3), min=1)).float() evt = (self.U(Bg, V, T) < ev_rate).float() ind = torch.where(self.U(Bg, V, 1) < 0.5, cal, evt) binary = is_cov & (self.U(Bg, V) < 0.4) cov_vals = torch.where(binary[..., None], ind, rows) # target response to covariates: lagged nonlinear effects beta = self.N(Bg, V, 1) * self.U(Bg, 1, 1, lo=0.5, hi=2.5) clag = torch.randint(0, 8, (Bg, V, 1), generator=self.g, device=self.dev) * (~known[..., None]).long() cidx = (self.t.long().view(1, 1, T) - clag).clamp_min(0) cl = torch.gather(cov_vals, 2, cidx) eff_in = torch.where(binary[..., None], cl, _z(cl)) eff = (beta * torch.where(self.U(Bg, V, 1) < 0.5, eff_in, torch.tanh(eff_in)) * is_cov[..., None].float()).sum(1, keepdim=True) tgt_mask = (~is_cov)[..., None].float() rows = torch.where(is_cov[..., None], cov_vals, rows + tgt_mask * eff * cov_g.view(Bg, 1, 1).float()) roles = torch.where(is_cov, torch.where(known, 2, 1), roles) bin_rows = binary.view(R) else: bin_rows = torch.zeros(R, dtype=torch.bool, device=self.dev) flat = rows.reshape(R, T) aff = self.affine(flat) flat = torch.where(bin_rows[:, None], flat, aff) return flat.view(Bg, V, T), roles def masks(self, vals: Tensor, roles: Tensor) -> Tensor: Bg, V, T = vals.shape obs = torch.ones(Bg, V, T, dtype=torch.bool, device=self.dev) # iid missing pm = self.U(Bg, V, 1, hi=0.15) * self.B(0.2, Bg, V, 1) obs &= self.U(Bg, V, T) >= pm # contiguous gaps gap = self.B(0.15, Bg, V, 1) > 0 g0 = self.U(Bg, V, 1) * T gl = self.logU(Bg, V, 1, lo=8, hi=T / 4) in_gap = (self.t.view(1, 1, T) >= g0) & (self.t.view(1, 1, T) < g0 + gl) obs &= ~(gap & in_gap) # late start (left padding), per group late = self.B(0.25, Bg, 1, 1) > 0 st = (self.U(Bg, 1, 1) * T * 0.8) obs &= ~(late & (self.t.view(1, 1, T) < st)) return obs class SynthStream: """Endless batches of roughly `tokens_per_batch` encoder patch tokens.""" V_CHOICES = (1, 2, 4, 8, 16) V_PROBS = (0.45, 0.1, 0.15, 0.15, 0.15) def __init__(self, T: int, tokens_per_batch: int, device="cpu", seed=0, patch=32, block=32, min_ctx=4, chunk=8): self.gen = SynthGen(T, device, seed) self.chunk = chunk # batches generated per call (amortises kernel-launch overhead on GPU) self.cache: dict[int, list] = {} self.T, self.patch, self.N = T, patch, T // patch self.tok = tokens_per_batch self.min_ctx = min_ctx self.dev = torch.device(device) self.rng = torch.Generator().manual_seed(seed + 1) def next_batch(self, V: int | None = None) -> dict: if V is None: i = torch.multinomial(torch.tensor(self.V_PROBS), 1, generator=self.rng).item() V = self.V_CHOICES[i] Bg = max(1, self.tok // (V * self.N)) if not self.cache.get(V): vals, roles = self.gen.group_batch(Bg * self.chunk, V) obs = self.gen.masks(vals, roles) vals = torch.nan_to_num(vals, nan=0.0, posinf=0.0, neginf=0.0) s = torch.randint(self.min_ctx, self.N - 1, (Bg * self.chunk,), generator=self.gen.g, device=self.dev) self.cache[V] = [ {"values": vals[i * Bg:(i + 1) * Bg], "observed": obs[i * Bg:(i + 1) * Bg], "roles": roles[i * Bg:(i + 1) * Bg], "s": s[i * Bg:(i + 1) * Bg]} for i in range(self.chunk) ] return self.cache[V].pop()