from __future__ import annotations import argparse import hashlib import json import math import os import time from dataclasses import dataclass from pathlib import Path from typing import Any import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch import torch.nn.functional as F from torch import nn from torch.optim import AdamW from torch.utils.data import DataLoader, WeightedRandomSampler from tqdm import tqdm from .data import SpectraListDataset, collect_mmu_desi, compute_sample_stats, split_indices, valid_pixel_mask from .metrics import LossConfig, masked_huber, redshift_losses, redshift_metrics from .model import fourier_loglam from .plots import plot_reconstruction_batch, plot_redshift_scatter @dataclass class RawCollatorConfig: target_length: int = 4096 min_scale: float = 1e-3 random_mask_ratio: float = 0.0 eval_mask_ratio: float = 0.25 mask_mode: str = "pixel" mask_span_min: int = 16 mask_span_max: int = 64 line_region_percentile: float = 90.0 augment_ood: bool = False crop_prob: float = 0.0 bad_window_prob: float = 0.0 throughput_prob: float = 0.0 noise_prob: float = 0.0 resolution_prob: float = 0.0 downsample_prob: float = 0.0 line_dropout_prob: float = 0.0 span_dropout_prob: float = 0.0 grid_jitter_prob: float = 0.0 grid_shift_frac: float = 0.0 grid_scale_frac: float = 0.0 grid_jitter_warmup_steps: int = 0 redshift_shift: float = 0.0 class RawSpectraCollator: def __init__(self, cfg: RawCollatorConfig, train: bool = True, seed: int = 17): self.cfg = cfg self.train = train self.seed = seed self.rng = np.random.default_rng(seed) self.batch_count = 0 self.jitter_scale = 1.0 def __call__(self, samples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: if self.train: self.batch_count += 1 warmup = max(0, int(self.cfg.grid_jitter_warmup_steps)) self.jitter_scale = min(1.0, self.batch_count / float(warmup)) if warmup > 0 else 1.0 items = [self._prepare_sample(s) for s in samples] x = np.stack([item["x"] for item in items], axis=0).astype(np.float32) valid = np.stack([item["valid"] for item in items], axis=0).astype(np.bool_) loglam = np.stack([item["loglam"] for item in items], axis=0).astype(np.float32) target_flux = np.stack([item["target_flux"] for item in items], axis=0).astype(np.float32) loss_mask = np.stack([item["loss_mask"] for item in items], axis=0).astype(np.bool_) line_weight = np.stack([item["line_weight"] for item in items], axis=0).astype(np.float32) line_region = np.stack([item["line_region"] for item in items], axis=0).astype(np.bool_) z = np.asarray([item["z"] for item in items], dtype=np.float32) y = np.asarray([item["y"] for item in items], dtype=np.float32) zwarn = np.asarray([item["zwarn"] for item in items], dtype=np.bool_) return { "x": torch.from_numpy(x), "valid": torch.from_numpy(valid), "loglam": torch.from_numpy(loglam), "target_flux": torch.from_numpy(target_flux), "loss_mask": torch.from_numpy(loss_mask), "line_weight": torch.from_numpy(line_weight), "line_region": torch.from_numpy(line_region), "z": torch.from_numpy(z), "y": torch.from_numpy(y), "zwarn": torch.from_numpy(zwarn), } def _prepare_sample(self, sample: dict[str, Any]) -> dict[str, Any]: rng = self.rng if self.train else self._eval_rng(sample) flux = np.asarray(sample["flux"], dtype=np.float32).copy() ivar = np.asarray(sample["ivar"], dtype=np.float32).copy() lam = np.asarray(sample["lambda"], dtype=np.float32) lsf = np.asarray(sample["lsf_sigma"], dtype=np.float32) bad = np.asarray(sample["bad_mask"], dtype=np.bool_).copy() if self.cfg.augment_ood: bad = self._augment_bad_windows(bad, rng) flux = self._augment_flux_calibration(flux, lam, rng) flux = self._augment_resolution(flux, rng) flux, ivar = self._augment_downsample_resample(flux, ivar, lam, rng) flux = self._augment_noise(flux, ivar, rng) valid = np.isfinite(flux) & np.isfinite(ivar) & np.isfinite(lam) & (ivar > 0) & (~bad) loglam = np.log(np.clip(lam.astype(np.float64), 1.0, None)).astype(np.float32) if valid.sum() < 16: valid = valid_pixel_mask(sample) grid_lo = float(np.nanmin(loglam)) grid_hi = float(np.nanmax(loglam)) if ( self.train and self.cfg.grid_jitter_prob > 0 and rng.random() < self.cfg.grid_jitter_prob * self.jitter_scale and math.isfinite(grid_lo) and math.isfinite(grid_hi) and grid_hi > grid_lo ): span = grid_hi - grid_lo shift = float(rng.normal(0.0, max(0.0, self.cfg.grid_shift_frac) * self.jitter_scale)) * span scale_max = max(0.0, self.cfg.grid_scale_frac) * self.jitter_scale scale_delta = float(rng.uniform(-scale_max, scale_max)) scaled_span = span * max(0.50, 1.0 + scale_delta) center = 0.5 * (grid_lo + grid_hi) + shift grid_lo = center - 0.5 * scaled_span grid_hi = center + 0.5 * scaled_span grid = np.linspace(grid_lo, grid_hi, self.cfg.target_length, dtype=np.float32) flux_grid = self._interp_valid(loglam, flux, valid, grid, fill=0.0) ivar_grid = self._interp_valid(loglam, ivar, valid, grid, fill=0.0) lsf_grid = self._interp_valid(loglam, lsf, valid, grid, fill=0.0) valid_grid = np.interp(grid, loglam, valid.astype(np.float32), left=0.0, right=0.0) > 0.5 center = float(np.nanmedian(flux_grid[valid_grid])) if valid_grid.any() else 0.0 dev = np.abs(flux_grid[valid_grid] - center) if valid_grid.any() else np.asarray([1.0], dtype=np.float32) scale = float(np.nanmedian(dev) * 1.4826) if not math.isfinite(scale) or scale < self.cfg.min_scale: scale = max(float(np.nanmedian(np.abs(flux_grid[valid_grid]))) if valid_grid.any() else 1.0, self.cfg.min_scale) norm_flux = np.arcsinh((flux_grid - center) / scale).astype(np.float32) norm_ivar = np.log1p(np.maximum(ivar_grid * scale * scale, 0.0)).astype(np.float32) norm_ivar = np.clip(norm_ivar / 8.0, 0.0, 4.0) lsf_norm = np.nan_to_num(lsf_grid / 3.0, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32) loglam_norm = ((grid - math.log(6000.0)) / 0.45).astype(np.float32) grad = np.gradient(norm_flux, grid).astype(np.float32) good_grad = np.abs(grad[valid_grid]) grad_scale = float(np.percentile(good_grad, 95)) if len(good_grad) else 1.0 if not math.isfinite(grad_scale) or grad_scale <= 0: grad_scale = 1.0 grad = np.clip(grad / grad_scale, -5.0, 5.0).astype(np.float32) abs_grad = np.abs(grad).astype(np.float32) target_flux = norm_flux.copy() line_weight = self._line_weights(abs_grad, valid_grid) line_region = self._line_region(abs_grad, valid_grid) corrupt = self._sample_input_dropout(abs_grad, valid_grid, rng) if corrupt.any(): norm_flux = norm_flux.copy() grad = grad.copy() abs_grad = abs_grad.copy() norm_flux[corrupt] = 0.0 grad[corrupt] = 0.0 abs_grad[corrupt] = 0.0 y = math.log1p(float(sample["z"])) if self.train and self.cfg.redshift_shift > 0: delta = float(self.rng.uniform(-self.cfg.redshift_shift, self.cfg.redshift_shift)) y = max(0.0, y + delta) grid = (grid + delta).astype(np.float32) loglam_norm = ((grid - math.log(6000.0)) / 0.45).astype(np.float32) x = np.stack( [ norm_flux, norm_ivar, valid_grid.astype(np.float32), lsf_norm, loglam_norm, grad, abs_grad, corrupt.astype(np.float32), ], axis=0, ) return { "x": x, "valid": valid_grid, "loglam": grid, "target_flux": target_flux, "loss_mask": corrupt & valid_grid, "line_weight": line_weight, "line_region": line_region, "z": sample["z"], "y": np.float32(y), "zwarn": sample["zwarn"], } def _eval_rng(self, sample: dict[str, Any]) -> np.random.Generator: object_id = str(sample.get("object_id", "")) lam = np.asarray(sample["lambda"], dtype=np.float32) key = f"{self.seed}|{object_id}|{float(sample['z']):.8g}|{len(lam)}|{float(lam[0]):.4f}|{float(lam[-1]):.4f}" digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest() return np.random.default_rng(int.from_bytes(digest, "little", signed=False)) def _interp_valid(self, x: np.ndarray, y: np.ndarray, valid: np.ndarray, x_new: np.ndarray, fill: float) -> np.ndarray: good = valid & np.isfinite(x) & np.isfinite(y) if good.sum() < 2: return np.full_like(x_new, fill, dtype=np.float32) return np.interp(x_new, x[good], y[good], left=fill, right=fill).astype(np.float32) def _augment_bad_windows(self, bad: np.ndarray, rng: np.random.Generator) -> np.ndarray: out = bad.copy() n = len(out) if rng.random() < self.cfg.crop_prob: frac = float(rng.uniform(0.62, 0.96)) width = max(32, int(n * frac)) start = int(rng.integers(0, max(1, n - width))) keep = np.zeros(n, dtype=np.bool_) keep[start : start + width] = True out |= ~keep if rng.random() < self.cfg.bad_window_prob: for _ in range(int(rng.integers(1, 5))): width = int(rng.integers(max(8, n // 240), max(12, n // 45))) start = int(rng.integers(0, max(1, n - width))) out[start : start + width] = True return out def _augment_flux_calibration(self, flux: np.ndarray, lam: np.ndarray, rng: np.random.Generator) -> np.ndarray: if rng.random() >= self.cfg.throughput_prob: return flux x = np.linspace(-1.0, 1.0, len(flux), dtype=np.float32) coeff = rng.normal(0.0, [0.05, 0.025, 0.015]).astype(np.float32) curve = 1.0 + coeff[0] * x + coeff[1] * (x * x - 0.33) + coeff[2] * np.sin(np.pi * x) return (flux * np.clip(curve, 0.65, 1.35)).astype(np.float32) def _augment_noise(self, flux: np.ndarray, ivar: np.ndarray, rng: np.random.Generator) -> np.ndarray: if rng.random() >= self.cfg.noise_prob: return flux sigma = np.zeros_like(flux, dtype=np.float32) good = np.isfinite(ivar) & (ivar > 0) sigma[good] = 1.0 / np.sqrt(np.maximum(ivar[good], 1e-8)) scale = float(rng.uniform(0.15, 0.75)) return (flux + rng.normal(0.0, sigma * scale).astype(np.float32)).astype(np.float32) def _augment_resolution(self, flux: np.ndarray, rng: np.random.Generator) -> np.ndarray: if rng.random() >= self.cfg.resolution_prob: return flux finite = np.isfinite(flux) fill = float(np.nanmedian(flux[finite])) if finite.any() else 0.0 base = np.nan_to_num(flux, nan=fill, posinf=fill, neginf=fill).astype(np.float32) sigma = float(rng.uniform(0.6, 3.0)) radius = max(2, int(math.ceil(4.0 * sigma))) x = np.arange(-radius, radius + 1, dtype=np.float32) kernel = np.exp(-0.5 * (x / sigma) ** 2) kernel = (kernel / kernel.sum()).astype(np.float32) padded = np.pad(base, (radius, radius), mode="edge") return np.convolve(padded, kernel, mode="valid").astype(np.float32) def _augment_downsample_resample( self, flux: np.ndarray, ivar: np.ndarray, lam: np.ndarray, rng: np.random.Generator, ) -> tuple[np.ndarray, np.ndarray]: if rng.random() >= self.cfg.downsample_prob: return flux, ivar n = len(flux) if n < 32: return flux, ivar factor = int(rng.choice(np.asarray([2, 3, 4, 6, 8], dtype=np.int64))) offset = int(rng.integers(0, factor)) idx = np.arange(offset, n, factor, dtype=np.int64) if len(idx) < 4: return flux, ivar lam_good = np.asarray(lam[idx], dtype=np.float32) flux_good = np.asarray(flux[idx], dtype=np.float32) ivar_good = np.asarray(ivar[idx], dtype=np.float32) good = np.isfinite(lam_good) & np.isfinite(flux_good) & np.isfinite(ivar_good) if np.count_nonzero(good) < 4: return flux, ivar lam_good = lam_good[good] order = np.argsort(lam_good) lam_good = lam_good[order] flux_good = flux_good[good][order] ivar_good = ivar_good[good][order] flux_out = np.interp(lam, lam_good, flux_good, left=flux_good[0], right=flux_good[-1]).astype(np.float32) ivar_out = np.interp(lam, lam_good, ivar_good, left=0.0, right=0.0).astype(np.float32) ivar_out *= float(rng.uniform(0.25, 0.85)) return flux_out, ivar_out def _sample_input_dropout(self, abs_grad: np.ndarray, valid: np.ndarray, rng: np.random.Generator) -> np.ndarray: corrupt = np.zeros_like(valid, dtype=np.bool_) if valid.sum() < 16: return corrupt n = len(valid) valid_idx = np.where(valid)[0] ratio = self.cfg.random_mask_ratio if self.train else self.cfg.eval_mask_ratio if ratio > 0: n_rand = max(1, int(round(len(valid_idx) * min(float(ratio), 1.0)))) if self.cfg.mask_mode == "pixel": corrupt[rng.choice(valid_idx, size=min(n_rand, len(valid_idx)), replace=False)] = True else: line_bias = self.cfg.mask_mode in {"line_span", "mixed_span"} self._add_spans_to_mask(corrupt, valid, abs_grad, n_rand, rng, line_bias=line_bias) if self.train and rng.random() < self.cfg.span_dropout_prob: for _ in range(int(rng.integers(1, 4))): width = int(rng.integers(max(4, n // 220), max(8, n // 55))) start = int(rng.integers(0, max(1, n - width))) corrupt[start : start + width] |= valid[start : start + width] if self.train and rng.random() < self.cfg.line_dropout_prob: score = abs_grad.copy() score[~valid] = 0.0 if np.count_nonzero(score) > 0: k = max(4, n // 96) peaks = np.argsort(score)[-k:] for j in peaks: width = int(rng.integers(max(2, n // 900), max(4, n // 280))) lo = max(0, int(j) - width) hi = min(n, int(j) + width + 1) corrupt[lo:hi] |= valid[lo:hi] return corrupt & valid def _add_spans_to_mask( self, corrupt: np.ndarray, valid: np.ndarray, abs_grad: np.ndarray, target_count: int, rng: np.random.Generator, *, line_bias: bool, ) -> None: valid_idx = np.where(valid)[0] if len(valid_idx) == 0: return lo_w = max(1, int(self.cfg.mask_span_min)) hi_w = max(lo_w + 1, int(self.cfg.mask_span_max) + 1) probs = None if line_bias: score = abs_grad[valid_idx].astype(np.float64) positive = score[np.isfinite(score) & (score > 0)] scale = float(np.percentile(positive, 90)) if len(positive) else 1.0 if not math.isfinite(scale) or scale <= 0: scale = 1.0 score = np.clip(score / scale, 0.0, 5.0) + 0.05 probs = score / score.sum() max_tries = max(32, target_count * 4) tries = 0 while int(np.count_nonzero(corrupt & valid)) < target_count and tries < max_tries: tries += 1 center = int(rng.choice(valid_idx, p=probs)) width = int(rng.integers(lo_w, hi_w)) lo = max(0, center - width // 2) hi = min(len(valid), lo + width) corrupt[lo:hi] |= valid[lo:hi] def _line_weights(self, abs_grad: np.ndarray, valid: np.ndarray) -> np.ndarray: weight = np.ones_like(abs_grad, dtype=np.float32) if valid.sum() < 16: return weight scale = float(np.percentile(abs_grad[valid], 90)) if math.isfinite(scale) and scale > 0: weight += 2.0 * np.clip(abs_grad / scale, 0.0, 2.0) weight[~valid] = 1.0 return np.clip(weight, 1.0, 5.0).astype(np.float32) def _line_region(self, abs_grad: np.ndarray, valid: np.ndarray) -> np.ndarray: region = np.zeros_like(valid, dtype=np.bool_) if valid.sum() < 16: return region pct = min(max(float(self.cfg.line_region_percentile), 0.0), 100.0) thresh = float(np.percentile(abs_grad[valid], pct)) if math.isfinite(thresh) and thresh > 0: region = (abs_grad >= thresh) & valid return region.astype(np.bool_) class ConvBlock(nn.Module): def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 7, stride: int = 1, dropout: float = 0.0): super().__init__() padding = kernel_size // 2 self.net = nn.Sequential( nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False), nn.BatchNorm1d(out_channels), nn.GELU(), nn.Dropout(dropout), nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False), nn.BatchNorm1d(out_channels), ) self.skip = ( nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False) if stride != 1 or in_channels != out_channels else nn.Identity() ) self.act = nn.GELU() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.act(self.net(x) + self.skip(x)) class LayerScaleEncoderLayer(nn.Module): def __init__(self, d_model: int, heads: int, dropout: float, layerscale_init: float): super().__init__() self.norm1 = nn.LayerNorm(d_model) self.self_attn = nn.MultiheadAttention(d_model, heads, dropout=dropout, batch_first=True) self.dropout1 = nn.Dropout(dropout) self.norm2 = nn.LayerNorm(d_model) self.linear1 = nn.Linear(d_model, d_model * 4) self.linear2 = nn.Linear(d_model * 4, d_model) self.dropout = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.act = nn.GELU() init = float(layerscale_init) self.ls1 = nn.Parameter(torch.full((d_model,), init)) self.ls2 = nn.Parameter(torch.full((d_model,), init)) def forward( self, src: torch.Tensor, src_mask: torch.Tensor | None = None, src_key_padding_mask: torch.Tensor | None = None, is_causal: bool = False, ) -> torch.Tensor: q = self.norm1(src) attn, _ = self.self_attn( q, q, q, attn_mask=src_mask, key_padding_mask=src_key_padding_mask, need_weights=False, is_causal=is_causal, ) src = src + self.ls1 * self.dropout1(attn) ff = self.linear2(self.dropout(self.act(self.linear1(self.norm2(src))))) return src + self.ls2 * self.dropout2(ff) class HybridSpecZ(nn.Module): def __init__( self, in_channels: int = 8, d_model: int = 256, conv_width: int = 128, layers: int = 5, heads: int = 8, dropout: float = 0.1, fourier_freqs: int = 32, z_bins: int = 64, y_min: float = 0.0, y_max: float = math.log1p(6.0), prediction_mode: str = "regression", bin_temperature: float = 1.0, residual_scale: float = 0.06, candidate_topk: int = 5, stem_stride: int = 8, rec_hidden_mult: int = 0, rec_refine_width: int = 16, rec_refine_kernel: int = 5, layerscale_init: float = 0.0, ): super().__init__() allowed_modes = { "regression", "softbin", "hybrid", "bin_residual", "ranked_bin_residual", "candidate_rerank", "calibrated_bin_residual", } if prediction_mode not in allowed_modes: raise ValueError(f"prediction_mode must be one of {sorted(allowed_modes)}, got {prediction_mode!r}") self.fourier_freqs = fourier_freqs self.z_bins = z_bins self.y_min = y_min self.y_max = y_max self.prediction_mode = prediction_mode self.bin_temperature = bin_temperature self.residual_scale = residual_scale self.candidate_topk = max(1, min(int(candidate_topk), z_bins)) if stem_stride not in {4, 8}: raise ValueError(f"stem_stride must be 4 or 8, got {stem_stride}") self.stem_stride = int(stem_stride) self.rec_pixels_per_token = int(stem_stride) self.stride_stages = int(round(math.log2(self.stem_stride))) bin_width = (y_max - y_min) / z_bins centers = torch.linspace(y_min + 0.5 * bin_width, y_max - 0.5 * bin_width, z_bins) self.register_buffer("z_bin_centers", centers, persistent=False) if self.stem_stride == 8: self.stem = nn.Sequential( ConvBlock(in_channels, conv_width, kernel_size=9, stride=2, dropout=dropout * 0.5), ConvBlock(conv_width, conv_width, kernel_size=7, stride=2, dropout=dropout * 0.5), ConvBlock(conv_width, d_model, kernel_size=7, stride=2, dropout=dropout * 0.5), ConvBlock(d_model, d_model, kernel_size=5, stride=1, dropout=dropout * 0.5), ) else: self.stem = nn.Sequential( ConvBlock(in_channels, conv_width, kernel_size=9, stride=2, dropout=dropout * 0.5), ConvBlock(conv_width, d_model, kernel_size=7, stride=2, dropout=dropout * 0.5), ConvBlock(d_model, d_model, kernel_size=5, stride=1, dropout=dropout * 0.5), ) self.pos_proj = nn.Sequential(nn.Linear(fourier_freqs * 2, d_model), nn.LayerNorm(d_model)) self.cls = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) # The model never receives true z; this learned query is the always-masked z token. self.z_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) if layerscale_init > 0: enc_layer = LayerScaleEncoderLayer(d_model, heads, dropout, layerscale_init) else: enc_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=heads, dim_feedforward=d_model * 4, dropout=dropout, batch_first=True, norm_first=True, activation="gelu", ) self.encoder = nn.TransformerEncoder(enc_layer, num_layers=layers) self.pool_gate = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, 1)) head_dim = d_model * 5 self.z_head = nn.Sequential(nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, 2)) self.z_bin_head = nn.Sequential(nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, z_bins)) self.z_candidate_head = nn.Sequential( nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, z_bins), ) self.z_rerank_head = nn.Sequential( nn.LayerNorm(head_dim + 3), nn.Linear(head_dim + 3, max(64, d_model // 2)), nn.GELU(), nn.Dropout(dropout), nn.Linear(max(64, d_model // 2), 1), ) self.z_calib_head = nn.Sequential( nn.LayerNorm(head_dim + 3), nn.Linear(head_dim + 3, max(64, d_model // 2)), nn.GELU(), nn.Dropout(dropout), nn.Linear(max(64, d_model // 2), 1), ) nn.init.zeros_(self.z_calib_head[-1].weight) nn.init.zeros_(self.z_calib_head[-1].bias) if rec_hidden_mult > 0: rec_hidden = int(d_model * rec_hidden_mult) self.rec_head = nn.Sequential( nn.LayerNorm(d_model), nn.Linear(d_model, rec_hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(rec_hidden, self.rec_pixels_per_token), ) else: self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, self.rec_pixels_per_token)) rec_pad = int(rec_refine_kernel) // 2 self.rec_refine = nn.Sequential( nn.Conv1d(1, rec_refine_width, kernel_size=rec_refine_kernel, padding=rec_pad), nn.GELU(), nn.Conv1d(rec_refine_width, 1, kernel_size=rec_refine_kernel, padding=rec_pad), ) def forward(self, x: torch.Tensor, valid: torch.Tensor, loglam: torch.Tensor) -> dict[str, torch.Tensor]: bsz = x.shape[0] h = self.stem(x).transpose(1, 2) tok_valid = valid.float().unsqueeze(1) tok_loglam = loglam.unsqueeze(1) for _ in range(self.stride_stages): tok_valid = F.avg_pool1d(tok_valid, kernel_size=2, stride=2, ceil_mode=True) tok_loglam = F.avg_pool1d(tok_loglam, kernel_size=2, stride=2, ceil_mode=True) tok_valid = tok_valid.squeeze(1) > 0.20 tok_loglam = tok_loglam.squeeze(1) if tok_valid.shape[1] != h.shape[1]: tok_valid = tok_valid[:, : h.shape[1]] tok_loglam = tok_loglam[:, : h.shape[1]] h = h[:, : tok_valid.shape[1]] h = h + self.pos_proj(fourier_loglam(tok_loglam, self.fourier_freqs)) cls = self.cls.expand(bsz, -1, -1) z_query = self.z_query.expand(bsz, -1, -1) src = torch.cat([cls, z_query, h], dim=1) special_valid = torch.ones((bsz, 2), dtype=torch.bool, device=x.device) src_valid = torch.cat([special_valid, tok_valid], dim=1) padding = ~src_valid memory = self.encoder(src, src_key_padding_mask=padding) spec = memory[:, 2:] spec_valid = src_valid[:, 2:] spec_mask = spec_valid.unsqueeze(-1) rec = self.rec_head(spec).reshape(bsz, -1) rec = rec + self.rec_refine(rec.unsqueeze(1)).squeeze(1) if rec.shape[1] > x.shape[-1]: rec = rec[:, : x.shape[-1]] elif rec.shape[1] < x.shape[-1]: rec = F.pad(rec, (0, x.shape[-1] - rec.shape[1])) denom = spec_valid.float().sum(dim=1).clamp_min(1.0).unsqueeze(-1) mean_pool = (spec * spec_mask.float()).sum(dim=1) / denom max_pool = spec.masked_fill(~spec_mask, -1e4).max(dim=1).values gate_logits = self.pool_gate(spec).squeeze(-1).masked_fill(~spec_valid, -1e4) gate = torch.softmax(gate_logits, dim=1) attn_pool = torch.einsum("bn,bnd->bd", gate, spec) feat = torch.cat([memory[:, 0], memory[:, 1], mean_pool, max_pool, attn_pool], dim=-1) z_params = self.z_head(feat) z_bin_logits = self.z_bin_head(feat) candidate_residual = self.residual_scale * torch.tanh(self.z_candidate_head(feat)) centers = self.z_bin_centers.to(dtype=z_bin_logits.dtype, device=z_bin_logits.device) candidate_y = (centers.unsqueeze(0) + candidate_residual).clamp(self.y_min, self.y_max) topk_logits, topk_bins = torch.topk(z_bin_logits, k=self.candidate_topk, dim=-1) candidate_topk_y = candidate_y.gather(1, topk_bins) rank = torch.linspace(0.0, 1.0, self.candidate_topk, device=x.device, dtype=feat.dtype).view(1, self.candidate_topk, 1) rerank_feat = feat.unsqueeze(1).expand(-1, self.candidate_topk, -1) rerank_in = torch.cat( [ rerank_feat, candidate_topk_y.to(dtype=feat.dtype).unsqueeze(-1), topk_logits.to(dtype=feat.dtype).unsqueeze(-1), rank.expand(bsz, -1, -1), ], dim=-1, ) rerank_logits = self.z_rerank_head(rerank_in).squeeze(-1) rerank_idx = rerank_logits.argmax(dim=-1, keepdim=True) y_reranked = candidate_topk_y.gather(1, rerank_idx).squeeze(1) y_reg = z_params[:, 0] bin_prob = torch.softmax(z_bin_logits / max(self.bin_temperature, 1e-4), dim=-1) y_bin = (bin_prob * self.z_bin_centers.to(dtype=bin_prob.dtype, device=bin_prob.device)).sum(dim=-1) y_ranked = (bin_prob * candidate_y.to(dtype=bin_prob.dtype)).sum(dim=-1) y_legacy_bin_residual = y_bin + self.residual_scale * torch.tanh(y_reg) calib_in = torch.cat( [ feat, y_legacy_bin_residual.to(dtype=feat.dtype).unsqueeze(-1), y_ranked.to(dtype=feat.dtype).unsqueeze(-1), candidate_topk_y[:, 0].to(dtype=feat.dtype).unsqueeze(-1), ], dim=-1, ) y_calibrated = y_legacy_bin_residual + self.residual_scale * torch.tanh(self.z_calib_head(calib_in).squeeze(-1)) if self.prediction_mode == "regression": y_pred = y_reg elif self.prediction_mode == "softbin": y_pred = y_bin elif self.prediction_mode == "hybrid": y_pred = 0.35 * y_reg + 0.65 * y_bin elif self.prediction_mode == "ranked_bin_residual": y_pred = 0.5 * y_legacy_bin_residual + 0.5 * y_ranked elif self.prediction_mode == "candidate_rerank": y_pred = y_reranked elif self.prediction_mode == "calibrated_bin_residual": y_pred = y_calibrated else: y_pred = y_legacy_bin_residual y_pred = y_pred.clamp(self.y_min, self.y_max) return { "rec": rec, "y_mu": y_pred, "y_pred": y_pred, "y_reg": y_reg, "y_bin": y_bin, "y_ranked": y_ranked, "y_top1_candidate": candidate_topk_y[:, 0], "y_reranked": y_reranked, "y_calibrated": y_calibrated, "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0), "z_bin_logits": z_bin_logits, "z_feat": feat, "candidate_y": candidate_y, "candidate_topk_y": candidate_topk_y, "candidate_topk_bins": topk_bins, "candidate_topk_logits": topk_logits, "rerank_logits": rerank_logits, } def y_to_bin(self, y: torch.Tensor) -> torch.Tensor: scaled = (y - self.y_min) / max(self.y_max - self.y_min, 1e-6) return torch.clamp((scaled * self.z_bins).long(), 0, self.z_bins - 1) class WavelengthTokenSpecZ(HybridSpecZ): """Transformer encoder over wavelength-conditioned tokens instead of a conv pixel stem.""" def __init__( self, in_channels: int = 8, d_model: int = 256, conv_width: int = 128, layers: int = 5, heads: int = 8, dropout: float = 0.1, fourier_freqs: int = 32, z_bins: int = 64, y_min: float = 0.0, y_max: float = math.log1p(6.0), prediction_mode: str = "regression", bin_temperature: float = 1.0, residual_scale: float = 0.06, candidate_topk: int = 5, token_stride: int = 8, rec_hidden_mult: int = 0, rec_refine_width: int = 16, rec_refine_kernel: int = 5, layerscale_init: float = 0.0, ): super().__init__( in_channels=in_channels, d_model=d_model, conv_width=conv_width, layers=layers, heads=heads, dropout=dropout, fourier_freqs=fourier_freqs, z_bins=z_bins, y_min=y_min, y_max=y_max, prediction_mode=prediction_mode, bin_temperature=bin_temperature, residual_scale=residual_scale, candidate_topk=candidate_topk, stem_stride=8, rec_hidden_mult=rec_hidden_mult, rec_refine_width=rec_refine_width, rec_refine_kernel=rec_refine_kernel, layerscale_init=layerscale_init, ) self.token_stride = max(1, int(token_stride)) self.rec_pixels_per_token = self.token_stride self.stem = nn.Identity() self.input_proj = nn.Sequential( nn.Linear(in_channels + fourier_freqs * 2, d_model), nn.LayerNorm(d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, d_model), ) if rec_hidden_mult > 0: rec_hidden = int(d_model * rec_hidden_mult) self.rec_head = nn.Sequential( nn.LayerNorm(d_model), nn.Linear(d_model, rec_hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(rec_hidden, self.rec_pixels_per_token), ) else: self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, self.rec_pixels_per_token)) def _pool_wavelength_tokens( self, x: torch.Tensor, valid: torch.Tensor, loglam: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, channels, length = x.shape stride = self.token_stride pad = (-length) % stride if pad: x = F.pad(x, (0, pad)) valid = F.pad(valid.float(), (0, pad)).bool() loglam = torch.cat([loglam, loglam[:, -1:].expand(-1, pad)], dim=1) token_count = x.shape[-1] // stride x_group = x.reshape(bsz, channels, token_count, stride) valid_group = valid.reshape(bsz, 1, token_count, stride).float() loglam_group = loglam.reshape(bsz, 1, token_count, stride) counts = valid_group.sum(dim=-1) denom = counts.clamp_min(1.0) token_x = (x_group * valid_group).sum(dim=-1) / denom token_loglam = (loglam_group * valid_group).sum(dim=-1) / denom fallback_loglam = loglam_group.mean(dim=-1) token_loglam = torch.where(counts > 0, token_loglam, fallback_loglam).squeeze(1) token_valid = counts.squeeze(1) > 0 return token_x.transpose(1, 2), token_valid, token_loglam def forward(self, x: torch.Tensor, valid: torch.Tensor, loglam: torch.Tensor) -> dict[str, torch.Tensor]: bsz = x.shape[0] token_x, tok_valid, tok_loglam = self._pool_wavelength_tokens(x, valid, loglam) h = self.input_proj(torch.cat([token_x, fourier_loglam(tok_loglam, self.fourier_freqs)], dim=-1)) cls = self.cls.expand(bsz, -1, -1) z_query = self.z_query.expand(bsz, -1, -1) src = torch.cat([cls, z_query, h], dim=1) special_valid = torch.ones((bsz, 2), dtype=torch.bool, device=x.device) src_valid = torch.cat([special_valid, tok_valid], dim=1) padding = ~src_valid memory = self.encoder(src, src_key_padding_mask=padding) spec = memory[:, 2:] spec_valid = src_valid[:, 2:] spec_mask = spec_valid.unsqueeze(-1) rec = self.rec_head(spec).reshape(bsz, -1) rec = rec + self.rec_refine(rec.unsqueeze(1)).squeeze(1) if rec.shape[1] > x.shape[-1]: rec = rec[:, : x.shape[-1]] elif rec.shape[1] < x.shape[-1]: rec = F.pad(rec, (0, x.shape[-1] - rec.shape[1])) denom = spec_valid.float().sum(dim=1).clamp_min(1.0).unsqueeze(-1) mean_pool = (spec * spec_mask.float()).sum(dim=1) / denom max_pool = spec.masked_fill(~spec_mask, -1e4).max(dim=1).values gate_logits = self.pool_gate(spec).squeeze(-1).masked_fill(~spec_valid, -1e4) gate = torch.softmax(gate_logits, dim=1) attn_pool = torch.einsum("bn,bnd->bd", gate, spec) feat = torch.cat([memory[:, 0], memory[:, 1], mean_pool, max_pool, attn_pool], dim=-1) z_params = self.z_head(feat) z_bin_logits = self.z_bin_head(feat) candidate_residual = self.residual_scale * torch.tanh(self.z_candidate_head(feat)) centers = self.z_bin_centers.to(dtype=z_bin_logits.dtype, device=z_bin_logits.device) candidate_y = (centers.unsqueeze(0) + candidate_residual).clamp(self.y_min, self.y_max) topk_logits, topk_bins = torch.topk(z_bin_logits, k=self.candidate_topk, dim=-1) candidate_topk_y = candidate_y.gather(1, topk_bins) rank = torch.linspace(0.0, 1.0, self.candidate_topk, device=x.device, dtype=feat.dtype).view(1, self.candidate_topk, 1) rerank_feat = feat.unsqueeze(1).expand(-1, self.candidate_topk, -1) rerank_in = torch.cat( [ rerank_feat, candidate_topk_y.to(dtype=feat.dtype).unsqueeze(-1), topk_logits.to(dtype=feat.dtype).unsqueeze(-1), rank.expand(bsz, -1, -1), ], dim=-1, ) rerank_logits = self.z_rerank_head(rerank_in).squeeze(-1) rerank_idx = rerank_logits.argmax(dim=-1, keepdim=True) y_reranked = candidate_topk_y.gather(1, rerank_idx).squeeze(1) y_reg = z_params[:, 0] bin_prob = torch.softmax(z_bin_logits / max(self.bin_temperature, 1e-4), dim=-1) y_bin = (bin_prob * self.z_bin_centers.to(dtype=bin_prob.dtype, device=bin_prob.device)).sum(dim=-1) y_ranked = (bin_prob * candidate_y.to(dtype=bin_prob.dtype)).sum(dim=-1) y_legacy_bin_residual = y_bin + self.residual_scale * torch.tanh(y_reg) calib_in = torch.cat( [ feat, y_legacy_bin_residual.to(dtype=feat.dtype).unsqueeze(-1), y_ranked.to(dtype=feat.dtype).unsqueeze(-1), candidate_topk_y[:, 0].to(dtype=feat.dtype).unsqueeze(-1), ], dim=-1, ) y_calibrated = y_legacy_bin_residual + self.residual_scale * torch.tanh(self.z_calib_head(calib_in).squeeze(-1)) if self.prediction_mode == "regression": y_pred = y_reg elif self.prediction_mode == "softbin": y_pred = y_bin elif self.prediction_mode == "hybrid": y_pred = 0.35 * y_reg + 0.65 * y_bin elif self.prediction_mode == "ranked_bin_residual": y_pred = 0.5 * y_legacy_bin_residual + 0.5 * y_ranked elif self.prediction_mode == "candidate_rerank": y_pred = y_reranked elif self.prediction_mode == "calibrated_bin_residual": y_pred = y_calibrated else: y_pred = y_legacy_bin_residual y_pred = y_pred.clamp(self.y_min, self.y_max) return { "rec": rec, "y_mu": y_pred, "y_pred": y_pred, "y_reg": y_reg, "y_bin": y_bin, "y_ranked": y_ranked, "y_top1_candidate": candidate_topk_y[:, 0], "y_reranked": y_reranked, "y_calibrated": y_calibrated, "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0), "z_bin_logits": z_bin_logits, "z_feat": feat, "candidate_y": candidate_y, "candidate_topk_y": candidate_topk_y, "candidate_topk_bins": topk_bins, "candidate_topk_logits": topk_logits, "rerank_logits": rerank_logits, } def move_to_device(batch: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]: return {k: v.to(device, non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()} def limit_batch_examples(batch: dict[str, torch.Tensor], max_examples: int | None, seen_examples: int) -> dict[str, torch.Tensor] | None: if max_examples is None or max_examples <= 0: return batch remaining = int(max_examples) - int(seen_examples) if remaining <= 0: return None bsz = int(batch["y"].shape[0]) if remaining >= bsz: return batch return {k: v[:remaining] if torch.is_tensor(v) and v.shape[:1] == (bsz,) else v for k, v in batch.items()} def load_checkpoint_into_model(model: nn.Module, state: dict[str, torch.Tensor], allow_mismatched: bool = False) -> None: if not allow_mismatched: try: model.load_state_dict(state, strict=True) except RuntimeError: missing, unexpected = model.load_state_dict(state, strict=False) print(f"RESUME_NONSTRICT missing={list(missing)} unexpected={list(unexpected)}") return target_state = model.state_dict() compatible = {} skipped = [] for key, value in state.items(): target = target_state.get(key) if target is not None and tuple(target.shape) == tuple(value.shape): compatible[key] = value else: skipped.append(key) missing, unexpected = model.load_state_dict(compatible, strict=False) print( "RESUME_FILTERED " f"loaded={len(compatible)} skipped={len(skipped)} " f"missing={list(missing)} unexpected={list(unexpected)} skipped_keys={skipped[:20]}" ) def configure_trainable_parameters(model: nn.Module, freeze_mode: str, train_top_layers: int, train_layernorms: bool) -> int: if freeze_mode == "none": for param in model.parameters(): param.requires_grad = True elif freeze_mode == "rerank": for param in model.parameters(): param.requires_grad = False for name, param in model.named_parameters(): if name.startswith("z_rerank_head"): param.requires_grad = True elif freeze_mode == "calib": for param in model.parameters(): param.requires_grad = False for name, param in model.named_parameters(): if name.startswith("z_calib_head"): param.requires_grad = True elif freeze_mode == "adapter": for param in model.parameters(): param.requires_grad = False train_prefixes = ( "stem", "input_proj", "pos_proj", "pool_gate", "z_head", "z_bin_head", "z_candidate_head", "z_rerank_head", "z_calib_head", "rec_head", "rec_refine", "cls", "z_query", ) for name, param in model.named_parameters(): if name.startswith(train_prefixes): param.requires_grad = True if train_layernorms and (".norm" in name or name.endswith("norm.weight") or name.endswith("norm.bias")): param.requires_grad = True layers = getattr(getattr(model, "encoder", None), "layers", None) if layers is not None and train_top_layers > 0: for layer in list(layers)[-int(train_top_layers) :]: for param in layer.parameters(): param.requires_grad = True else: raise ValueError(f"Unknown freeze mode {freeze_mode!r}") return sum(p.numel() for p in model.parameters() if p.requires_grad) def replay_loss( student_out: dict[str, torch.Tensor], teacher_out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], *, y_weight: float, bin_weight: float, clean_only: bool, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: y_student = student_out.get("y_pred", student_out["y_mu"]).float() y_teacher = teacher_out.get("y_pred", teacher_out["y_mu"]).float().detach() mask = torch.isfinite(batch["y"]) if clean_only: mask = mask & (~batch["zwarn"].bool()) if mask.sum() == 0: zero = y_student.sum() * 0.0 return zero, {"replay_y": zero.detach(), "replay_bin": zero.detach()} replay_y = F.smooth_l1_loss(y_student[mask], y_teacher[mask], beta=0.01) replay_bin = y_student.sum() * 0.0 if bin_weight > 0 and "z_bin_logits" in student_out and "z_bin_logits" in teacher_out: student_logp = F.log_softmax(student_out["z_bin_logits"][mask].float(), dim=-1) teacher_p = F.softmax(teacher_out["z_bin_logits"][mask].float().detach(), dim=-1) replay_bin = F.kl_div(student_logp, teacher_p, reduction="batchmean") total = y_weight * replay_y + bin_weight * replay_bin return total, {"replay_y": replay_y.detach(), "replay_bin": replay_bin.detach()} def redshift_total_loss(model: HybridSpecZ, out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], cfg: LossConfig) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: parts = redshift_losses(model, out, batch["y"], batch["zwarn"], cfg) if "rec" in out and "target_flux" in batch and "loss_mask" in batch: line_weight = batch.get("line_weight") if line_weight is not None: line_weight = line_weight.pow(cfg.line_weight_power) rec = masked_huber(out["rec"], batch["target_flux"], batch["loss_mask"], weight=line_weight) else: rec = parts["z_huber"].sum() * 0.0 total = ( cfg.rec_weight * rec + cfg.z_weight * parts["z_huber"] + cfg.z_bin_weight * parts["z_bin"] + cfg.z_candidate_weight * parts["z_candidate"] + cfg.z_rerank_weight * parts["z_rerank"] + cfg.z_nll_weight * parts["z_nll"] ) metrics = {"loss": total.detach(), "rec": rec.detach(), **{k: v.detach() for k, v in parts.items()}} return total, metrics def plot_spectra_batch(path: str | Path, batch: dict[str, torch.Tensor], y_pred: np.ndarray, max_items: int = 4) -> None: path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) x = batch["x"].detach().cpu().numpy() loglam = batch["loglam"].detach().cpu().numpy() valid = batch["valid"].detach().cpu().numpy() z = batch["z"].detach().cpu().numpy() bsz = min(max_items, x.shape[0]) fig, axes = plt.subplots(bsz, 1, figsize=(13, 3.0 * bsz), squeeze=False) for i in range(bsz): ax = axes[i, 0] wave = np.exp(loglam[i]) good = valid[i].astype(bool) ax.plot(wave[good], x[i, 0, good], color="black", linewidth=0.8, label="input flux") ax.plot(wave[good], x[i, 6, good], color="#1f77b4", linewidth=0.6, alpha=0.55, label="line score") masked = x[i, 7] > 0 if masked.any(): ax.scatter(wave[masked], np.zeros(masked.sum()), s=5, color="#d62728", alpha=0.55, label="redshift dropout") ax.set_title(f"z true={z[i]:.5f} z pred={np.expm1(y_pred[i]):.5f}") ax.set_ylabel("normalized") ax.grid(alpha=0.2) if i == 0: ax.legend(loc="best", fontsize=8) axes[-1, 0].set_xlabel("wavelength Angstrom") fig.tight_layout() fig.savefig(path, dpi=150) plt.close(fig) def add_redshift_slice_metrics(metrics: dict[str, float], prefix: str, y_true: np.ndarray, y_pred: np.ndarray) -> None: z_true = np.expm1(y_true) z_pred = np.expm1(y_pred) slices = { "z_lt_0p4": z_true < 0.4, "z_0p4_1p0": (z_true >= 0.4) & (z_true < 1.0), "z_1p0_2p0": (z_true >= 1.0) & (z_true < 2.0), "z_gte_2p0": z_true >= 2.0, } for name, mask in slices.items(): count = int(np.count_nonzero(mask)) metrics[f"{prefix}/{name}_count"] = float(count) if count >= 5: err = z_pred[mask] - z_true[mask] denom = 1.0 + z_true[mask] metrics[f"{prefix}/{name}_mae_z"] = float(np.mean(np.abs(err))) metrics[f"{prefix}/{name}_bias_z"] = float(np.mean(err)) metrics[f"{prefix}/{name}_cat_0p05"] = float(np.mean(np.abs(err / denom) > 0.05)) def add_candidate_metrics( metrics: dict[str, float], prefix: str, y_true: np.ndarray, candidate_y: np.ndarray, candidate_bins: np.ndarray | None, *, z_bins: int, y_min: float, y_max: float, ) -> None: if candidate_y.size == 0: return z_true = np.expm1(y_true) z_candidate = np.expm1(candidate_y) abs_dz = np.abs(z_candidate - z_true[:, None]) norm_dz = abs_dz / (1.0 + z_true[:, None]) top_limits = [1, 3, 5] for k in top_limits: kk = min(k, candidate_y.shape[1]) best_abs = np.min(abs_dz[:, :kk], axis=1) best_norm = np.min(norm_dz[:, :kk], axis=1) metrics[f"{prefix}/candidate_top{kk}_best_mae_z"] = float(np.mean(best_abs)) metrics[f"{prefix}/candidate_top{kk}_hit_0p003"] = float(np.mean(best_norm <= 0.003)) metrics[f"{prefix}/candidate_top{kk}_hit_0p01"] = float(np.mean(best_norm <= 0.01)) metrics[f"{prefix}/candidate_top{kk}_hit_0p05"] = float(np.mean(best_norm <= 0.05)) if candidate_bins is not None and candidate_bins.size: scaled = (y_true - y_min) / max(y_max - y_min, 1e-6) true_bins = np.clip((scaled * z_bins).astype(np.int64), 0, z_bins - 1) for k in top_limits: kk = min(k, candidate_bins.shape[1]) metrics[f"{prefix}/candidate_top{kk}_bin_hit"] = float(np.mean(np.any(candidate_bins[:, :kk] == true_bins[:, None], axis=1))) @torch.no_grad() def evaluate( model: HybridSpecZ, loader: DataLoader, loss_cfg: LossConfig, device: torch.device, run_dir: Path, step: int, prefix: str = "val", max_batches: int | None = 50, max_examples: int | None = None, ) -> dict[str, float]: model.eval() losses: dict[str, list[float]] = {} y_true_all: list[np.ndarray] = [] y_pred_all: list[np.ndarray] = [] candidate_y_all: list[np.ndarray] = [] candidate_bins_all: list[np.ndarray] = [] y_true_clean: list[np.ndarray] = [] y_pred_clean: list[np.ndarray] = [] candidate_y_clean: list[np.ndarray] = [] candidate_bins_clean: list[np.ndarray] = [] zwarn_all: list[np.ndarray] = [] first_batch = None first_pred = None first_rec = None seen_examples = 0 for bi, batch in enumerate(loader): if max_batches is not None and max_batches > 0 and bi >= max_batches: break batch = limit_batch_examples(batch, max_examples, seen_examples) if batch is None: break seen_examples += int(batch["y"].shape[0]) batch = move_to_device(batch, device) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): out = model(batch["x"], batch["valid"], batch["loglam"]) _, parts = redshift_total_loss(model, out, batch, loss_cfg) y_pred = out.get("y_pred", out["y_mu"]) for k, v in parts.items(): losses.setdefault(k, []).append(float(v.detach().cpu())) if "rec" in out and "target_flux" in batch and "loss_mask" in batch: rec_err = F.smooth_l1_loss(out["rec"].float(), batch["target_flux"].float(), reduction="none", beta=0.5) loss_mask = batch["loss_mask"].bool() line_region = batch.get("line_region") if line_region is not None: line_mask = loss_mask & line_region.bool() cont_mask = loss_mask & (~line_region.bool()) for name, mask in (("rec_line", line_mask), ("rec_continuum", cont_mask)): denom = mask.float().sum() if float(denom.detach().cpu()) > 0: losses.setdefault(name, []).append(float(((rec_err * mask.float()).sum() / denom.clamp_min(1.0)).detach().cpu())) context_mask = batch["valid"].bool() & (~loss_mask) denom = context_mask.float().sum(dim=1).clamp_min(1.0) baseline = (batch["target_flux"].float() * context_mask.float()).sum(dim=1, keepdim=True) / denom.unsqueeze(1) baseline_err = F.smooth_l1_loss(baseline.expand_as(batch["target_flux"]).float(), batch["target_flux"].float(), reduction="none", beta=0.5) mask_denom = loss_mask.float().sum().clamp_min(1.0) losses.setdefault("rec_mean_baseline", []).append(float(((baseline_err * loss_mask.float()).sum() / mask_denom).detach().cpu())) finite = torch.isfinite(batch["y"]).detach().cpu().numpy() clean = ((~batch["zwarn"].bool()) & torch.isfinite(batch["y"])).detach().cpu().numpy() zw = batch["zwarn"].detach().cpu().numpy().astype(bool) if finite.any(): y_true_all.append(batch["y"].detach().cpu().numpy()[finite]) y_pred_all.append(y_pred.float().detach().cpu().numpy()[finite]) zwarn_all.append(zw[finite]) if "candidate_topk_y" in out: candidate_y_all.append(out["candidate_topk_y"].float().detach().cpu().numpy()[finite]) if "candidate_topk_bins" in out: candidate_bins_all.append(out["candidate_topk_bins"].detach().cpu().numpy()[finite]) if clean.any(): y_true_clean.append(batch["y"].detach().cpu().numpy()[clean]) y_pred_clean.append(y_pred.float().detach().cpu().numpy()[clean]) if "candidate_topk_y" in out: candidate_y_clean.append(out["candidate_topk_y"].float().detach().cpu().numpy()[clean]) if "candidate_topk_bins" in out: candidate_bins_clean.append(out["candidate_topk_bins"].detach().cpu().numpy()[clean]) if first_batch is None: first_batch = {k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in batch.items()} first_pred = y_pred.float().detach().cpu().numpy() if "rec" in out: first_rec = out["rec"].float().detach().cpu().numpy() metrics = {f"{prefix}/{k}": float(np.mean(v)) for k, v in losses.items()} if y_true_all: y_true = np.concatenate(y_true_all) y_pred = np.concatenate(y_pred_all) for k, v in redshift_metrics(y_true, y_pred).items(): metrics[f"{prefix}/{k}"] = v add_redshift_slice_metrics(metrics, prefix, y_true, y_pred) if candidate_y_all: candidate_y_np = np.concatenate(candidate_y_all) candidate_bins_np = np.concatenate(candidate_bins_all) if candidate_bins_all else None add_candidate_metrics( metrics, prefix, y_true, candidate_y_np, candidate_bins_np, z_bins=model.z_bins, y_min=model.y_min, y_max=model.y_max, ) metrics[f"{prefix}/z_count"] = float(len(y_true)) metrics[f"{prefix}/zwarn_fraction"] = float(np.mean(np.concatenate(zwarn_all))) if zwarn_all else 0.0 plot_redshift_scatter(run_dir / "plots" / f"{prefix}_redshift_step_{step:06d}.png", y_true, y_pred) if y_true_clean: clean_true = np.concatenate(y_true_clean) clean_pred = np.concatenate(y_pred_clean) if len(clean_true) >= 5: for k, v in redshift_metrics(clean_true, clean_pred).items(): metrics[f"{prefix}_clean/{k}"] = v if candidate_y_clean: candidate_y_clean_np = np.concatenate(candidate_y_clean) candidate_bins_clean_np = np.concatenate(candidate_bins_clean) if candidate_bins_clean else None add_candidate_metrics( metrics, f"{prefix}_clean", clean_true, candidate_y_clean_np, candidate_bins_clean_np, z_bins=model.z_bins, y_min=model.y_min, y_max=model.y_max, ) metrics[f"{prefix}_clean/z_count"] = float(len(clean_true)) if first_batch is not None and first_pred is not None: if first_rec is not None and "target_flux" in first_batch and "loss_mask" in first_batch: plot_reconstruction_batch( run_dir / "plots" / f"{prefix}_reconstruction_step_{step:06d}.png", first_batch["loglam"].numpy(), first_batch["target_flux"].numpy(), first_rec, first_batch["loss_mask"].numpy(), first_batch["valid"].numpy(), first_batch["z"].numpy(), np.expm1(first_pred), ) plot_spectra_batch(run_dir / "plots" / f"{prefix}_spectra_step_{step:06d}.png", first_batch, first_pred) model.train() return metrics def make_loader( samples: list[dict[str, Any]], indices: np.ndarray, cfg: RawCollatorConfig, args: argparse.Namespace, train: bool, sampler: WeightedRandomSampler | None = None, ) -> DataLoader: return DataLoader( SpectraListDataset(samples, indices), batch_size=args.batch_size, shuffle=train and sampler is None, sampler=sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=RawSpectraCollator(cfg, train=train, seed=args.seed + (0 if train else 1000)), ) def checkpoint_score( mode: str, val_metrics: dict[str, float], ood_metrics: dict[str, float] | None, z_alpha: float = 0.6, desi_mae_ceiling: float = 0.0, desi_mae_penalty: float = 0.0, ) -> float: def score_prefix(metrics: dict[str, float], prefix: str) -> float: z_score = ( metrics.get(f"{prefix}/nmad", math.inf) + metrics.get(f"{prefix}/cat_0p01", 1.0) + metrics.get(f"{prefix}/mae_log1p", 1.0) ) rec_score = metrics.get(f"{prefix}/rec") if rec_score is None or not math.isfinite(float(rec_score)): return z_score alpha = min(max(float(z_alpha), 0.0), 1.0) return alpha * z_score + (1.0 - alpha) * float(rec_score) val_score = score_prefix(val_metrics, "val") if mode == "rec": return float(val_metrics.get("val/rec", math.inf)) if mode == "val" or ood_metrics is None: score = val_score if desi_mae_ceiling > 0 and desi_mae_penalty > 0: val_mae = float(val_metrics.get("val/mae_z", 0.0)) score += float(desi_mae_penalty) * max(0.0, val_mae - float(desi_mae_ceiling)) return score ood_score = score_prefix(ood_metrics, "ood") if mode == "ood": score = ood_score else: score = 0.5 * val_score + 0.5 * ood_score if desi_mae_ceiling > 0 and desi_mae_penalty > 0: val_mae = float(val_metrics.get("val/mae_z", 0.0)) score += float(desi_mae_penalty) * max(0.0, val_mae - float(desi_mae_ceiling)) return score def scheduled_lr(base_lr: float, min_lr: float, step: int, total_steps: int, warmup_steps: int) -> float: if warmup_steps > 0 and step <= warmup_steps: return base_lr * float(step) / float(max(1, warmup_steps)) if min_lr < 0 or total_steps <= warmup_steps: return base_lr progress = (step - warmup_steps) / float(max(1, total_steps - warmup_steps)) progress = min(max(progress, 0.0), 1.0) return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress)) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--dataset-name", default="MultimodalUniverse/desi") parser.add_argument("--max-samples", type=int, default=4096) parser.add_argument("--cache-dir", default="/workspace/native_specz_mae/cache") parser.add_argument("--hf-cache-dir", default=os.environ.get("HF_DATASETS_CACHE", "/workspace/hf_cache/datasets")) parser.add_argument("--run-dir", default="/workspace/runs/hybrid_specz") parser.add_argument("--resume-checkpoint", default="") parser.add_argument("--allow-mismatched-checkpoint", action="store_true") parser.add_argument("--refresh-data", action="store_true") parser.add_argument("--epochs", type=int, default=8) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--num-workers", type=int, default=2) parser.add_argument("--target-length", type=int, default=4096) parser.add_argument("--architecture", choices=["conv", "wave_token"], default="conv") parser.add_argument("--d-model", type=int, default=256) parser.add_argument("--conv-width", type=int, default=128) parser.add_argument("--layers", type=int, default=5) parser.add_argument("--heads", type=int, default=8) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--z-bins", type=int, default=64) parser.add_argument("--stem-stride", type=int, choices=[4, 8], default=8) parser.add_argument("--token-stride", type=int, default=8) parser.add_argument("--rec-hidden-mult", type=int, default=0) parser.add_argument("--rec-refine-width", type=int, default=16) parser.add_argument("--rec-refine-kernel", type=int, default=5) parser.add_argument("--layerscale-init", type=float, default=0.0) parser.add_argument( "--prediction-mode", choices=[ "regression", "softbin", "hybrid", "bin_residual", "ranked_bin_residual", "candidate_rerank", "calibrated_bin_residual", ], default="regression", ) parser.add_argument("--bin-temperature", type=float, default=1.0) parser.add_argument("--residual-scale", type=float, default=0.06) parser.add_argument("--candidate-topk", type=int, default=5) parser.add_argument("--lr", type=float, default=2e-4) parser.add_argument("--min-lr", type=float, default=-1.0) parser.add_argument("--warmup-steps", type=int, default=0) parser.add_argument("--weight-decay", type=float, default=0.03) parser.add_argument("--grad-clip", type=float, default=1.0) parser.add_argument("--grad-accum-steps", type=int, default=1) parser.add_argument("--eval-every", type=int, default=100) parser.add_argument("--eval-max-val", type=int, default=800) parser.add_argument("--eval-max-ood", type=int, default=480) parser.add_argument("--max-steps", type=int, default=0) parser.add_argument("--checkpoint-score", choices=["val", "ood", "combined", "rec"], default="combined") parser.add_argument("--score-z-alpha", type=float, default=0.6) parser.add_argument("--desi-mae-ceiling", type=float, default=0.0) parser.add_argument("--desi-mae-penalty", type=float, default=0.0) parser.add_argument("--objective", choices=["joint", "rec_only", "z_only"], default="joint") parser.add_argument("--freeze-mode", choices=["none", "adapter", "rerank", "calib"], default="none") parser.add_argument("--train-top-layers", type=int, default=0) parser.add_argument("--train-layernorms", action="store_true") parser.add_argument("--replay-checkpoint", default="") parser.add_argument("--replay-y-weight", type=float, default=0.0) parser.add_argument("--replay-bin-weight", type=float, default=0.0) parser.add_argument("--replay-clean-only", action="store_true") parser.add_argument("--balance-redshift", action="store_true") parser.add_argument("--train-clean-only", action="store_true") parser.add_argument("--clean-sample-boost", type=float, default=1.0) parser.add_argument("--augment-ood", action="store_true") parser.add_argument("--eval-ood", action="store_true") parser.add_argument("--random-mask-ratio", type=float, default=0.0) parser.add_argument("--eval-mask-ratio", type=float, default=0.25) parser.add_argument("--mask-mode", choices=["pixel", "span", "line_span", "mixed_span"], default="pixel") parser.add_argument("--mask-span-min", type=int, default=16) parser.add_argument("--mask-span-max", type=int, default=64) parser.add_argument("--line-region-percentile", type=float, default=90.0) parser.add_argument("--crop-prob", type=float, default=0.0) parser.add_argument("--bad-window-prob", type=float, default=0.0) parser.add_argument("--throughput-prob", type=float, default=0.0) parser.add_argument("--noise-prob", type=float, default=0.0) parser.add_argument("--resolution-prob", type=float, default=0.0) parser.add_argument("--downsample-prob", type=float, default=0.0) parser.add_argument("--line-dropout-prob", type=float, default=0.0) parser.add_argument("--span-dropout-prob", type=float, default=0.0) parser.add_argument("--grid-jitter-prob", type=float, default=0.0) parser.add_argument("--grid-shift-frac", type=float, default=0.0) parser.add_argument("--grid-scale-frac", type=float, default=0.0) parser.add_argument("--grid-jitter-warmup-steps", type=int, default=0) parser.add_argument("--redshift-shift", type=float, default=0.0) parser.add_argument("--rec-weight", type=float, default=0.0) parser.add_argument("--z-weight", type=float, default=1.0) parser.add_argument("--z-bin-weight", type=float, default=0.25) parser.add_argument("--z-candidate-weight", type=float, default=0.0) parser.add_argument("--z-rerank-weight", type=float, default=0.0) parser.add_argument("--z-nll-weight", type=float, default=0.05) parser.add_argument("--zwarn-weight", type=float, default=0.3) parser.add_argument("--high-z-boost", type=float, default=1.0) parser.add_argument("--high-z-threshold", type=float, default=1.0) parser.add_argument("--clean-z-only", action="store_true") parser.add_argument("--seed", type=int, default=17) args = parser.parse_args() torch.manual_seed(args.seed) np.random.seed(args.seed) run_dir = Path(args.run_dir) / time.strftime("%Y%m%d_%H%M%S") run_dir.mkdir(parents=True, exist_ok=True) (run_dir / "args.json").write_text(json.dumps(vars(args), indent=2), encoding="utf-8") samples = collect_mmu_desi( cache_file=Path(args.cache_dir) / f"desi_{args.max_samples}.pt", max_samples=args.max_samples, dataset_name=args.dataset_name, hf_cache_dir=args.hf_cache_dir, refresh=args.refresh_data, ) stats = compute_sample_stats(samples) (run_dir / "data_stats.json").write_text(json.dumps(stats.__dict__, indent=2), encoding="utf-8") print("DATA_STATS", json.dumps(stats.__dict__, sort_keys=True)) train_idx, val_idx = split_indices(len(samples), val_fraction=0.15, seed=args.seed) if args.train_clean_only: clean_train = np.asarray([not bool(samples[int(i)]["zwarn"]) for i in train_idx], dtype=np.bool_) train_idx = train_idx[clean_train] if len(train_idx) == 0: raise RuntimeError("No clean ZWARN==0 samples are available for --train-clean-only.") print(f"TRAIN_CLEAN_ONLY n_train={len(train_idx)}") sampler = None if args.balance_redshift or args.clean_sample_boost != 1.0: weights = np.ones(len(train_idx), dtype=np.float32) y_train = np.asarray([np.log1p(float(samples[int(i)]["z"])) for i in train_idx], dtype=np.float32) if args.balance_redshift: bins = np.linspace(float(y_train.min()), float(y_train.max()) + 1e-6, 28) bin_id = np.clip(np.digitize(y_train, bins) - 1, 0, len(bins) - 2) counts = np.bincount(bin_id, minlength=len(bins) - 1).astype(np.float32) weights *= 1.0 / np.maximum(counts[bin_id], 1.0) if args.clean_sample_boost != 1.0: clean = np.asarray([not bool(samples[int(i)]["zwarn"]) for i in train_idx], dtype=np.bool_) weights *= np.where(clean, float(args.clean_sample_boost), 1.0).astype(np.float32) weights = weights / weights.mean() sampler = WeightedRandomSampler(torch.as_tensor(weights, dtype=torch.double), num_samples=len(weights), replacement=True) train_cfg = RawCollatorConfig( target_length=args.target_length, random_mask_ratio=args.random_mask_ratio, eval_mask_ratio=args.eval_mask_ratio, mask_mode=args.mask_mode, mask_span_min=args.mask_span_min, mask_span_max=args.mask_span_max, line_region_percentile=args.line_region_percentile, augment_ood=args.augment_ood, crop_prob=args.crop_prob, bad_window_prob=args.bad_window_prob, throughput_prob=args.throughput_prob, noise_prob=args.noise_prob, resolution_prob=args.resolution_prob, downsample_prob=args.downsample_prob, line_dropout_prob=args.line_dropout_prob, span_dropout_prob=args.span_dropout_prob, grid_jitter_prob=args.grid_jitter_prob, grid_shift_frac=args.grid_shift_frac, grid_scale_frac=args.grid_scale_frac, grid_jitter_warmup_steps=args.grid_jitter_warmup_steps, redshift_shift=args.redshift_shift, ) val_cfg = RawCollatorConfig( target_length=args.target_length, eval_mask_ratio=args.eval_mask_ratio, mask_mode=args.mask_mode, mask_span_min=args.mask_span_min, mask_span_max=args.mask_span_max, line_region_percentile=args.line_region_percentile, ) ood_cfg = RawCollatorConfig( target_length=args.target_length, eval_mask_ratio=args.eval_mask_ratio, mask_mode=args.mask_mode, mask_span_min=args.mask_span_min, mask_span_max=args.mask_span_max, line_region_percentile=args.line_region_percentile, augment_ood=True, crop_prob=0.65, bad_window_prob=0.45, throughput_prob=0.65, noise_prob=0.35, resolution_prob=0.45, downsample_prob=0.35, grid_jitter_prob=0.65, grid_shift_frac=0.04, grid_scale_frac=0.08, ) train_loader = make_loader(samples, train_idx, train_cfg, args, train=True, sampler=sampler) val_loader = make_loader(samples, val_idx, val_cfg, args, train=False) ood_loader = make_loader(samples, val_idx, ood_cfg, args, train=False) if args.eval_ood else None device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_kwargs = dict( d_model=args.d_model, conv_width=args.conv_width, layers=args.layers, heads=args.heads, dropout=args.dropout, z_bins=args.z_bins, rec_hidden_mult=args.rec_hidden_mult, rec_refine_width=args.rec_refine_width, rec_refine_kernel=args.rec_refine_kernel, layerscale_init=args.layerscale_init, prediction_mode=args.prediction_mode, bin_temperature=args.bin_temperature, residual_scale=args.residual_scale, candidate_topk=args.candidate_topk, ) if args.architecture == "wave_token": model = WavelengthTokenSpecZ(token_stride=args.token_stride, **model_kwargs).to(device) else: model = HybridSpecZ(stem_stride=args.stem_stride, **model_kwargs).to(device) if args.resume_checkpoint: ckpt = torch.load(args.resume_checkpoint, map_location=device, weights_only=False) state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt load_checkpoint_into_model(model, state, allow_mismatched=args.allow_mismatched_checkpoint) print(f"RESUME_CHECKPOINT {args.resume_checkpoint}") n_params = sum(p.numel() for p in model.parameters()) print(f"MODEL_PARAMS {n_params}") teacher_model = None if args.replay_y_weight > 0 or args.replay_bin_weight > 0: replay_path = args.replay_checkpoint or args.resume_checkpoint if not replay_path: raise RuntimeError("--replay-checkpoint or --resume-checkpoint is required when replay weights are nonzero.") if args.architecture == "wave_token": teacher_model = WavelengthTokenSpecZ(token_stride=args.token_stride, **model_kwargs).to(device) else: teacher_model = HybridSpecZ(stem_stride=args.stem_stride, **model_kwargs).to(device) ckpt = torch.load(replay_path, map_location=device, weights_only=False) state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt load_checkpoint_into_model(teacher_model, state, allow_mismatched=args.allow_mismatched_checkpoint) teacher_model.eval() for param in teacher_model.parameters(): param.requires_grad = False print(f"REPLAY_TEACHER {replay_path}") trainable_params = configure_trainable_parameters(model, args.freeze_mode, args.train_top_layers, args.train_layernorms) print(f"TRAINABLE_PARAMS {trainable_params} freeze_mode={args.freeze_mode} train_top_layers={args.train_top_layers}") opt_params = [p for p in model.parameters() if p.requires_grad] if not opt_params: raise RuntimeError("No trainable parameters remain after freeze configuration.") optimizer = AdamW(opt_params, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95)) rec_weight = args.rec_weight z_weight = args.z_weight z_bin_weight = args.z_bin_weight z_candidate_weight = args.z_candidate_weight z_rerank_weight = args.z_rerank_weight z_nll_weight = args.z_nll_weight if args.objective == "rec_only": rec_weight = rec_weight if rec_weight > 0 else 1.0 z_weight = 0.0 z_bin_weight = 0.0 z_candidate_weight = 0.0 z_rerank_weight = 0.0 z_nll_weight = 0.0 elif args.objective == "z_only": rec_weight = 0.0 loss_cfg = LossConfig( rec_weight=rec_weight, z_weight=z_weight, z_bin_weight=z_bin_weight, z_candidate_weight=z_candidate_weight, z_rerank_weight=z_rerank_weight, z_nll_weight=z_nll_weight, zwarn_weight=args.zwarn_weight, clean_z_only=args.clean_z_only, high_z_boost=args.high_z_boost, high_z_threshold=math.log1p(args.high_z_threshold), ) best_score = math.inf global_step = 0 micro_step = 0 grad_accum_steps = max(1, int(args.grad_accum_steps)) total_train_steps = args.max_steps if args.max_steps else int(math.ceil(len(train_loader) / grad_accum_steps) * args.epochs) model.train() optimizer.zero_grad(set_to_none=True) for epoch in range(args.epochs): pbar = tqdm(train_loader, desc=f"hybrid epoch {epoch}") for batch in pbar: micro_step += 1 batch = move_to_device(batch, device) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): out = model(batch["x"], batch["valid"], batch["loglam"]) loss, parts = redshift_total_loss(model, out, batch, loss_cfg) if teacher_model is not None: with torch.no_grad(): teacher_out = teacher_model(batch["x"], batch["valid"], batch["loglam"]) replay, replay_parts = replay_loss( out, teacher_out, batch, y_weight=args.replay_y_weight, bin_weight=args.replay_bin_weight, clean_only=args.replay_clean_only, ) loss = loss + replay parts = {**parts, "loss": parts["loss"] + replay.detach(), "replay": replay.detach(), **replay_parts} (loss / grad_accum_steps).backward() if micro_step % grad_accum_steps != 0: pbar.set_postfix( loss=float(parts["loss"].detach().cpu()), rec=float(parts["rec"].detach().cpu()), huber=float(parts["z_huber"].detach().cpu()), replay=float(parts.get("replay", parts["loss"].sum() * 0.0).detach().cpu()), accum=f"{micro_step % grad_accum_steps}/{grad_accum_steps}", ) continue next_step = global_step + 1 lr_now = scheduled_lr(args.lr, args.min_lr, next_step, total_train_steps, int(args.warmup_steps)) for group in optimizer.param_groups: group["lr"] = lr_now grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() optimizer.zero_grad(set_to_none=True) global_step = next_step pbar.set_postfix( loss=float(parts["loss"].detach().cpu()), rec=float(parts["rec"].detach().cpu()), huber=float(parts["z_huber"].detach().cpu()), replay=float(parts.get("replay", parts["loss"].sum() * 0.0).detach().cpu()), lr=lr_now, grad=float(grad_norm.detach().cpu()) if torch.is_tensor(grad_norm) else float(grad_norm), ) if global_step == 1 or global_step % args.eval_every == 0: val_metrics = evaluate( model, val_loader, loss_cfg, device, run_dir, global_step, prefix="val", max_batches=None, max_examples=args.eval_max_val, ) print("VAL", global_step, json.dumps(val_metrics, sort_keys=True)) ood_metrics = None if ood_loader is not None: ood_metrics = evaluate( model, ood_loader, loss_cfg, device, run_dir, global_step, prefix="ood", max_batches=None, max_examples=args.eval_max_ood, ) print("OOD", global_step, json.dumps(ood_metrics, sort_keys=True)) score = checkpoint_score( args.checkpoint_score, val_metrics, ood_metrics, z_alpha=args.score_z_alpha, desi_mae_ceiling=args.desi_mae_ceiling, desi_mae_penalty=args.desi_mae_penalty, ) if score < best_score: best_score = score best_metrics = {"step": global_step, "score": best_score, **val_metrics} if ood_metrics is not None: best_metrics.update(ood_metrics) torch.save( {"model": model.state_dict(), "args": vars(args), "step": global_step, "score": best_score, "metrics": best_metrics}, run_dir / "best.pt", ) (run_dir / "best_metrics.json").write_text(json.dumps(best_metrics, indent=2), encoding="utf-8") if args.max_steps and global_step >= args.max_steps: break if args.max_steps and global_step >= args.max_steps: break final_metrics = evaluate( model, val_loader, loss_cfg, device, run_dir, global_step, prefix="val", max_batches=None, max_examples=args.eval_max_val, ) if ood_loader is not None: final_metrics.update( evaluate( model, ood_loader, loss_cfg, device, run_dir, global_step, prefix="ood", max_batches=None, max_examples=args.eval_max_ood, ) ) torch.save({"model": model.state_dict(), "args": vars(args), "step": global_step, "metrics": final_metrics}, run_dir / "last.pt") (run_dir / "final_metrics.json").write_text(json.dumps(final_metrics, indent=2), encoding="utf-8") print("FINAL", json.dumps(final_metrics, sort_keys=True)) print("RUN_DIR", run_dir) if __name__ == "__main__": main()