1"""Student-t marginal head and masked negative log-likelihood (FP32 likelihood math).23Per future minute h: location mu_h = sigma_t * a_h, scale s_h = sigma_t * softplus(b_h) + eps, degrees of freedom4nu_h = 2 + softplus(c_h) (+ eps) so that the variance exists; sigma_t is the trailing scale computed at origin5(scale-equivariant parametrisation)."""6from __future__ import annotations78import math910import torch11import torch.nn as nn12import torch.nn.functional as F1314NU_MIN, SCALE_EPS = 2.05, 1e-7151617class StudentTHead(nn.Module):18 def __init__(self, d_model: int, learn_nu: bool = True, nu_fixed: float = 4.0):19 super().__init__()20 self.out = nn.Linear(d_model, 3)21 self.learn_nu, self.nu_fixed = learn_nu, nu_fixed22 nn.init.zeros_(self.out.bias); nn.init.normal_(self.out.weight, std=0.02)2324 def forward(self, h: torch.Tensor, sigma: torch.Tensor) -> dict[str, torch.Tensor]:25 """h: [B, H, d]; sigma: [B] -> dict(mu, scale, nu) each [B, H] in FP32."""26 raw = self.out(h).float()27 s = sigma.float().unsqueeze(-1)28 mu = s * raw[..., 0] * 0.129 scale = s * (F.softplus(raw[..., 1] + 0.5413) ) + SCALE_EPS # softplus(0.5413) ~= 1 -> starts at sigma_t30 if self.learn_nu:31 nu = NU_MIN + F.softplus(raw[..., 2] + 1.5) # starts near 2.05 + 1.7 = 3.7532 else:33 nu = torch.full_like(mu, self.nu_fixed)34 return {"mu": mu, "scale": scale, "nu": nu}353637def student_t_logpdf(y: torch.Tensor, mu: torch.Tensor, scale: torch.Tensor, nu: torch.Tensor) -> torch.Tensor:38 z = (y - mu) / scale39 return (torch.lgamma((nu + 1) / 2) - torch.lgamma(nu / 2) - 0.5 * torch.log(nu * math.pi) - torch.log(scale)40 - (nu + 1) / 2 * torch.log1p(z * z / nu))414243def masked_student_t_nll(y: torch.Tensor, mask: torch.Tensor, p: dict[str, torch.Tensor]) -> torch.Tensor:44 """Mean NLL over present targets (per-minute marginals). y, mask: [B, H]."""45 ll = student_t_logpdf(y.float(), p["mu"], p["scale"], p["nu"])46 m = mask.float()47 return -(ll * m).sum() / m.sum().clamp_min(1.0)484950def gaussian_nll(y: torch.Tensor, mask: torch.Tensor, p: dict[str, torch.Tensor]) -> torch.Tensor:51 z = (y.float() - p["mu"]) / p["scale"]52 ll = -0.5 * z * z - torch.log(p["scale"]) - 0.5 * math.log(2 * math.pi)53 m = mask.float()54 return -(ll * m).sum() / m.sum().clamp_min(1.0)555657def student_t_quantiles(p: dict[str, torch.Tensor], qs: torch.Tensor) -> torch.Tensor:58 """Quantiles via torch's StudentT icdf is unavailable; use scipy on CPU (evaluation only)."""59 import numpy as np60 from scipy import stats61 mu, sc, nu = (p[k].detach().cpu().double().numpy() for k in ("mu", "scale", "nu"))62 q = qs.cpu().double().numpy()63 out = stats.t.ppf(q[None, None, :], nu[..., None]) * sc[..., None] + mu[..., None]64 return torch.from_numpy(np.asarray(out, dtype=np.float32))656667def sample_student_t(p: dict[str, torch.Tensor], n: int, gen: torch.Generator | None = None) -> torch.Tensor:68 """Independent samples [n, B, H] (Stage A: marginals are independent across minutes)."""69 mu, sc, nu = p["mu"], p["scale"], p["nu"]70 shape = (n,) + mu.shape71 z = torch.randn(shape, generator=gen, device=mu.device)72 if not torch.isfinite(nu).all(): # Gaussian head (nu = inf): plain normal draws73 nu_f = torch.where(torch.isfinite(nu), nu, torch.full_like(nu, 1e6))74 g = torch.distributions.Gamma(nu_f / 2, nu_f / 2).sample((n,))75 g = torch.where(torch.isfinite(nu).unsqueeze(0), g, torch.ones_like(g))76 else:77 g = torch.distributions.Gamma(nu / 2, nu / 2).sample((n,)) # chi2_nu / nu78 return mu + sc * z / torch.sqrt(g)79