{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NativeSpecZ-296M — Demo Notebook (scale-up ablation)\n",
    "\n",
    "A 296M-parameter unimodal foundation model for DESI spectra, trained from scratch (no AION pretrained weights) with wavelength-grid jitter augmentation. Hits the spec's ~300M target and is competitive with AION-base on DESI in-distribution, but does NOT beat AION on out-of-distribution data — the smaller NativeSpecZ-FM-76M is the recommended headline.\n",
    "\n",
    "**Approach A**: z head trained jointly with the encoder.  \n",
    "**Approach B**: `[Z_MASK]` token always masked.\n",
    "\n",
    "NOTE: load with `strict=False` — the current `hybrid_redshift.py` has `z_rerank_head` / `z_calib_head` modules added after this checkpoint was trained; they are unused in the `bin_residual` prediction path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, json\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "sys.path.append(\"code\")\n",
    "from hybrid_redshift import HybridSpecZ, RawCollatorConfig, RawSpectraCollator, move_to_device\n",
    "from data import SpectraListDataset, parse_mmu_example\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Device: {device}\")\n",
    "torch.manual_seed(42); np.random.seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the 296M model (strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt = torch.load(\"weights/best.pt\", map_location=device, weights_only=False)\n",
    "a = ckpt[\"args\"]\n",
    "model = HybridSpecZ(\n",
    "    d_model=a[\"d_model\"], conv_width=a[\"conv_width\"], layers=a[\"layers\"],\n",
    "    heads=a[\"heads\"], dropout=a[\"dropout\"], z_bins=a[\"z_bins\"],\n",
    "    stem_stride=a[\"stem_stride\"], rec_hidden_mult=a[\"rec_hidden_mult\"],\n",
    "    rec_refine_width=a[\"rec_refine_width\"], rec_refine_kernel=a[\"rec_refine_kernel\"],\n",
    "    layerscale_init=a[\"layerscale_init\"], prediction_mode=a[\"prediction_mode\"],\n",
    "    bin_temperature=a[\"bin_temperature\"], residual_scale=a[\"residual_scale\"],\n",
    "    candidate_topk=a[\"candidate_topk\"],\n",
    ").to(device)\n",
    "missing, unexpected = model.load_state_dict(ckpt[\"model\"], strict=False)\n",
    "model.eval()\n",
    "print(f\"Params: {sum(p.numel() for p in model.parameters()):,}\")\n",
    "print(f\"Missing keys (unused rerank/calib heads): {len(missing)}; unexpected: {len(unexpected)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load DESI test data + run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from tqdm.auto import tqdm\n",
    "import torch.nn.functional as F\n",
    "\n",
    "ds = load_dataset(\"MultimodalUniverse/desi\", split=\"train\", streaming=True).shuffle(seed=4242, buffer_size=5000)\n",
    "samples = []\n",
    "for ex in ds:\n",
    "    p = parse_mmu_example(ex)\n",
    "    if p is not None: samples.append(p)\n",
    "    if len(samples) >= 500: break\n",
    "print(f\"Loaded {len(samples)} DESI spectra\")\n",
    "\n",
    "cfg = RawCollatorConfig(target_length=a[\"target_length\"], eval_mask_ratio=0.30, mask_mode=\"mixed_span\", mask_span_min=a[\"mask_span_min\"], mask_span_max=a[\"mask_span_max\"])\n",
    "loader = DataLoader(SpectraListDataset(samples, np.arange(len(samples))), batch_size=8, shuffle=False, num_workers=2, collate_fn=RawSpectraCollator(cfg, train=False, seed=31415))\n",
    "z_true_all, y_pred_all = [], []\n",
    "with torch.no_grad():\n",
    "    for batch in tqdm(loader):\n",
    "        batch = move_to_device(batch, device)\n",
    "        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type==\"cuda\"):\n",
    "            out = model(batch[\"x\"], batch[\"valid\"], batch[\"loglam\"])\n",
    "        z_true_all.append(batch[\"z\"].cpu().numpy())\n",
    "        y_pred_all.append(out.get(\"y_pred\", out[\"y_mu\"]).float().cpu().numpy())\n",
    "z_true = np.concatenate(z_true_all)\n",
    "z_pred = np.clip(np.expm1(np.concatenate(y_pred_all)), 0, 6)\n",
    "mae = float(np.mean(np.abs(z_pred - z_true)))\n",
    "r = float(np.corrcoef(z_true, z_pred)[0,1])\n",
    "print(f\"DESI MAE={mae:.4f}  Pearson r={r:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Held-out results (precomputed, n=2500/2000/2000) + comparison\n",
    "\n",
    "| Dataset | NativeSpecZ-296M | NativeSpecZ-76M (headline) | AION-base |\n",
    "|---|---:|---:|---:|\n",
    "| **DESI** | **0.067** | 0.069 | 0.074 |\n",
    "| SDSS (real non-DESI) | 0.314 | 0.382 | **0.127** |\n",
    "| VIPERS (real non-DESI) | 0.316 | **0.172** | 0.274 |\n",
    "\n",
    "The 296M ties the 76M on DESI (both beat AION) and improves SDSS over the 76M, but loses the 76M's strong VIPERS result. It beats AION-base on no OOD dataset, so the 76M remains the recommended headline; the 296M is the scale-up ablation that hits the ~300M target."
   ]
  }
 ],
 "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.11"}},
 "nbformat": 4,
 "nbformat_minor": 4
}