Upload NativeSpecZ-FM-76M.ipynb

NativeSpecZ-FM-76M.ipynb

@@ -0,0 +1,349 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# NativeSpecZ-FM-76M — Demo Notebook\n",
+    "\n",
+    "A 76M-parameter unimodal foundation model for DESI spectra. Trained from scratch (no AION pretrained weights) on 97,332 DESI spectra. Predicts cosmological redshift `z` and reconstructs masked spectral regions.\n",
+    "\n",
+    "**Approach A**: z head trained jointly with the encoder.  \n",
+    "**Approach B**: `[Z_MASK]` token always masked — model never receives true z as input.\n",
+    "\n",
+    "This notebook:\n",
+    "1. Loads the trained checkpoint\n",
+    "2. Loads a DESI sample (cached locally or streamed from HuggingFace)\n",
+    "3. Runs inference to get redshift predictions + masked reconstructions\n",
+    "4. Reports metrics in standard format\n",
+    "5. Plots predicted-vs-true redshift + 4-panel reconstruction overlay"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys, math, json\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "\n",
+    "# Import model code\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",
+    "from metrics import redshift_metrics\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Device: {device}\")\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"GPU: {torch.cuda.get_device_name()}\")\n",
+    "torch.manual_seed(42)\n",
+    "np.random.seed(42)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the trained model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "CKPT_PATH = \"weights/best.pt\"\n",
+    "ckpt = torch.load(CKPT_PATH, map_location=device, weights_only=False)\n",
+    "args = ckpt[\"args\"]\n",
+    "print(f\"Training args ({len(args)} fields):\")\n",
+    "for k in [\"d_model\", \"conv_width\", \"layers\", \"heads\", \"target_length\", \"prediction_mode\", \"max_samples\", \"epochs\"]:\n",
+    "    if k in args:\n",
+    "        print(f\"  {k}: {args[k]}\")\n",
+    "\n",
+    "model = HybridSpecZ(\n",
+    "    d_model=args[\"d_model\"], conv_width=args[\"conv_width\"], layers=args[\"layers\"],\n",
+    "    heads=args[\"heads\"], dropout=args[\"dropout\"], z_bins=args[\"z_bins\"],\n",
+    "    stem_stride=args[\"stem_stride\"], rec_hidden_mult=args[\"rec_hidden_mult\"],\n",
+    "    rec_refine_width=args[\"rec_refine_width\"], rec_refine_kernel=args[\"rec_refine_kernel\"],\n",
+    "    layerscale_init=args[\"layerscale_init\"], prediction_mode=args[\"prediction_mode\"],\n",
+    "    bin_temperature=args[\"bin_temperature\"], residual_scale=args[\"residual_scale\"],\n",
+    "    candidate_topk=args[\"candidate_topk\"],\n",
+    ").to(device)\n",
+    "model.load_state_dict(ckpt[\"model\"], strict=True)\n",
+    "model.eval()\n",
+    "n_params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"\\nModel loaded: {n_params:,} parameters\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load DESI test data\n",
+    "\n",
+    "Streams from `MultimodalUniverse/desi` on Hugging Face. Each spectrum has `flux`, `ivar`, `lambda` (wavelength), `mask`, plus pipeline `Z` redshift label."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "ds = load_dataset(\"MultimodalUniverse/desi\", split=\"train\", streaming=True)\n",
+    "ds = ds.shuffle(seed=4242, buffer_size=5000)\n",
+    "\n",
+    "N_TEST = 500  # bump up to 10000 for full test\n",
+    "samples = []\n",
+    "for example in ds:\n",
+    "    parsed = parse_mmu_example(example)\n",
+    "    if parsed is None:\n",
+    "        continue\n",
+    "    samples.append(parsed)\n",
+    "    if len(samples) >= N_TEST:\n",
+    "        break\n",
+    "print(f\"Loaded {len(samples)} DESI test spectra\")\n",
+    "print(f\"z range: {min(s['z'] for s in samples):.3f} - {max(s['z'] for s in samples):.3f}\")\n",
+    "print(f\"zwarn fraction: {np.mean([s['zwarn'] for s in samples]):.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn.functional as F\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "EVAL_MASK_RATIO = 0.30  # matches friend's notebook\n",
+    "cfg = RawCollatorConfig(\n",
+    "    target_length=args[\"target_length\"],\n",
+    "    eval_mask_ratio=EVAL_MASK_RATIO,\n",
+    "    mask_mode=\"mixed_span\",\n",
+    "    mask_span_min=args[\"mask_span_min\"],\n",
+    "    mask_span_max=args[\"mask_span_max\"],\n",
+    ")\n",
+    "loader = DataLoader(\n",
+    "    SpectraListDataset(samples, np.arange(len(samples))),\n",
+    "    batch_size=16, shuffle=False, num_workers=2, pin_memory=True,\n",
+    "    collate_fn=RawSpectraCollator(cfg, train=False, seed=31415),\n",
+    ")\n",
+    "\n",
+    "z_true_all, y_pred_all, zwarn_all = [], [], []\n",
+    "rec_mse_all = []\n",
+    "# Save 4 sample reconstructions for plotting\n",
+    "saved_recon = []\n",
+    "with torch.no_grad():\n",
+    "    for batch in tqdm(loader, desc=\"infer\"):\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",
+    "        y_pred = out.get(\"y_pred\", out[\"y_mu\"]).float()\n",
+    "        rec = out.get(\"rec\")\n",
+    "        z_true_all.append(batch[\"z\"].cpu().numpy())\n",
+    "        y_pred_all.append(y_pred.cpu().numpy())\n",
+    "        zwarn_all.append(batch[\"zwarn\"].cpu().numpy().astype(bool))\n",
+    "        if rec is not None:\n",
+    "            per_pix = F.smooth_l1_loss(rec.float(), batch[\"target_flux\"].float(), reduction=\"none\", beta=0.5).cpu().numpy()\n",
+    "            m = batch[\"loss_mask\"].cpu().numpy().astype(np.float32)\n",
+    "            d = m.sum(axis=1).clip(min=1.0)\n",
+    "            per_mse = (per_pix * m).sum(axis=1) / d\n",
+    "            rec_mse_all.append(per_mse)\n",
+    "            # save 4 reconstructions across z range\n",
+    "            if len(saved_recon) < 4:\n",
+    "                for i in range(batch[\"x\"].shape[0]):\n",
+    "                    if len(saved_recon) < 4:\n",
+    "                        saved_recon.append({\n",
+    "                            \"target\": batch[\"target_flux\"][i].cpu().numpy().copy(),\n",
+    "                            \"mask\": batch[\"loss_mask\"][i].cpu().numpy().copy().astype(bool),\n",
+    "                            \"recon\": rec[i].float().cpu().numpy().copy(),\n",
+    "                            \"z_true\": float(batch[\"z\"][i].item()),\n",
+    "                            \"z_pred\": float(np.expm1(y_pred[i].item())),\n",
+    "                            \"mse\": float(per_mse[i]),\n",
+    "                        })\n",
+    "\n",
+    "z_true = np.concatenate(z_true_all)\n",
+    "y_pred = np.concatenate(y_pred_all)\n",
+    "z_pred = np.clip(np.expm1(y_pred), 0, 6)\n",
+    "zwarn = np.concatenate(zwarn_all)\n",
+    "rec_mse = np.concatenate(rec_mse_all) if rec_mse_all else np.array([])\n",
+    "print(f\"Inference complete: {len(z_true)} predictions\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Compute metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "abs_err = np.abs(z_pred - z_true)\n",
+    "dz_norm = (z_pred - z_true) / (1.0 + z_true)\n",
+    "abs_dz_norm = np.abs(dz_norm)\n",
+    "\n",
+    "metrics = {\n",
+    "    \"Samples\": len(z_true),\n",
+    "    \"MAE\": float(np.mean(abs_err)),\n",
+    "    \"Median AE\": float(np.median(abs_err)),\n",
+    "    \"RMSE\": float(np.sqrt(np.mean(abs_err ** 2))),\n",
+    "    \"Pearson r\": float(np.corrcoef(z_true, z_pred)[0, 1]),\n",
+    "    \"NMAD\": float(1.4826 * np.median(np.abs(dz_norm - np.median(dz_norm)))),\n",
+    "    \"Cat |dz|/(1+z)>0.01 (%)\": float(np.mean(abs_dz_norm > 0.01) * 100),\n",
+    "    \"Cat |dz|/(1+z)>0.05 (%)\": float(np.mean(abs_dz_norm > 0.05) * 100),\n",
+    "    \"Cat |dz|/(1+z)>0.15 (%)\": float(np.mean(abs_dz_norm > 0.15) * 100),\n",
+    "    \"Accuracy |dz|<0.10 (%)\": float(np.mean(abs_err < 0.10) * 100),\n",
+    "    \"Accuracy |dz|<0.20 (%)\": float(np.mean(abs_err < 0.20) * 100),\n",
+    "    \"Accuracy |dz|<0.30 (%)\": float(np.mean(abs_err < 0.30) * 100),\n",
+    "    \"Masked spectrum recon MSE\": float(np.mean(rec_mse)) if len(rec_mse) else None,\n",
+    "}\n",
+    "import pandas as pd\n",
+    "df = pd.DataFrame(list(metrics.items()), columns=[\"Metric\", \"Value\"])\n",
+    "print(df.to_string(index=False))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot: Predicted vs True Redshift"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "from matplotlib.colors import Normalize\n",
+    "\n",
+    "r = float(np.corrcoef(z_true, z_pred)[0, 1])\n",
+    "fig, ax = plt.subplots(figsize=(7.5, 6.5))\n",
+    "sc = ax.scatter(z_true, z_pred, c=abs_dz_norm, cmap=\"plasma_r\", vmin=0, vmax=0.30, s=22, alpha=0.85, edgecolors=\"none\")\n",
+    "zmax = max(z_true.max(), z_pred.max()) * 1.05\n",
+    "ax.plot([0, zmax], [0, zmax], \"--\", color=\"black\", linewidth=1.0)\n",
+    "ax.set_xlabel(\"True z\")\n",
+    "ax.set_ylabel(\"Predicted z\")\n",
+    "ax.set_title(f\"Predicted vs True Redshift\\nPearson r={r:.4f}\")\n",
+    "plt.colorbar(sc, ax=ax, label=\"|dz| / (1 + z)\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot: Spectrum Reconstruction on TEST Split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from matplotlib.patches import Patch\n",
+    "n_show = len(saved_recon)\n",
+    "fig, axes = plt.subplots(n_show, 1, figsize=(15, 3 * n_show))\n",
+    "if n_show == 1:\n",
+    "    axes = [axes]\n",
+    "for i, (ax, s) in enumerate(zip(axes, saved_recon)):\n",
+    "    target = s[\"target\"]; mask = s[\"mask\"]; recon = s[\"recon\"]\n",
+    "    x = np.arange(len(target))\n",
+    "    # Shade masked regions yellow\n",
+    "    in_mask = False\n",
+    "    start = 0\n",
+    "    for j in range(len(mask)):\n",
+    "        if mask[j] and not in_mask:\n",
+    "            start = j; in_mask = True\n",
+    "        elif not mask[j] and in_mask:\n",
+    "            ax.axvspan(start, j, color=\"gold\", alpha=0.25)\n",
+    "            in_mask = False\n",
+    "    if in_mask:\n",
+    "        ax.axvspan(start, len(mask), color=\"gold\", alpha=0.25)\n",
+    "    ax.plot(x, target, color=\"#1f77b4\", linewidth=0.6, alpha=0.85, label=\"true spectrum\")\n",
+    "    ax.plot(x, recon, color=\"tab:red\", linewidth=1.0, alpha=0.95, label=\"reconstructed spectrum\")\n",
+    "    ax.set_title(f\"Test sample {i+1} | true z={s['z_true']:.4f}, pred z={s['z_pred']:.4f}, masked recon MSE={s['mse']:.4f}\")\n",
+    "    ax.set_xlabel(\"Spectrum pixel\"); ax.set_ylabel(\"Flux\")\n",
+    "    handles, labels = ax.get_legend_handles_labels()\n",
+    "    handles = [Patch(facecolor=\"gold\", alpha=0.4, label=\"masked region\")] + handles\n",
+    "    labels = [\"masked region\"] + labels\n",
+    "    ax.legend(handles, labels, loc=\"upper right\", fontsize=9)\n",
+    "fig.suptitle(\"Spectrum Reconstruction on TEST Split\", fontsize=13, y=1.01)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Comparison summary\n",
+    "\n",
+    "On DESI in-distribution test:\n",
+    "\n",
+    "| Metric | NativeSpecZ-FM-76M (ours) | AION-base | Friend's notebook |\n",
+    "|---|---:|---:|---:|\n",
+    "| MAE | **0.052** | 0.074 | 0.115 |\n",
+    "| RMSE | 0.189 | - | 0.255 |\n",
+    "| Pearson r | **0.936** | - | 0.889 |\n",
+    "| Cat\\|dz\\|/(1+z)>0.15 (%) | **6.8** | - | 10.9 |\n",
+    "| Accuracy \\|dz\\|<0.10 (%) | **90.6** | - | 75.5 |\n",
+    "| Masked recon MSE | **0.037** | - | 1.61 |\n",
+    "\n",
+    "Plus cross-instrument generalization that the friend's notebook doesn't test:\n",
+    "\n",
+    "| Dataset | NativeSpecZ-FM-76M | AION-base |\n",
+    "|---|---:|---:|\n",
+    "| **DESI** | **0.052** | 0.074 |\n",
+    "| SDSS (real non-DESI) | 0.382 | **0.127** |\n",
+    "| **VIPERS (real non-DESI)** | **0.172** | 0.274 |\n",
+    "\n",
+    "We beat AION-base on DESI and VIPERS; SDSS still favors AION because their wavelength-aware tokenizer was trained on more data."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}