code/data.py

from __future__ import annotations

import math
import os
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import torch
from datasets import load_dataset
from torch.utils.data import Dataset


REQUIRED_SPECTRUM_KEYS = ("flux", "ivar", "lambda", "mask")


@dataclass
class SampleStats:
    n_samples: int
    min_len: int
    max_len: int
    median_len: float
    min_lambda: float
    max_lambda: float
    valid_fraction: float
    z_min: float
    z_max: float
    z_median: float
    zwarn_fraction: float


def _as_array(value: Any, dtype: np.dtype) -> np.ndarray:
    arr = np.asarray(value, dtype=dtype)
    return np.ascontiguousarray(arr)


def parse_mmu_example(example: dict[str, Any]) -> dict[str, Any] | None:
    spectrum = example.get("spectrum")
    if not isinstance(spectrum, dict):
        return None
    if any(key not in spectrum for key in REQUIRED_SPECTRUM_KEYS):
        return None

    flux = _as_array(spectrum["flux"], np.float32)
    ivar = _as_array(spectrum["ivar"], np.float32)
    lam = _as_array(spectrum["lambda"], np.float32)
    bad_mask = _as_array(spectrum["mask"], np.bool_)
    if len(flux) == 0 or not (len(flux) == len(ivar) == len(lam) == len(bad_mask)):
        return None

    lsf = spectrum.get("lsf_sigma")
    lsf_sigma = _as_array(lsf, np.float32) if lsf is not None and len(lsf) == len(flux) else np.zeros_like(flux)

    z = float(example.get("Z", np.nan))
    zerr = float(example.get("ZERR", np.nan))
    zwarn = bool(example.get("ZWARN", True))
    if not math.isfinite(z) or z < -0.001:
        return None

    return {
        "flux": flux,
        "ivar": ivar,
        "lambda": lam,
        "bad_mask": bad_mask,
        "lsf_sigma": lsf_sigma,
        "z": np.float32(max(z, 0.0)),
        "zerr": np.float32(zerr if math.isfinite(zerr) else np.nan),
        "zwarn": zwarn,
        "object_id": str(example.get("object_id", "")),
    }


def collect_mmu_desi(
    cache_file: str | os.PathLike[str],
    max_samples: int,
    dataset_name: str = "MultimodalUniverse/desi",
    split: str = "train",
    seed: int = 17,
    hf_cache_dir: str | None = None,
    refresh: bool = False,
) -> list[dict[str, Any]]:
    cache_path = Path(cache_file)
    if cache_path.exists() and not refresh:
        payload = torch.load(cache_path, map_location="cpu", weights_only=False)
        samples = payload["samples"] if isinstance(payload, dict) and "samples" in payload else payload
        if len(samples) >= max_samples:
            return samples[:max_samples]

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    ds = load_dataset(dataset_name, split=split, streaming=True, cache_dir=hf_cache_dir)
    shuffle_buffer = min(max_samples, 10_000)
    print(f"COLLECT_SHUFFLE_BUFFER {shuffle_buffer}", flush=True)
    ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
    samples: list[dict[str, Any]] = []
    # Deterministic shard order is fine for smoke runs; shuffle after collection.
    for example in ds:
        parsed = parse_mmu_example(example)
        if parsed is None:
            continue
        samples.append(parsed)
        if len(samples) % 4096 == 0:
            print(f"COLLECTED_SAMPLES {len(samples)}/{max_samples}", flush=True)
        if len(samples) >= max_samples:
            break
    rng = random.Random(seed)
    rng.shuffle(samples)
    torch.save({"samples": samples, "dataset_name": dataset_name, "split": split}, cache_path)
    return samples


def compute_sample_stats(samples: Iterable[dict[str, Any]]) -> SampleStats:
    samples = list(samples)
    lengths = np.asarray([len(s["flux"]) for s in samples], dtype=np.int64)
    mins = np.asarray([np.nanmin(s["lambda"]) for s in samples], dtype=np.float32)
    maxs = np.asarray([np.nanmax(s["lambda"]) for s in samples], dtype=np.float32)
    z = np.asarray([s["z"] for s in samples], dtype=np.float32)
    zwarn = np.asarray([s["zwarn"] for s in samples], dtype=np.bool_)
    valid_fracs = []
    for s in samples:
        valid = valid_pixel_mask(s)
        valid_fracs.append(float(valid.mean()) if len(valid) else 0.0)
    return SampleStats(
        n_samples=len(samples),
        min_len=int(lengths.min()),
        max_len=int(lengths.max()),
        median_len=float(np.median(lengths)),
        min_lambda=float(np.nanmin(mins)),
        max_lambda=float(np.nanmax(maxs)),
        valid_fraction=float(np.mean(valid_fracs)),
        z_min=float(np.nanmin(z)),
        z_max=float(np.nanmax(z)),
        z_median=float(np.nanmedian(z)),
        zwarn_fraction=float(np.mean(zwarn)),
    )


