File size: 4,366 Bytes
48532fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | 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)),
}
|