1"""Stage B: dependence-aware joint path head (Gaussian copula with low-rank + diagonal correlation over the H future2minutes, Student-t marginals from Stage A). Training: copula log-likelihood of the probability-integral transforms of3the observed targets with the marginals frozen. Sampling: z = L eps + sqrt(psi) eta, standardised, u = Phi(z),4y_h = F_t^{-1}(u_h; mu_h, s_h, nu_h). Missing targets are marginalised exactly by replacing their rows/columns of the5correlation matrix with the identity. All linear algebra in FP32 on CPU when the device lacks Cholesky support."""6from __future__ import annotations78import math910import numpy as np11import torch12import torch.nn as nn13import torch.nn.functional as F14from scipy import stats1516SQRT2 = math.sqrt(2.0)171819class LowRankCopulaHead(nn.Module):20 def __init__(self, d_model: int, horizon: int, rank: int = 4):21 super().__init__()22 self.rank, self.H = rank, horizon23 self.load = nn.Linear(d_model, rank)24 self.logpsi = nn.Linear(d_model, 1)25 nn.init.normal_(self.load.weight, std=0.02); nn.init.zeros_(self.load.bias)26 nn.init.zeros_(self.logpsi.weight); nn.init.zeros_(self.logpsi.bias)2728 def correlation(self, h: torch.Tensor) -> torch.Tensor:29 """h: [B, H, d] -> correlation matrices [B, H, H] (FP32)."""30 lam = self.load(h.float()) # [B, H, k]31 psi = F.softplus(self.logpsi(h.float())).squeeze(-1) + 1e-3 # [B, H]32 cov = lam @ lam.transpose(1, 2) + torch.diag_embed(psi)33 d = torch.rsqrt(torch.diagonal(cov, dim1=1, dim2=2))34 return cov * d.unsqueeze(1) * d.unsqueeze(2)3536 @staticmethod37 def _masked_corr(R: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:38 m = mask.float()39 return R * m.unsqueeze(1) * m.unsqueeze(2) + torch.diag_embed(1.0 - m)4041 def nll(self, h: torch.Tensor, z: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:42 """Gaussian-copula negative log density (per observed dimension): z = Phi^{-1}(u) [B, H], mask [B, H]."""43 R = self._masked_corr(self.correlation(h), mask)44 dev = R.device45 Rc = R.cpu() if dev.type == "mps" else R46 zc = (z.float() * mask.float()).unsqueeze(-1)47 zc = zc.cpu() if dev.type == "mps" else zc48 L, info = torch.linalg.cholesky_ex(Rc)49 logdet = 2.0 * torch.log(torch.diagonal(L, dim1=1, dim2=2)).sum(-1)50 sol = torch.cholesky_solve(zc, L)51 quad = (zc * sol).sum(dim=(1, 2))52 # copula density: -0.5 logdet - 0.5 z'(R^-1 - I) z53 ll = -0.5 * logdet - 0.5 * (quad - (zc * zc).sum(dim=(1, 2)))54 n = mask.float().sum(-1).clamp_min(1.0)55 n = n.cpu() if dev.type == "mps" else n56 return -(ll / n).mean().to(dev)5758 def sample_z(self, h: torch.Tensor, n: int, gen: torch.Generator | None = None) -> torch.Tensor:59 """Correlated standard normals [n, B, H] from the low-rank structure (no Cholesky needed)."""60 lam = self.load(h.float()); psi = F.softplus(self.logpsi(h.float())).squeeze(-1) + 1e-361 B, H, k = lam.shape62 eps = torch.randn((n, B, k), generator=gen, device=lam.device)63 eta = torch.randn((n, B, H), generator=gen, device=lam.device)64 z = torch.einsum("bhk,nbk->nbh", lam, eps) + torch.sqrt(psi).unsqueeze(0) * eta65 sd = torch.sqrt((lam * lam).sum(-1) + psi).unsqueeze(0)66 return z / sd676869def pit_student_t(y: np.ndarray, mu: np.ndarray, scale: np.ndarray, nu: np.ndarray) -> np.ndarray:70 y, mu, scale, nu = (np.asarray(a, dtype=np.float64) for a in (y, mu, scale, nu))71 return stats.t.cdf((y - mu) / scale, nu)727374def z_from_pit(u: np.ndarray) -> np.ndarray:75 u = np.clip(u, 1e-6, 1 - 1e-6)76 return stats.norm.ppf(u)777879def samples_from_copula(z: torch.Tensor, mu: torch.Tensor, scale: torch.Tensor, nu: torch.Tensor) -> torch.Tensor:80 """z [n, B, H] correlated normals -> Student-t marginal samples [n, B, H] via the PIT (CPU numpy)."""81 # float64 throughout: scipy's t.ppf returns NaN for float32 inputs near 0.5 (observed with scipy 1.17)82 u = (0.5 * (1.0 + torch.erf(z.detach().cpu().double() / SQRT2))).clamp(1e-6, 1 - 1e-6).numpy()83 mu_, sc_, nu_ = (t.detach().cpu().double().numpy() for t in (mu, scale, nu))84 y = stats.t.ppf(u, nu_[None]) * sc_[None] + mu_[None]85 return torch.from_numpy(np.asarray(y, dtype=np.float32))86