code/metrics.py

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_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_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
    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_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_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)),
    }