Upload code/data.py

code/data.py

@@ -0,0 +1,407 @@
+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