Upload code/hybrid_redshift.py

code/hybrid_redshift.py

@@ -0,0 +1,1278 @@
+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()