def valid_pixel_mask(sample: dict[str, Any]) -> np.ndarray:
    flux = sample["flux"]
    ivar = sample["ivar"]
    lam = sample["lambda"]
    bad = sample["bad_mask"]
    return np.isfinite(flux) & np.isfinite(ivar) & np.isfinite(lam) & (ivar > 0) & (~bad)


class SpectraListDataset(Dataset):
    def __init__(self, samples: list[dict[str, Any]], indices: np.ndarray):
        self.samples = samples
        self.indices = np.asarray(indices, dtype=np.int64)

    def __len__(self) -> int:
        return int(len(self.indices))

    def __getitem__(self, idx: int) -> dict[str, Any]:
        return self.samples[int(self.indices[idx])]


def split_indices(n: int, val_fraction: float = 0.15, seed: int = 17) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(seed)
    perm = rng.permutation(n)
    n_val = max(1, int(round(n * val_fraction)))
    return perm[n_val:], perm[:n_val]


@dataclass
class CollatorConfig:
    num_patches: int = 256
    random_mask_ratio: float = 0.25
    span_mask_prob: float = 0.65
    line_mask_prob: float = 0.45
    arm_dropout_prob: float = 0.25
    min_scale: float = 1e-3
    max_mask_ratio: float = 0.70
    augment_ood: bool = False


