Upload via Autoresearch export_data

51bb0b7 verified 12 days ago

14.5 kB

"""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()

	"""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()