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 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) def __call__(self, samples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: 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 = np.linspace(float(np.nanmin(loglam)), float(np.nanmax(loglam)), 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"} 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), ) 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) 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) 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 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_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, } 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) 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 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_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) -> 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: return val_score ood_score = score_prefix(ood_metrics, "ood") if mode == "ood": return ood_score return 0.5 * val_score + 0.5 * ood_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("--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("--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"], 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("--objective", choices=["joint", "rec_only", "z_only"], default="joint") 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("--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-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, 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, ) 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 = HybridSpecZ( d_model=args.d_model, conv_width=args.conv_width, layers=args.layers, heads=args.heads, dropout=args.dropout, z_bins=args.z_bins, stem_stride=args.stem_stride, 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, ).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}") optimizer = AdamW(model.parameters(), 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_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_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_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) (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()), 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()), 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) 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()