class SpectraCollator:
    def __init__(self, cfg: CollatorConfig, train: bool = True, seed: int = 17):
        self.cfg = cfg
        self.train = train
        self.rng = np.random.default_rng(seed)

    def __call__(self, samples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
        patched = [self._patch_sample(s) for s in samples]
        bsz = len(patched)
        p = self.cfg.num_patches
        fdim = patched[0]["enc_features"].shape[-1]
        max_visible = max(x["enc_features"].shape[0] for x in patched)

        enc_features = np.zeros((bsz, max_visible, fdim), dtype=np.float32)
        enc_loglam = np.zeros((bsz, max_visible), dtype=np.float32)
        enc_padding = np.ones((bsz, max_visible), dtype=np.bool_)

        target_loglam = np.zeros((bsz, p), dtype=np.float32)
        target_aux = np.zeros((bsz, p, 4), dtype=np.float32)
        target_flux = np.zeros((bsz, p), dtype=np.float32)
        loss_mask = np.zeros((bsz, p), dtype=np.bool_)
        valid_patch = np.zeros((bsz, p), dtype=np.bool_)
        corrupt_patch = np.zeros((bsz, p), dtype=np.bool_)
        line_weight = np.ones((bsz, p), dtype=np.float32)
        z = np.zeros((bsz,), dtype=np.float32)
        y = np.zeros((bsz,), dtype=np.float32)
        zwarn = np.zeros((bsz,), dtype=np.bool_)

        for i, item in enumerate(patched):
            nvis = item["enc_features"].shape[0]
            enc_features[i, :nvis] = item["enc_features"]
            enc_loglam[i, :nvis] = item["enc_loglam"]
            enc_padding[i, :nvis] = False
            target_loglam[i] = item["target_loglam"]
            target_aux[i] = item["target_aux"]
            target_flux[i] = item["target_flux"]
            loss_mask[i] = item["loss_mask"]
            valid_patch[i] = item["valid_patch"]
            corrupt_patch[i] = item["corrupt_patch"]
            line_weight[i] = item["line_weight"]
            z[i] = item["z"]
            y[i] = math.log1p(float(item["z"]))
            zwarn[i] = item["zwarn"]

        return {
            "enc_features": torch.from_numpy(enc_features),
            "enc_loglam": torch.from_numpy(enc_loglam),
            "enc_padding": torch.from_numpy(enc_padding),
            "target_loglam": torch.from_numpy(target_loglam),
            "target_aux": torch.from_numpy(target_aux),
            "target_flux": torch.from_numpy(target_flux),
            "loss_mask": torch.from_numpy(loss_mask),
            "valid_patch": torch.from_numpy(valid_patch),
            "corrupt_patch": torch.from_numpy(corrupt_patch),
            "line_weight": torch.from_numpy(line_weight),
            "z": torch.from_numpy(z),
            "y": torch.from_numpy(y),
            "zwarn": torch.from_numpy(zwarn),
        }

    def _patch_sample(self, sample: dict[str, Any]) -> dict[str, Any]:
        sample = self._maybe_augment(sample)
        flux = sample["flux"]
        ivar = sample["ivar"]
        lam = sample["lambda"]
        lsf = sample["lsf_sigma"]
        valid = valid_pixel_mask(sample)
        n = len(flux)
        p = self.cfg.num_patches
        edges = np.linspace(0, n, p + 1, dtype=np.int64)

        patch_id = np.searchsorted(edges[1:-1], np.arange(n), side="right")
        patch_valid = np.zeros(p, dtype=np.bool_)
        target_loglam = np.zeros(p, dtype=np.float32)
        target_aux = np.zeros((p, 4), dtype=np.float32)
        raw_patch_flux = np.zeros(p, dtype=np.float32)
        raw_patch_ivar = np.zeros(p, dtype=np.float32)
        raw_patch_lsf = np.zeros(p, dtype=np.float32)
        valid_frac = np.zeros(p, dtype=np.float32)

        safe_loglam = np.log(np.clip(lam.astype(np.float64), 1.0, None)).astype(np.float32)
        for j in range(p):
            lo, hi = int(edges[j]), int(edges[j + 1])
            idx = np.arange(lo, hi)
            if len(idx) == 0:
                continue
            v = valid[idx]
            valid_frac[j] = float(v.mean())
            patch_valid[j] = bool(v.sum() >= max(2, int(0.2 * len(idx))))
            target_loglam[j] = float(np.nanmedian(safe_loglam[idx]))
            width = float(np.nanmax(safe_loglam[idx]) - np.nanmin(safe_loglam[idx])) if len(idx) > 1 else 0.0
            if v.any():
                raw_patch_flux[j] = float(np.nanmedian(flux[idx][v]))
                raw_patch_ivar[j] = float(np.nanmedian(ivar[idx][v]))
                raw_patch_lsf[j] = float(np.nanmedian(lsf[idx][v]))
            target_aux[j] = [valid_frac[j], math.log1p(max(raw_patch_ivar[j], 0.0)), raw_patch_lsf[j], width]

        corrupt = self._sample_corruption(raw_patch_flux, patch_valid)
        pixel_visible = valid & (~corrupt[patch_id])
        if pixel_visible.sum() < 16:
            pixel_visible = valid

        center = float(np.nanmedian(flux[pixel_visible])) if pixel_visible.any() else 0.0
        abs_dev = np.abs(flux[pixel_visible] - center) if pixel_visible.any() else np.asarray([1.0], dtype=np.float32)
        scale = float(np.nanmedian(abs_dev) * 1.4826)
        if not math.isfinite(scale) or scale < self.cfg.min_scale:
            scale = max(float(np.nanmedian(np.abs(flux[valid]))) if valid.any() else 1.0, self.cfg.min_scale)

        norm_patch_flux = np.arcsinh((raw_patch_flux - center) / scale).astype(np.float32)
        norm_ivar = np.log1p(np.maximum(raw_patch_ivar * scale * scale, 0.0)).astype(np.float32)
        line_weight = self._line_weights(norm_patch_flux, patch_valid)

        visible_patch = patch_valid & (~corrupt)
        # Encoder never receives masked target flux. It receives visible patches only.
        enc_idx = np.where(visible_patch)[0]
        if len(enc_idx) == 0:
            enc_idx = np.where(patch_valid)[0][:1]

        enc_features = np.stack(
            [
                norm_patch_flux[enc_idx],
                norm_ivar[enc_idx],
                valid_frac[enc_idx],
                raw_patch_lsf[enc_idx],
                np.zeros(len(enc_idx), dtype=np.float32),
                target_aux[enc_idx, 3],
            ],
            axis=-1,
        ).astype(np.float32)

        loss_mask = patch_valid & corrupt
        return {
            "enc_features": enc_features,
            "enc_loglam": target_loglam[enc_idx],
            "target_loglam": target_loglam,
            "target_aux": target_aux.astype(np.float32),
            "target_flux": norm_patch_flux,
            "loss_mask": loss_mask,
            "valid_patch": patch_valid,
            "corrupt_patch": corrupt,
            "line_weight": line_weight,
            "z": sample["z"],
            "zwarn": sample["zwarn"],
        }

    def _sample_corruption(self, patch_flux: np.ndarray, valid_patch: np.ndarray) -> np.ndarray:
        p = len(valid_patch)
        corrupt = np.zeros(p, dtype=np.bool_)
        valid_idx = np.where(valid_patch)[0]
        if len(valid_idx) == 0:
            return corrupt

        ratio = self.cfg.random_mask_ratio if self.train else 0.35
        if ratio <= 0 and self.cfg.span_mask_prob <= 0 and self.cfg.line_mask_prob <= 0 and self.cfg.arm_dropout_prob <= 0:
            return corrupt
        n_rand = max(1, int(round(len(valid_idx) * ratio)))
        corrupt[self.rng.choice(valid_idx, size=min(n_rand, len(valid_idx)), replace=False)] = True

        if self.train and self.rng.random() < self.cfg.span_mask_prob:
            for _ in range(int(self.rng.integers(1, 4))):
                width = int(self.rng.integers(max(3, p // 80), max(4, p // 18)))
                start = int(self.rng.integers(0, max(1, p - width)))
                corrupt[start : start + width] |= valid_patch[start : start + width]

        if self.train and self.rng.random() < self.cfg.arm_dropout_prob:
            arm = int(self.rng.integers(0, 4))
            if arm == 0:
                sl = slice(0, p // 3)
            elif arm == 1:
                sl = slice(p // 3, 2 * p // 3)
            elif arm == 2:
                sl = slice(2 * p // 3, p)
            else:
                lo = int(self.rng.integers(p // 5, p // 2))
                hi = min(p, lo + int(self.rng.integers(p // 12, p // 5)))
                sl = slice(lo, hi)
            corrupt[sl] |= valid_patch[sl]

        if self.train and self.rng.random() < self.cfg.line_mask_prob:
            score = np.zeros(p, dtype=np.float32)
            good = valid_patch & np.isfinite(patch_flux)
            if good.sum() > 4:
                grad = np.abs(np.gradient(np.nan_to_num(patch_flux, nan=0.0)))
                score[good] = grad[good]
                top = np.argsort(score)[-max(2, p // 24) :]
                for j in top:
                    lo, hi = max(0, j - 1), min(p, j + 2)
                    corrupt[lo:hi] |= valid_patch[lo:hi]

        max_allowed = max(1, int(valid_patch.sum() * self.cfg.max_mask_ratio))
        cur = np.where(corrupt & valid_patch)[0]
        if len(cur) > max_allowed:
            keep = self.rng.choice(cur, size=max_allowed, replace=False)
            next_corrupt = np.zeros_like(corrupt)
            next_corrupt[keep] = True
            corrupt = next_corrupt
        return corrupt & valid_patch

    def _line_weights(self, patch_flux: np.ndarray, valid_patch: np.ndarray) -> np.ndarray:
        w = np.ones_like(patch_flux, dtype=np.float32)
        if valid_patch.sum() < 4:
            return w
        grad = np.abs(np.gradient(np.nan_to_num(patch_flux, nan=0.0))).astype(np.float32)
        scale = np.percentile(grad[valid_patch], 90) if valid_patch.any() else 1.0
        if scale > 0:
            w += 2.0 * np.clip(grad / scale, 0.0, 2.0)
        w[~valid_patch] = 1.0
        return np.clip(w, 1.0, 5.0)

    def _maybe_augment(self, sample: dict[str, Any]) -> dict[str, Any]:
        if not (self.train and self.cfg.augment_ood):
            return sample
        out = {k: v for k, v in sample.items()}
        bad = np.array(sample["bad_mask"], copy=True)
        lam = sample["lambda"]
        n = len(lam)
        if self.rng.random() < 0.45:
            frac = float(self.rng.uniform(0.65, 0.95))
            width = max(32, int(n * frac))
            start = int(self.rng.integers(0, max(1, n - width)))
            keep = np.zeros(n, dtype=np.bool_)
            keep[start : start + width] = True
            bad |= ~keep
        if self.rng.random() < 0.30:
            for _ in range(int(self.rng.integers(1, 4))):
                width = int(self.rng.integers(max(8, n // 200), max(12, n // 45)))
                start = int(self.rng.integers(0, max(1, n - width)))
                bad[start : start + width] = True
        out["bad_mask"] = bad
        return out