| """GiftEvalPretrain stream v2 (Modular-TS v3): resolution + mixup augmentation on real data. |
| |
| Adds to FastGiftEvalPretrain (run1): |
| * TTM-style "diverse resolution sampling" (Ekambaram et al., arXiv 2401.03955): with prob |
| `ds_p` a univariate batch is drawn at a coarser resolution, mean-pooling k in {2,4,8} |
| consecutive raw steps (NaN-aware), creating new frequencies from the same corpus; |
| * TSMixup (Chronos, arXiv 2403.07815): with prob `mix_p` rows are replaced by a convex |
| combination of 2-3 standardised real series (Dirichlet weights), rescaled; |
| * flatter dataset temperature (ds_temp 0.3 instead of 0.5) so small datasets are seen more. |
| Contamination filter and GPU-resident layout are inherited unchanged. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| from modular_ts.realdata import FastGiftEvalPretrain |
|
|
|
|
| class FastGiftEvalPretrain2(FastGiftEvalPretrain): |
| def __init__(self, *a, ds_p=0.3, ds_factors=(2, 4, 8), mix_p=0.2, **k): |
| k.setdefault("ds_temp", 0.3) |
| super().__init__(*a, **k) |
| self.ds_p, self.ds_factors, self.mix_p = ds_p, ds_factors, mix_p |
|
|
| def _gather_k(self, off, L, kf): |
| """Gather T*kf raw steps ending at a random point and mean-pool by kf.""" |
| T, R = self.T, off.shape[0] |
| TK = T * kf |
| st = torch.where(L >= TK, (self._rand(R) * (L - TK + 1)).long(), L - TK) |
| ar = torch.arange(TK, device=self.dev) |
| pos = st[:, None] + ar[None] |
| valid = (pos >= 0) & (pos < L[:, None]) |
| idx = (off[:, None] + pos.clamp(min=0)).clamp_max(self.flat.numel() - 1) |
| v = self.flat[idx] |
| ok = valid & torch.isfinite(v) |
| v = torch.where(ok, v, torch.zeros_like(v)).view(R, T, kf) |
| c = ok.view(R, T, kf).sum(-1) |
| out = v.sum(-1) / c.clamp_min(1) |
| obs = c > 0 |
| return torch.where(obs, out, torch.zeros_like(out)), obs |
|
|
| def _mixup(self, v, o): |
| R, T = v.shape |
| both = o & o.roll(1, 0) & o.roll(2, 0) |
| def zs(x, m): |
| n = m.sum(-1, keepdim=True).clamp_min(1) |
| mu = (x * m).sum(-1, keepdim=True) / n |
| sd = (((x - mu) ** 2) * m).sum(-1, keepdim=True).div(n).sqrt().clamp_min(1e-6) |
| return (x - mu) / sd, sd, mu |
| z0, sd0, mu0 = zs(v, o) |
| z1, _, _ = zs(v.roll(1, 0), o.roll(1, 0)) |
| z2, _, _ = zs(v.roll(2, 0), o.roll(2, 0)) |
| w = torch.distributions.Dirichlet(torch.full((3,), 1.5, device=self.dev)).sample((R,)) |
| w[:, 2] *= (self._rand(R) < 0.5).float() |
| w = w / w.sum(-1, keepdim=True) |
| m = w[:, :1] * z0 + w[:, 1:2] * z1 + w[:, 2:3] * z2 |
| m = m * sd0 + mu0 |
| sel = (self._rand(R) < self.mix_p)[:, None] |
| return torch.where(sel, torch.where(both, m, 0.0), v), torch.where(sel, both, o) |
|
|
| @torch.no_grad() |
| def next_batch(self, V=None): |
| if (self.mv is not None and self.rng.random() < self.mv_frac) or self.rng.random() >= self.ds_p: |
| b = super().next_batch(V) |
| if b["values"].shape[1] == 1 and self.mix_p > 0: |
| v, o = self._mixup(b["values"][:, 0], b["observed"][:, 0]) |
| b["values"], b["observed"] = v[:, None], o[:, None] |
| b["s"] = self._s(b["observed"]) |
| return b |
| kf = self.rng.choice(self.ds_factors) |
| Bg = max(1, self.tok // self.N) |
| c = self._pick(self.u_w, self.u_start, self.u_cnt, Bg) |
| off, L = self.u_off[c], self.u_len[c] |
| v, o = self._gather_k(off, L, kf) |
| o = o & (torch.arange(self.T, device=self.dev)[None] >= (self.T - self._keep_len(Bg))[:, None]) |
| v = torch.where(o, v, torch.zeros_like(v)) |
| if self.mix_p > 0: |
| v, o = self._mixup(v, o) |
| v, o = v[:, None], o[:, None] |
| return {"values": v, "observed": o, "roles": torch.zeros(Bg, 1, dtype=torch.long, device=self.dev), "s": self._s(o)} |
|
|
