Upload code/metrics.py

code/metrics.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+@dataclass
+class LossConfig:
+    rec_weight: float = 1.0
+    z_weight: float = 2.0
+    z_bin_weight: float = 0.2
+    z_candidate_weight: float = 0.0
+    z_rerank_weight: float = 0.0
+    z_nll_weight: float = 0.05
+    line_weight_power: float = 1.0
+    clean_z_only: bool = False
+    zwarn_weight: float = 0.3
+    high_z_boost: float = 1.0
+    high_z_threshold: float = math.log1p(1.0)
+
+
+def masked_huber(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, weight: torch.Tensor | None = None) -> torch.Tensor:
+    loss = F.smooth_l1_loss(pred, target, reduction="none", beta=0.5)
+    m = mask.float()
+    if weight is not None:
+        loss = loss * weight.float()
+    denom = m.sum().clamp_min(1.0)
+    return (loss * m).sum() / denom
+
+
+def redshift_losses(model, out: dict[str, torch.Tensor], y: torch.Tensor, zwarn: torch.Tensor, cfg: LossConfig) -> dict[str, torch.Tensor]:
+    clean = (~zwarn.bool()) & torch.isfinite(y) if cfg.clean_z_only else torch.isfinite(y)
+    y_pred_all = out.get("y_pred", out["y_mu"])
+    if clean.sum() == 0:
+        zero = y_pred_all.sum() * 0.0
+        return {"z_huber": zero, "z_bin": zero, "z_candidate": zero, "z_rerank": zero, "z_nll": zero}
+    y_mu = y_pred_all[clean]
+    y_true = y[clean]
+    y_logvar = out["y_logvar"][clean]
+    sample_weight = torch.where(zwarn[clean].bool(), torch.full_like(y_true, cfg.zwarn_weight), torch.ones_like(y_true))
+    if cfg.high_z_boost != 1.0:
+        high_z_weight = torch.where(
+            y_true >= cfg.high_z_threshold,
+            torch.full_like(y_true, cfg.high_z_boost),
+            torch.ones_like(y_true),
+        )
+        sample_weight = sample_weight * high_z_weight
+    huber = F.smooth_l1_loss(y_mu, y_true, beta=0.01, reduction="none")
+    z_huber = (huber * sample_weight).sum() / sample_weight.sum().clamp_min(1.0)
+    bins = model.y_to_bin(y_true)
+    z_bin_each = F.cross_entropy(out["z_bin_logits"][clean], bins, reduction="none")
+    z_bin = (z_bin_each * sample_weight).sum() / sample_weight.sum().clamp_min(1.0)
+    if "candidate_y" in out:
+        candidate_y = out["candidate_y"][clean]
+        true_candidate_y = candidate_y.gather(1, bins.unsqueeze(1)).squeeze(1)
+        z_candidate_each = F.smooth_l1_loss(true_candidate_y, y_true, beta=0.01, reduction="none")
+        z_candidate = (z_candidate_each * sample_weight).sum() / sample_weight.sum().clamp_min(1.0)
+    else:
+        z_candidate = y_mu.sum() * 0.0
+    if "rerank_logits" in out and "candidate_topk_y" in out:
+        candidate_topk_y = out["candidate_topk_y"][clean]
+        target_rank = torch.abs(candidate_topk_y - y_true.unsqueeze(1)).argmin(dim=1)
+        z_rerank_each = F.cross_entropy(out["rerank_logits"][clean], target_rank, reduction="none")
+        z_rerank = (z_rerank_each * sample_weight).sum() / sample_weight.sum().clamp_min(1.0)
+    else:
+        z_rerank = y_mu.sum() * 0.0
+    z_nll_each = 0.5 * (torch.exp(-y_logvar) * (y_mu - y_true).pow(2) + y_logvar)
+    z_nll = (z_nll_each * sample_weight).sum() / sample_weight.sum().clamp_min(1.0)
+    return {"z_huber": z_huber, "z_bin": z_bin, "z_candidate": z_candidate, "z_rerank": z_rerank, "z_nll": z_nll}
+
+
+def total_loss(model, out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], cfg: LossConfig) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    line_weight = batch["line_weight"].pow(cfg.line_weight_power)
+    rec = masked_huber(out["rec"], batch["target_flux"], batch["loss_mask"], weight=line_weight)
+    z_parts = redshift_losses(model, out, batch["y"], batch["zwarn"], cfg)
+    total = (
+        cfg.rec_weight * rec
+        + cfg.z_weight * z_parts["z_huber"]
+        + cfg.z_bin_weight * z_parts["z_bin"]
+        + cfg.z_candidate_weight * z_parts["z_candidate"]
+        + cfg.z_rerank_weight * z_parts["z_rerank"]
+        + cfg.z_nll_weight * z_parts["z_nll"]
+    )
+    metrics = {"loss": total.detach(), "rec": rec.detach(), **{k: v.detach() for k, v in z_parts.items()}}
+    return total, metrics
+
+
+def redshift_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]:
+    z_true = np.expm1(y_true)
+    z_pred = np.expm1(y_pred)
+    dz_norm = (z_pred - z_true) / (1.0 + z_true)
+    med = np.nanmedian(dz_norm)
+    nmad = 1.4826 * np.nanmedian(np.abs(dz_norm - med))
+    return {
+        "mae_z": float(np.nanmean(np.abs(z_pred - z_true))),
+        "mae_log1p": float(np.nanmean(np.abs(y_pred - y_true))),
+        "rmse_z": float(math.sqrt(np.nanmean((z_pred - z_true) ** 2))),
+        "nmad": float(nmad),
+        "cat_0p003": float(np.nanmean(np.abs(dz_norm) > 0.003)),
+        "cat_0p01": float(np.nanmean(np.abs(dz_norm) > 0.01)),
+        "cat_0p05": float(np.nanmean(np.abs(dz_norm) > 0.05)),
+        "pred_std_z": float(np.nanstd(z_pred)),
+        "true_std_z": float(np.nanstd(z_true)),
+    }