1"""MRPT: Multi-Resolution Patch Transformer with a direct 60-step decoder and Student-t / Gaussian / quantile heads."""2from __future__ import annotations34from dataclasses import dataclass, field, asdict56import torch7import torch.nn as nn89from ..data.windows import F_BAR, F_GLOBAL10from .patches import PatchEmbed, sinusoidal11from .fusion import SharedFusion, CrossAttentionFusion12from .regime import RegimeEmbedding13from .student_t import StudentTHead14from .joint_path import LowRankCopulaHead1516BRANCH_SPECS = {"A": dict(patch=5, res=0), "B": dict(patch=3, res=1), "C": dict(patch=2, res=2)}171819@dataclass20class MRPTConfig:21 d_model: int = 25622 n_layers: int = 623 n_heads: int = 824 d_ff: int = 102425 dropout: float = 0.126 dec_layers: int = 227 horizon: int = 6028 branches: list[str] = field(default_factory=lambda: ["A", "B", "C"])29 fusion: str = "shared" # shared | cross30 head: str = "student_t" # student_t | gaussian | quantile31 regime: str = "stats" # stats | learned | none32 n_symbols: int = 0 # > 0 enables the ticker-identity embedding (ablation only)33 copula_rank: int = 0 # > 0 enables the Stage-B joint path head34 quantiles: list[float] = field(default_factory=lambda: [0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975])3536 @classmethod37 def preset(cls, name: str, **kw) -> "MRPTConfig":38 base = {"small": dict(d_model=192, n_layers=4, n_heads=6, d_ff=768, dec_layers=2),39 "base": dict(d_model=384, n_layers=8, n_heads=8, d_ff=1536, dec_layers=2),40 "spec": dict(d_model=256, n_layers=6, n_heads=8, d_ff=1024, dec_layers=2),41 "tiny": dict(d_model=64, n_layers=2, n_heads=4, d_ff=128, dec_layers=1)}[name]42 base.update(kw)43 return cls(**base)444546class MRPT(nn.Module):47 def __init__(self, cfg: MRPTConfig):48 super().__init__()49 self.cfg = cfg50 d = cfg.d_model51 self.embeds = nn.ModuleDict({b: PatchEmbed(F_BAR, BRANCH_SPECS[b]["patch"], d, 3, cfg.dropout)52 for b in cfg.branches})53 if cfg.fusion == "shared":54 self.fusion = SharedFusion(d, cfg.n_layers, cfg.n_heads, cfg.d_ff, cfg.dropout)55 else:56 self.fusion = CrossAttentionFusion(d, cfg.n_layers, cfg.n_heads, cfg.d_ff, cfg.dropout, len(cfg.branches))57 self.glob_proj = nn.Sequential(nn.Linear(F_GLOBAL, d), nn.GELU(), nn.Linear(d, d))58 self.regime = RegimeEmbedding(F_GLOBAL, d, cfg.regime)59 self.fut_pos = nn.Embedding(cfg.horizon, d)60 self.fut_cal = nn.Linear(4, d)61 self.sym_emb = nn.Embedding(cfg.n_symbols, d) if cfg.n_symbols > 0 else None62 dec_layer = nn.TransformerDecoderLayer(d, cfg.n_heads, cfg.d_ff, cfg.dropout, activation="gelu",63 batch_first=True, norm_first=True)64 self.decoder = nn.TransformerDecoder(dec_layer, cfg.dec_layers)65 self.dec_norm = nn.LayerNorm(d)66 if cfg.head == "student_t":67 self.head = StudentTHead(d, learn_nu=True)68 elif cfg.head == "gaussian":69 self.head = StudentTHead(d, learn_nu=False, nu_fixed=float("inf"))70 elif cfg.head == "quantile":71 self.head = nn.Linear(d, len(cfg.quantiles))72 else:73 raise ValueError(cfg.head)74 self.copula = LowRankCopulaHead(d, cfg.horizon, cfg.copula_rank) if cfg.copula_rank > 0 else None7576 def n_params(self, trainable_only: bool = True) -> int:77 return sum(p.numel() for p in self.parameters() if p.requires_grad or not trainable_only)7879 def encode(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:80 streams = []81 for b in self.cfg.branches:82 tok, valid = self.embeds[b](batch["x" + b], batch["m" + b].bool(), BRANCH_SPECS[b]["res"])83 streams.append((tok, valid))84 # global context token (history-only statistics) prepended to the first stream85 g = self.glob_proj(batch["glob"].float()).unsqueeze(1)86 t0, v0 = streams[0]87 streams[0] = (torch.cat([g, t0], 1), torch.cat([torch.ones_like(v0[:, :1]), v0], 1))88 return self.fusion(streams)8990 def decode(self, memory: torch.Tensor, mem_valid: torch.Tensor, batch: dict[str, torch.Tensor]) -> torch.Tensor:91 B = memory.shape[0]92 H = self.cfg.horizon93 q = self.fut_pos.weight[:H].unsqueeze(0).expand(B, H, -1) + self.fut_cal(batch["fut_cal"].float())94 r = self.regime(batch["glob"].float())95 if r is not None:96 q = q + r.unsqueeze(1)97 if self.sym_emb is not None:98 q = q + self.sym_emb(batch["sym"]).unsqueeze(1)99 kpm = ~mem_valid100 if kpm.all(dim=1).any():101 kpm = kpm.clone(); kpm[kpm.all(dim=1), 0] = False102 h = self.decoder(q, memory, memory_key_padding_mask=kpm)103 return self.dec_norm(h)104105 def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:106 memory, valid = self.encode(batch)107 h = self.decode(memory, valid, batch)108 if self.cfg.head == "quantile":109 s = batch["sigma"].float().view(-1, 1, 1)110 q = self.head(h).float() * s111 return {"quantiles": torch.cummax(q, dim=-1).values if False else q, "hidden": h}112 out = self.head(h, batch["sigma"])113 out["hidden"] = h114 return out115