1"""Batched prediction over an OriginIndex: Student-t parameters per future minute, Monte Carlo cumulative-return2samples at each horizon (Stage A: independent minutes; Stage B: copula-correlated), plus metadata for scoring."""3from __future__ import annotations45import numpy as np6import torch7from torch.utils.data import DataLoader89from ..data.windows import OriginIndex10from ..train.sampling import OriginDataset, collate11from ..train.trainer import pick_device, to_device12from ..models.student_t import sample_student_t13from ..models.joint_path import samples_from_copula141516@torch.no_grad()17def predict_index(model, index: OriginIndex, horizons: list[int], n_samples: int = 256, batch_size: int = 128,18 feature_set: str = "full", shuffle_minutes: bool = False, device: str | None = None, num_workers: int = 4,19 joint: bool = False, seed: int = 0, keep_paths: bool = False) -> dict[str, np.ndarray]:20 dev = pick_device(device)21 model.to(dev).eval()22 ds = OriginDataset(index, feature_set, shuffle_minutes)23 loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate)24 gen = torch.Generator(device=dev); gen.manual_seed(seed)25 out = {k: [] for k in ("mu", "scale", "nu", "y", "y_mask", "sigma", "origin", "sym")}26 cum = {H: [] for H in horizons}; realised = {H: [] for H in horizons}; rv = {H: [] for H in horizons}27 rv_pred = {H: [] for H in horizons}28 path_sc = {k: [] for k in ("energy", "variogram", "mdd_crps", "rv_crps")}29 for b in loader:30 bd = to_device(b, dev)31 o = model(bd)32 if "quantiles" in o:33 raise NotImplementedError("sample-based scoring needs a parametric head")34 mu, sc, nu = o["mu"], o["scale"], o["nu"]35 if joint and getattr(model, "copula", None) is not None:36 z = model.copula.sample_z(o["hidden"], n_samples, gen)37 S = samples_from_copula(z, mu, sc, nu) # [n, B, H] cpu38 else:39 S = sample_student_t({"mu": mu, "scale": sc, "nu": nu}, n_samples, gen).cpu()40 S = S.numpy()41 ym = b["y_mask"].numpy().astype(bool); y = b["y"].numpy()42 for H in horizons:43 Sm = S[:, :, :H] * ym[None, :, :H] # missing future minutes contribute nothing (same as realised)44 cum[H].append(Sm.sum(-1)); realised[H].append((y[:, :H] * ym[:, :H]).sum(-1))45 rv[H].append(((y[:, :H] * ym[:, :H]) ** 2).sum(-1)); rv_pred[H].append((Sm ** 2).sum(-1))46 if keep_paths:47 from ..eval.scores import energy_score, variogram_score, crps_samples, path_stats48 Y = y * ym; Sp = S * ym[None]49 path_sc["energy"].append(energy_score(Y, Sp)); path_sc["variogram"].append(variogram_score(Y, Sp))50 st = path_stats(Sp); ob = path_stats(Y[None])51 path_sc["mdd_crps"].append(crps_samples(ob["mdd"][0], st["mdd"])); path_sc["rv_crps"].append(crps_samples(ob["rv"][0], st["rv"]))52 for k in ("mu", "scale", "nu"):53 out[k].append(o[k].float().cpu().numpy())54 out["y"].append(y); out["y_mask"].append(ym); out["sigma"].append(b["sigma"].numpy())55 out["origin"].append(b["origin"].numpy()); out["sym"].append(b["sym"].numpy())56 res = {k: np.concatenate(v) for k, v in out.items()}57 for H in horizons:58 res[f"cum_samples_{H}"] = np.concatenate(cum[H], axis=1) # [n, N]59 res[f"cum_{H}"] = np.concatenate(realised[H]); res[f"rv_{H}"] = np.concatenate(rv[H])60 res[f"rv_samples_{H}"] = np.concatenate(rv_pred[H], axis=1)61 if keep_paths:62 for k, v in path_sc.items():63 res[k] = np.concatenate(v)64 return res65