1from __future__ import annotations23import torch45from ..models.student_t import masked_student_t_nll, gaussian_nll678def pinball_loss(y: torch.Tensor, mask: torch.Tensor, q: torch.Tensor, levels: list[float]) -> torch.Tensor:9 """y, mask: [B, H]; q: [B, H, Q]."""10 lv = torch.tensor(levels, device=y.device, dtype=torch.float32).view(1, 1, -1)11 diff = y.float().unsqueeze(-1) - q12 loss = torch.maximum(lv * diff, (lv - 1) * diff).mean(-1)13 m = mask.float()14 return (loss * m).sum() / m.sum().clamp_min(1.0)151617def forecast_loss(out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], head: str,18 levels: list[float] | None = None, aux_vol_weight: float = 0.0) -> tuple[torch.Tensor, dict]:19 y, m = batch["y"].float(), batch["y_mask"]20 if head == "student_t":21 loss = masked_student_t_nll(y, m, out)22 elif head == "gaussian":23 loss = gaussian_nll(y, m, out)24 else:25 loss = pinball_loss(y, m, out["quantiles"], levels)26 logs = {"nll": float(loss.detach())}27 if aux_vol_weight > 0 and head != "quantile":28 # auxiliary: predicted variance of the H-sum vs realised squared sum (QLIKE-like, stabilised)29 var = (out["scale"] ** 2 * (out["nu"] / (out["nu"] - 2)).clamp(max=50)).sum(-1) if head == "student_t" \30 else (out["scale"] ** 2).sum(-1)31 rv = ((y * m.float()) ** 2).sum(-1)32 aux = (torch.log(var + 1e-12) + rv / (var + 1e-12)).mean()33 loss = loss + aux_vol_weight * aux34 logs["aux_vol"] = float(aux.detach())35 return loss, logs36